From a7f75f117aca893eeae6570548b96cef372f2699 Mon Sep 17 00:00:00 2001 From: ParkiratS Date: Wed, 24 Jun 2026 07:43:12 -0400 Subject: [PATCH 01/84] Removed existing transpiler --- README.md | 29 +- cactus-engine/src/model.cpp | 6 +- cactus-engine/src/model_npu.cpp | 2 +- cactus-engine/src/transcribe.cpp | 2 +- cactus-engine/tests/test_llm.cpp | 4 +- docs/cactus_engine.md | 5 +- docs/cactus_quants.md | 16 +- docs/cactus_transpiler.md | 520 -- docs/finetuning.md | 11 +- list.txt | 152 + mkdocs.yml | 1 - python/README.md | 4 +- python/cactus/cli/__init__.py | 134 +- python/cactus/cli/clean.py | 2 - python/cactus/cli/common.py | 6 - python/cactus/cli/convert.py | 99 +- python/cactus/cli/list.py | 32 +- python/cactus/cli/model.py | 319 +- python/cactus/cli/run.py | 7 +- python/cactus/cli/serve.py | 2 +- python/cactus/cli/transcribe.py | 4 +- python/cactus/cli/transpile.py | 66 - python/cactus/cli/upload.py | 53 - python/cactus/convert/compat.py | 179 +- .../cactus/convert/tests/test_naming_qdq.py | 4 +- python/cactus/server.py | 2 +- python/cactus/transpile/README.md | 1024 --- python/cactus/transpile/__init__.py | 1 - python/cactus/transpile/aten_ops.py | 184 - python/cactus/transpile/audio_preprocess.py | 400 -- .../cactus/transpile/canonicalize/cleanup.py | 704 -- python/cactus/transpile/canonicalize/utils.py | 277 - python/cactus/transpile/capture_jax.py | 1663 ----- python/cactus/transpile/capture_pytorch.py | 375 - .../cactus/transpile/component_partition.py | 335 - python/cactus/transpile/component_pipeline.py | 213 - python/cactus/transpile/component_plan.py | 239 - python/cactus/transpile/fusion/__init__.py | 47 - python/cactus/transpile/fusion/attention.py | 550 -- python/cactus/transpile/fusion/common.py | 73 - python/cactus/transpile/fusion/conv.py | 123 - python/cactus/transpile/fusion/deltanet.py | 357 - python/cactus/transpile/fusion/linear.py | 58 - python/cactus/transpile/fusion/lstm.py | 216 - python/cactus/transpile/fusion/mlp.py | 53 - .../cactus/transpile/fusion/rel_pos_bias.py | 179 - python/cactus/transpile/fusion/rms_norm.py | 133 - python/cactus/transpile/fusion/rope.py | 252 - python/cactus/transpile/graph_ir.py | 201 - python/cactus/transpile/hf_model.py | 3406 --------- python/cactus/transpile/import_ir.py | 277 - python/cactus/transpile/import_semantics.py | 337 - python/cactus/transpile/importers.py | 1869 ----- .../cactus/transpile/jax_semantic_rewrites.py | 875 --- .../cactus/transpile/jax_user_graph_bundle.py | 575 -- python/cactus/transpile/lower.py | 3437 --------- python/cactus/transpile/media_limits.py | 35 - python/cactus/transpile/model_adapters.py | 6326 ----------------- python/cactus/transpile/model_patterns.py | 64 - python/cactus/transpile/model_profiles.py | 305 - python/cactus/transpile/multimodal_runtime.py | 722 -- python/cactus/transpile/normalize.py | 29 - python/cactus/transpile/npu/README.md | 86 - python/cactus/transpile/npu/__init__.py | 3 - python/cactus/transpile/npu/audio.py | 115 - .../transpile/npu/coremltools_patches.py | 434 -- python/cactus/transpile/npu/pipeline.py | 94 - python/cactus/transpile/npu/source.py | 118 - python/cactus/transpile/npu/vision.py | 126 - python/cactus/transpile/ops.py | 255 - python/cactus/transpile/optimize_graph.py | 1653 ----- python/cactus/transpile/runtime_compat.py | 376 - python/cactus/transpile/runtime_support.py | 178 - python/cactus/transpile/tdt_runtime.py | 1028 --- python/cactus/transpile/weight_binding.py | 317 - python/cactus/transpile/weight_compat.py | 547 -- python/pyproject.toml | 1 - python/tests/test_bindings.py | 57 +- python/tests/test_capture_jax.py | 197 - .../test_capture_jax_user_graph_bundle.py | 196 - python/tests/test_cli_transpile_defaults.py | 157 - python/tests/test_component_partition.py | 55 - python/tests/test_encoder_cross_kv_route.py | 79 - python/tests/test_nomic_text_embedding.py | 37 - .../tests/test_optimize_gemma4_attention.py | 328 - python/tests/test_server_live.py | 18 +- python/tests/test_transpile_aten_ops.py | 124 - python/tests/test_transpile_layer_keys.py | 68 - python/tests/test_transpile_weight_compat.py | 361 - python/tests/transpile/__init__.py | 1 - python/tests/transpile/tools/__init__.py | 1 - .../tools/compare_gemma4_multimodal.py | 932 --- .../tools/compare_parakeet_performance.py | 792 --- .../transpile/tools/inspect_gemma4_cleanup.py | 246 - .../tools/inspect_transpile_cleanup.py | 144 - .../transpile/tools/torch_capture_smoke.py | 79 - .../transpile/tools/transpile_hf_causal_lm.py | 20 - 97 files changed, 448 insertions(+), 36350 deletions(-) delete mode 100644 docs/cactus_transpiler.md create mode 100644 list.txt delete mode 100644 python/cactus/cli/transpile.py delete mode 100644 python/cactus/cli/upload.py delete mode 100644 python/cactus/transpile/README.md delete mode 100644 python/cactus/transpile/__init__.py delete mode 100644 python/cactus/transpile/aten_ops.py delete mode 100644 python/cactus/transpile/audio_preprocess.py delete mode 100644 python/cactus/transpile/canonicalize/cleanup.py delete mode 100644 python/cactus/transpile/canonicalize/utils.py delete mode 100644 python/cactus/transpile/capture_jax.py delete mode 100644 python/cactus/transpile/capture_pytorch.py delete mode 100644 python/cactus/transpile/component_partition.py delete mode 100644 python/cactus/transpile/component_pipeline.py delete mode 100644 python/cactus/transpile/component_plan.py delete mode 100644 python/cactus/transpile/fusion/__init__.py delete mode 100644 python/cactus/transpile/fusion/attention.py delete mode 100644 python/cactus/transpile/fusion/common.py delete mode 100644 python/cactus/transpile/fusion/conv.py delete mode 100644 python/cactus/transpile/fusion/deltanet.py delete mode 100644 python/cactus/transpile/fusion/linear.py delete mode 100644 python/cactus/transpile/fusion/lstm.py delete mode 100644 python/cactus/transpile/fusion/mlp.py delete mode 100644 python/cactus/transpile/fusion/rel_pos_bias.py delete mode 100644 python/cactus/transpile/fusion/rms_norm.py delete mode 100644 python/cactus/transpile/fusion/rope.py delete mode 100644 python/cactus/transpile/graph_ir.py delete mode 100644 python/cactus/transpile/hf_model.py delete mode 100644 python/cactus/transpile/import_ir.py delete mode 100644 python/cactus/transpile/import_semantics.py delete mode 100644 python/cactus/transpile/importers.py delete mode 100644 python/cactus/transpile/jax_semantic_rewrites.py delete mode 100644 python/cactus/transpile/jax_user_graph_bundle.py delete mode 100644 python/cactus/transpile/lower.py delete mode 100644 python/cactus/transpile/media_limits.py delete mode 100644 python/cactus/transpile/model_adapters.py delete mode 100644 python/cactus/transpile/model_patterns.py delete mode 100644 python/cactus/transpile/model_profiles.py delete mode 100644 python/cactus/transpile/multimodal_runtime.py delete mode 100644 python/cactus/transpile/normalize.py delete mode 100644 python/cactus/transpile/npu/README.md delete mode 100644 python/cactus/transpile/npu/__init__.py delete mode 100644 python/cactus/transpile/npu/audio.py delete mode 100644 python/cactus/transpile/npu/coremltools_patches.py delete mode 100644 python/cactus/transpile/npu/pipeline.py delete mode 100644 python/cactus/transpile/npu/source.py delete mode 100644 python/cactus/transpile/npu/vision.py delete mode 100644 python/cactus/transpile/ops.py delete mode 100644 python/cactus/transpile/optimize_graph.py delete mode 100644 python/cactus/transpile/runtime_compat.py delete mode 100644 python/cactus/transpile/runtime_support.py delete mode 100644 python/cactus/transpile/tdt_runtime.py delete mode 100644 python/cactus/transpile/weight_binding.py delete mode 100644 python/cactus/transpile/weight_compat.py delete mode 100644 python/tests/test_capture_jax.py delete mode 100644 python/tests/test_capture_jax_user_graph_bundle.py delete mode 100644 python/tests/test_cli_transpile_defaults.py delete mode 100644 python/tests/test_component_partition.py delete mode 100644 python/tests/test_encoder_cross_kv_route.py delete mode 100644 python/tests/test_nomic_text_embedding.py delete mode 100644 python/tests/test_optimize_gemma4_attention.py delete mode 100644 python/tests/test_transpile_aten_ops.py delete mode 100644 python/tests/test_transpile_layer_keys.py delete mode 100644 python/tests/test_transpile_weight_compat.py delete mode 100644 python/tests/transpile/__init__.py delete mode 100644 python/tests/transpile/tools/__init__.py delete mode 100644 python/tests/transpile/tools/compare_gemma4_multimodal.py delete mode 100644 python/tests/transpile/tools/compare_parakeet_performance.py delete mode 100644 python/tests/transpile/tools/inspect_gemma4_cleanup.py delete mode 100644 python/tests/transpile/tools/inspect_transpile_cleanup.py delete mode 100644 python/tests/transpile/tools/torch_capture_smoke.py delete mode 100644 python/tests/transpile/tools/transpile_hf_causal_lm.py diff --git a/README.md b/README.md index 916487946..0874fc0bb 100644 --- a/README.md +++ b/README.md @@ -15,7 +15,7 @@ A hybrid edge-cloud AI engine for mobile devices & wearables. - **Low RAM:** zero-copy memory mapping ensures 10x lower RAM use than other engines - **Multimodal:** one engine for speech, vision, and language models - **Cloud fallback:** automatically route requests to cloud models if needed -- **Model-Agnostic:** Custom PyTorch models can be exported to the Cactus runtime. +- **Portable bundles:** download prebuilt Cactus bundles for local on-device inference ``` ┌─────────────────┐ @@ -33,10 +33,6 @@ A hybrid edge-cloud AI engine for mobile devices & wearables. ┌─────────────────┐ │ Cactus Quants │ ←── Cactus Quants at 4-bit uniform matches f16. └─────────────────┘ - │ -┌─────────────────┐ -│Cactus Transpiler│ ←── Transpiles custom PyTorch model to Cactus. -└─────────────────┘ ``` ## Quick Demo (Mac) @@ -148,10 +144,10 @@ graph.hard_reset(); ## Supported Models -- Any HuggigFace model can be converted using `cactus convert [HF-Name]`, though experimental. -- Liquid, Gemma. whisper. parakeet and Qwen model families are especially tested. -- Some models have been pre-uploaded [here](https://huggingface.co/Cactus-Compute), just run `cactus download [HF-Name]`. -- `cactus run [HF-Name]` albeit first downloads or convert the model if not found. +- Liquid, Gemma, Whisper, Parakeet, and Qwen model families are especially tested. +- Prebuilt runnable bundles are published [here](https://huggingface.co/Cactus-Compute); run `cactus download [HF-Name]`. +- `cactus run [HF-Name]` downloads a matching prebuilt bundle when needed. +- `cactus convert [HF-Name]` currently quantizes weights to Cactus CQ format only; local runtime bundle generation is unavailable while the graph builder is being rewritten. ## Learn More @@ -162,7 +158,6 @@ graph.hard_reset(); | [Cactus Kernels](/docs/cactus_kernels.md) | C++ | ARM NEON SIMD kernels for matmul, attention, convolution, quantization, DSP, image processing | | [Cactus Quants](/docs/cactus_quants.md) | C++ | Rotation-and-codebook quantization from 4-bit to 1-bit for all weight tensors | | [Cactus Hybrid](/docs/cactus_hybrid.md) | C/Python | Route hard queries to the cloud automatically based on local model confidence | -| [Cactus Transpiler](/docs/cactus_transpiler.md) | Python | Convert any PyTorch model to a Cactus runtime graph for on-device inference | | [Python Package](/python/) | Python | Python package and CLI | ## Bindings @@ -213,23 +208,17 @@ graph.hard_reset(); │ --token HuggingFace token (gated models) │ │ --reconvert force local rebuild from source │ │ │ -│ cactus download [model] get a bundle (prebuilt, else build) │ +│ cactus download [model] get a prebuilt bundle │ │ --bits 1|2|3|4 CQ quantization (default: 4) │ │ --platform auto|cpu|apple target platform (default: auto) │ │ --token HuggingFace token (gated models) │ -│ --reconvert force local rebuild from source │ +│ --reconvert refresh cached bundle │ │ │ -│ cactus convert [dir] HuggingFace -> runnable cactus bundle │ -│ (CQ weights + runtime graph) │ +│ cactus convert [dir] HuggingFace -> Cactus CQ weights │ │ --bits 1|2|3|4 CQ quantization (default: 4) │ -│ --platform auto|cpu|apple target platform (default: auto) │ │ --token HuggingFace token (gated models) │ -│ --reconvert force local rebuild from source │ +│ --reconvert force weight conversion from source │ │ --lora merge a LoRA adapter before converting │ -│ --weights-only stop after CQ weights (skip the graph) │ -│ --dynamic-batch dynamic-batch decoder graph (Gemma4) │ -│ --max-slots KV-cache slots for batched decode │ -│ --artifact-dir bundle output (default: weights/) │ │ │ │ cactus serve [model] OpenAI-compatible local HTTP server │ │ --host bind address (default: 127.0.0.1) │ diff --git a/cactus-engine/src/model.cpp b/cactus-engine/src/model.cpp index 37cb4a3a4..da02085c6 100644 --- a/cactus-engine/src/model.cpp +++ b/cactus-engine/src/model.cpp @@ -1352,7 +1352,7 @@ void Model::run_full_context_text() { size_t input_node = static_cast(encoder_->runtime_input_node_ids[input_ids_idx]); const auto& input_desc = encoder_->graph->get_output_buffer(input_node); if (context_tokens_.size() > input_desc.total_size) { - throw std::runtime_error("context exceeds transpiled text_lm_encoder capacity"); + throw std::runtime_error("context exceeds graph-bundle text_lm_encoder capacity"); } std::fill(encoder_->input_buffers[input_ids_idx].begin(), encoder_->input_buffers[input_ids_idx].end(), 0); std::fill(encoder_->input_buffers[attention_mask_idx].begin(), encoder_->input_buffers[attention_mask_idx].end(), 0); @@ -4213,8 +4213,8 @@ std::unique_ptr create_model(const std::string& bundle_dir) { fs::path manifest = fs::path(bundle_dir) / "components" / "manifest.json"; if (!fs::exists(manifest)) { CACTUS_LOG_ERROR("model", - "Not a transpiled bundle (no components/manifest.json at " << bundle_dir << "). " - "Run `cactus convert ` to produce one."); + "Not a runnable Cactus bundle (no components/manifest.json at " << bundle_dir << "). " + "Run `cactus download ` to fetch one."); return nullptr; } return std::make_unique(); diff --git a/cactus-engine/src/model_npu.cpp b/cactus-engine/src/model_npu.cpp index 8c421dee1..31944f5f6 100644 --- a/cactus-engine/src/model_npu.cpp +++ b/cactus-engine/src/model_npu.cpp @@ -212,7 +212,7 @@ bool Model::vision_encode_via_npu(const std::vector& pixel_values, const bool package_takes_positions = npu_vision_encoder_->has_input("pixel_position_ids"); if (package_takes_positions != (pixel_position_ids != nullptr)) { CACTUS_LOG_WARN("model", "NPU vision encoder and pixel_position_ids mismatch; " - "falling back to CPU vision encoder (re-transpile with --npu to fix)"); + "falling back to CPU vision encoder (rebuild the bundle with NPU vision support to fix)"); return false; } diff --git a/cactus-engine/src/transcribe.cpp b/cactus-engine/src/transcribe.cpp index cafd98c67..ae04a7958 100644 --- a/cactus-engine/src/transcribe.cpp +++ b/cactus-engine/src/transcribe.cpp @@ -92,7 +92,7 @@ CACTUS_FFI_EXPORT int cactus_preprocess_audio_features( cactus::engine::Config cfg; cfg.model_type = cactus::engine::Config::ModelType::GEMMA4; cfg.audio_input_feat_size = static_cast(bins); - // The Python transpiled runtime shape-specializes Gemma4 audio to + // The graph-bundle runtime shape-specializes Gemma4 audio to // its requested capacity (currently capped at 30s). Do not apply // the native interactive model's shorter audio_soft_tokens limit // here, otherwise long-file multimodal prompts are silently diff --git a/cactus-engine/tests/test_llm.cpp b/cactus-engine/tests/test_llm.cpp index 9225579c1..78fe7f95d 100644 --- a/cactus-engine/tests/test_llm.cpp +++ b/cactus-engine/tests/test_llm.cpp @@ -876,7 +876,7 @@ bool test_decode_batch() { auto two = handle->model->decode_batch(std::vector(2, seed), M); if (two.empty()) { - std::cout << " [WARN] bundle not dynamic-batch capable (need `cactus convert --dynamic-batch`); skipping\n"; + std::cout << " [WARN] bundle not dynamic-batch capable; skipping\n"; cactus_destroy(model); return true; } @@ -921,7 +921,7 @@ bool test_generate_batch_ragged() { auto batched = handle->model->generate_batch({a, b}, M); if (batched.empty()) { - std::cout << " [WARN] bundle not dynamic-batch capable (need `cactus convert --dynamic-batch`); skipping\n"; + std::cout << " [WARN] bundle not dynamic-batch capable; skipping\n"; cactus_destroy(model); return true; } diff --git a/docs/cactus_engine.md b/docs/cactus_engine.md index 6d2f20d7b..8b8447578 100644 --- a/docs/cactus_engine.md +++ b/docs/cactus_engine.md @@ -29,8 +29,9 @@ graph + manifest) from best build for your host, e.g. Apple Silicon on macOS); pass `--platform cpu` for the generic build. The result can be loaded directly via `cactus_init()`. -For models not on Cactus-Compute, build a bundle from source with -`cactus convert ` (quantizes the weights and builds the runtime graph). +For models not on Cactus-Compute, `cactus convert ` can still quantize +weights to Cactus CQ format, but local runtime bundle generation is unavailable +while the graph builder is being rewritten. ## Types diff --git a/docs/cactus_quants.md b/docs/cactus_quants.md index a656803bd..5b8e96536 100644 --- a/docs/cactus_quants.md +++ b/docs/cactus_quants.md @@ -183,20 +183,18 @@ Three patterns are visible: ## Using CQ Weights -`cactus convert` quantizes the weights to CQ and builds the runtime bundle -(graph + manifest) in one step. Run via `cactus run`: -`cactus convert ` builds a runnable bundle locally; `cactus run` runs a -bundle path or a model id. +`cactus convert` quantizes HuggingFace weights to Cactus CQ format. Runtime bundle +generation is unavailable while the graph builder is being rewritten, so `cactus run` +expects either a local runnable bundle or a prebuilt bundle from Cactus-Compute. ```bash -cactus convert google/gemma-4-E2B-it ./gemma4-bundle -cactus convert google/gemma-4-E2B-it ./gemma4-bundle --bits 2 -cactus run ./gemma4-bundle +cactus convert google/gemma-4-E2B-it ./gemma4-weights +cactus convert google/gemma-4-E2B-it ./gemma4-weights --bits 2 +cactus run google/gemma-4-E2B-it ``` For models on huggingface.co/Cactus-Compute, `cactus run ` (or `cactus -download `) fetches the pre-built bundle, building locally if none is -published. +download `) fetches the pre-built bundle. The CQ matmul kernels live in `cactus-kernels/src/matmul.cpp`. At inference, the kernel applies a Walsh-Hadamard transform to the FP16 activation (not the weight), diff --git a/docs/cactus_transpiler.md b/docs/cactus_transpiler.md deleted file mode 100644 index 86af3f948..000000000 --- a/docs/cactus_transpiler.md +++ /dev/null @@ -1,520 +0,0 @@ ---- -title: "Cactus Transpiler" -description: "Convert any PyTorch model to run on-device with Cactus. Step-by-step guide from a simple PyTorch model to on-device inference." -keywords: ["transpiler", "PyTorch", "torch.export", "on-device AI", "model conversion", "custom models"] ---- - -# Cactus Transpiler - -The Cactus Transpiler takes a PyTorch model and turns it into a Cactus runtime graph -that runs on ARM devices. You write normal PyTorch, and the transpiler handles the rest: -capturing the computation graph, fusing ops, quantizing weights, and producing a bundle -that loads with zero-copy memory mapping. - -## From PyTorch to On-Device in 2 Steps - -### Step 1: Write a PyTorch Model - -Any `torch.nn.Module` works. Here's a small feedforward classifier: - -```python -import torch -import torch.nn as nn - -class TextClassifier(nn.Module): - def __init__(self, vocab_size=32000, embed_dim=256, num_classes=4): - super().__init__() - self.embedding = nn.Embedding(vocab_size, embed_dim) - self.fc1 = nn.Linear(embed_dim, 128) - self.fc2 = nn.Linear(128, num_classes) - - def forward(self, input_ids): - x = self.embedding(input_ids) # (batch, seq, embed_dim) - x = x.mean(dim=1) # (batch, embed_dim) - x = torch.nn.functional.relu(self.fc1(x)) - x = self.fc2(x) - return x - -model = TextClassifier() -torch.save(model.state_dict(), "classifier.pt") -``` - -### Step 2: Convert and Run - -`cactus convert` quantizes the weights to Cactus format (writing a -`weights_manifest.json`) and then captures the graph, optimizes it, and saves a -self-contained bundle — all in one step: - -```bash -cactus convert ./classifier.pt ./classifier-cactus --bits 4 - -# Run the converted bundle -cactus run ./classifier-cactus -``` - -That's it. The output is a self-contained bundle with memory-mapped weights that -loads in milliseconds on any ARM device. - ---- - -## How It Works - -The transpiler is a six-stage compiler pipeline: - -``` -PyTorch nn.Module - | - v -+------------------+ -| 1. Canonicalize | Wrap the model in a task-specific adapter -| model | (causal LM, encoder, multimodal, etc.) -+--------+---------+ - | - v -+------------------+ -| 2. torch.export | Capture the forward pass as an FX graph -| | (no Python control flow, tensor-only) -+--------+---------+ - | - v -+------------------+ -| 3. Import to IR | FX nodes -> canonical IRGraph -| | Bind converted weight files from manifest -+--------+---------+ - | - v -+------------------+ -| 4. Canonicalize | Normalize ops, constant-fold, remove no-ops, -| IR | insert precision casts, dead code elimination -+--------+---------+ - | - v -+------------------+ -| 5. Fuse patterns | Collapse subgraphs into single ops: -| | RMSNorm, RoPE, attention, MLP, LSTM, etc. -+--------+---------+ - | - v -+------------------+ -| 6. Lower to | IRNodes -> CactusGraph ops -| CactusGraph | Constants -> mmap_weights() or inline tensors -+------------------+ - | - v - Saved bundle: graph.cactus + manifest.json + weights -``` - -### Stage 1: Canonicalize the Model - -The transpiler doesn't try to handle every model's forward signature as-is. It wraps -the model in a thin adapter that exposes a stable, task-specific interface: - -| Task | Input | Output | -|------|-------|--------| -| `causal_lm_logits` | `input_ids` | next-token logits | -| `multimodal_causal_lm_logits` | `input_ids`, image/audio features | next-token logits | -| `encoder_hidden_states` | `input_features` | hidden states | -| `ctc_logits` | `input_features` | CTC logits | -| `seq2seq_transcription` | `input_features` | transcription token IDs | -| `tdt_transcription` | `input_features` | TDT token IDs | - -This is where model-family-specific knowledge lives (Gemma4, Qwen, LFM2, Parakeet, etc.). -The adapter also injects `use_cache=False` and `return_dict=False` so the export -boundary stays tensor-only. - -### Stage 2: Capture with `torch.export` - -The adapted model is captured using `torch.export`, producing an FX graph where every -operation is a canonical ATen function call (e.g. `aten.linear.default`, `aten.silu.default`). - -This is the generality boundary. After capture, the transpiler works with ATen op names, -not layer names. It doesn't care whether your model calls its attention layer `self_attn` -or `mha` — it recognizes the computation pattern. - -### Stage 3: Import to IR - -The FX graph is converted into Cactus's intermediate representation (`IRGraph`), a simple -dataflow graph of `IRNode`s connected by `IRValue`s. - -During import, constants that correspond to converted weight files are resolved through -`weights_manifest.json`. This metadata follows the value through the pipeline so that -lowering can use `mmap_weights()` instead of embedding the tensor into the graph. - -### Stage 4: Canonicalize the IR - -Multi-pass cleanup: - -- Rename ops to canonical spellings -- Rewrite `transpose`/`movedim` into `permute` -- Constant-fold trivial expressions -- Materialize `ones`, `zeros`, `arange` constants -- Remove no-op casts, views, slices -- Insert FP16 precision casts where the runtime requires them -- Dead code elimination - -### Stage 5: Fuse Patterns - -The optimizer recognizes computational subgraphs and collapses them into single -high-level ops: - -| Pattern | Fused Op | -|---------|----------| -| `variance -> rsqrt -> mul -> mul` | `rms_norm` | -| `cos/sin interleave on position freqs` | `rope` | -| `Q @ K^T -> scale -> mask -> softmax -> @ V` | `attention` | -| `gate_proj * up_proj -> activation -> down_proj` | `dense_mlp_tq_fused` | -| `conv -> batchnorm -> activation` | `conv_module` | -| `input/forget/cell/output gates` | `lstm_cell` | - -These fused ops map directly to optimized Cactus kernels. - -### Stage 6: Lower to CactusGraph - -Each `IRNode` becomes one or more `CactusGraph` operations. Constants become: - -- `graph.mmap_weights(path)` — zero-copy memory-mapped from disk (most weights) -- Inline scalar values -- Materialized tensor nodes (small constants baked into the graph) - -The output is a `TranspiledGraph` containing the runtime graph, input/output handles, -and weight binding metadata. - ---- - -## Converting HuggingFace Models - -The most common use case is converting a model from HuggingFace. `cactus convert` -quantizes the weights and builds the runtime graph in one step: - -```bash -# 1. Convert (CQ weights + runtime bundle) -cactus convert google/gemma-4-E2B-it ./gemma4-weights --bits 4 - -# 2. Run -cactus run ./gemma4-weights --prompt "Hello!" -``` - -### Text Models - -```bash -cactus convert Qwen/Qwen3-0.6B ./qwen3-weights --bits 4 \ - --prompt "The capital of France is" -``` - -### Multimodal Models (Gemma4) - -Multimodal models are automatically split into components (`vision_encoder`, -`audio_encoder`, `lm_encoder`, `decoder`) that are each independently captured -and optimized: - -```bash -cactus convert google/gemma-4-E2B-it ./gemma4-weights --bits 4 \ - --task multimodal_causal_lm_logits \ - --image-file photo.jpg \ - --audio-file speech.wav \ - --prompt "Describe what you see and hear." -``` - -### Speech Models (Parakeet TDT) - -```bash -cactus convert nvidia/parakeet-tdt-0.6b-v3 ./parakeet-weights --bits 4 \ - --audio-file recording.wav \ - --task tdt_transcription -``` - ---- - -## Running Saved Bundles - -Transpiled bundles are saved to `./weights//` by default (alongside their CQ -weights). Run them later without re-transpiling: - -```bash -# Text -cactus run ./weights/qwen3-0.6b --prompt "Write a haiku" - -# Multimodal -cactus run ./weights/gemma-4-e2b-it \ - --image photo.jpg \ - --audio speech.wav \ - --prompt "What do you see?" - -# Audio -cactus run ./weights/parakeet-tdt-0.6b-v3 \ - --audio meeting.wav -``` - -`cactus run` accepts either a HuggingFace model id (downloads the bundle) or a -local bundle directory path (detected by the presence of -`components/manifest.json`). The standalone `run-transpiled` subcommand was -removed — `cactus run` handles both cases. - ---- - -## Using the Python API Directly - -You can also drive the transpiler from Python: - -```python -from cactus.transpile.capture_pytorch import capture_model -from cactus.transpile.lower import transpile_ir -from cactus.bindings.cactus import Graph -import torch -import numpy as np - -# Define your model -class MyModel(torch.nn.Module): - def __init__(self): - super().__init__() - self.linear = torch.nn.Linear(64, 32) - - def forward(self, x): - return torch.relu(self.linear(x)) - -model = MyModel().eval() -example_input = torch.randn(1, 64) - -# Capture and transpile -captured = capture_model(model, (example_input,)) -transpiled = transpile_ir(captured.ir_graph) - -# Run -transpiled.set_input(0, np.random.randn(1, 64).astype(np.float16)) -outputs = transpiled.execute() -result = outputs[0].numpy() -print("Output shape:", result.shape) # (1, 32) -``` - -### Generic JAX User Graphs - -JAX and Flax models can be bundled by providing params and one or more graph -entrypoints. The caller owns tokenization, masks, sampling, and any prefill/decode -loop; Cactus captures the supplied functions, writes FP16 mmap weights, and saves -component graphs. - -```python -import jax.numpy as jnp -import numpy as np - -from cactus.transpile.capture_jax import JaxGraphSpec -from cactus.transpile.jax_user_graph_bundle import build_jax_user_graph_bundle -from cactus.transpile.jax_user_graph_bundle import load_jax_user_graph_bundle - - -def model(params, x): - return jnp.maximum(x @ params["w"] + params["b"], 0) - - -params = {"w": jnp.ones((4, 3), jnp.float16), "b": jnp.zeros((3,), jnp.float16)} -example_x = jnp.ones((2, 4), jnp.float16) - -result = build_jax_user_graph_bundle( - params=params, - output_dir="/tmp/my-jax-bundle", - model_id="my-jax-model", - specs=(JaxGraphSpec(name="forward", fn=model, example_args=(example_x,)),), -) - -loaded = load_jax_user_graph_bundle(result.output_dir) -y = loaded.execute("forward", np.ones((2, 4), np.float16))[0].numpy() -``` - -For encoder-decoder models, provide each graph boundary explicitly: - -```python -params, model, tokenizer, config = load_needle_model() -src = tokenizer.encode("What time is it in Tokyo?") -tgt = tokenizer.encode("") -src_mask = make_padding_mask(src, config.pad_token_id) -tgt_mask = make_causal_mask(tgt.shape[1]) -encoder_out, enc_mask = model.apply({"params": params}, src, src_mask, method=model.encode_text) - -result = build_jax_user_graph_bundle( - params=params, - output_dir="weights/needle", - model_id="needle", - task="encoder-decoder", - specs=( - JaxGraphSpec( - name="encoder", - fn=lambda params, src, src_mask: model.apply( - {"params": params}, src, src_mask, method=model.encode_text - ), - example_args=(src, src_mask), - output_names=("encoder_out", "encoder_mask"), - ), - JaxGraphSpec( - name="decoder_prefill", - fn=lambda params, tgt, encoder_out, tgt_mask, cross_mask: model.apply( - {"params": params}, tgt, encoder_out, tgt_mask, cross_mask, method=model.decode - ), - example_args=(tgt, encoder_out, tgt_mask, enc_mask), - output_names=("logits",), - ), - ), -) - -loaded = load_jax_user_graph_bundle(result.output_dir) -encoder_out, enc_mask = loaded.execute("encoder", src, src_mask) -logits = loaded.execute("decoder_prefill", tgt, encoder_out, tgt_mask, enc_mask)[0].numpy() -next_token = int(np.argmax(logits[0, -1])) -``` - -Each component is saved under `components//` with `graph.cactus`, -`raw_ir.json`, and `optimized_ir.json`. For decode, expose a separate -`decoder_step` graph. The generic JAX path does not create or manage internal -KV cache yet; component output `Tensor`s can be passed directly into another -component. Use `loaded.reset()` between independent generations. - ---- - -## Artifact Layout - -A transpiled bundle looks like this: - -``` -weights// - raw_ir.json # IR before optimization (debugging) - optimized_ir.json # IR after fusion passes (debugging) - graph.cactus # serialized runtime graph - graph_bindings.json # weight binding metadata - result.json # execution results (if --execute-after-transpile) - components/ # for multimodal models - manifest.json # component order, input/output names - vision_encoder/ - graph.cactus - bound_constants/ - audio_encoder/ - graph.cactus - lm_encoder/ - graph.cactus - decoder/ - graph.cactus -``` - -The `manifest.json` tells the runtime loader how to chain the components and where -to find the memory-mapped weight files. The raw/optimized IR JSONs are also valid -reload targets — the runtime can re-lower from saved IR as a fallback. - ---- - -## CLI Reference - -### `cactus convert` - -```bash -cactus convert [output-dir] [options] -``` - -Quantizes the weights to CQ and builds the runtime graph. Pass `--weights-only` to -stop after the CQ weights. The graph-build options below are the same ones the -transpiler accepts: - -| Option | Description | -|--------|-------------| -| `--bits 1\|2\|3\|4` | CQ quantization bits (default: 4) | -| `--weights-only` | Stop after CQ quantization; skip the runtime graph | -| `--dynamic-batch` | Emit a dynamic-batch decoder graph for batched/continuous decode (Gemma4) | -| `--max-slots ` | KV-cache slot-pool capacity for batched decode (with `--dynamic-batch`) | -| `--weights-dir ` | Path to converted CQ weights (default: `weights/`) | -| `--task ` | Force task type (default: `auto` — inferred from model config). Choices: `causal_lm_logits`, `multimodal_causal_lm_logits`, `ctc_logits`, `encoder_hidden_states`, `seq2seq_transcription`, `tdt_transcription` | -| `--prompt ` | Representative prompt for shape capture | -| `--image-file ` | Representative image (repeatable) | -| `--audio-file ` | Representative audio file | -| `--max-new-tokens ` | Generation room to preallocate for causal decode graphs | -| `--artifact-dir ` | Output directory (default: `weights/`) | -| `--execute-after-transpile` | Run the graph after saving | -| `--skip-reference-compare` | Skip PyTorch vs transpiled comparison | -| `--component-pipeline auto\|on\|off` | Use split component graph transpilation when supported | -| `--components ` | Comma-separated component subset for component-pipeline models | -| `--trust-remote-code` | Allow HF remote code during the build | -| `--local-files-only` | Require HF files to already be local | -| `--allow-unconverted-weights` | Build against an unconverted source checkpoint | -| `--no-fuse-rms-norm` | Disable RMSNorm fusion | -| `--no-fuse-rope` | Disable RoPE fusion | -| `--no-fuse-attention` | Disable attention fusion | - -NPU emission (CoreML `.mlpackage`s for Apple Silicon audio and vision encoders) -is available through the `--npu`, `--npu-quantize`, `--npu-audio-quantize`, and -`--npu-vision-quantize` flags on `cactus convert`. - -### `cactus run` - -```bash -cactus run [options] -``` - -`model_id` may be a HuggingFace model id (downloads the matching bundle from -huggingface.co/Cactus-Compute) or a local path to a bundle directory. - -| Option | Description | -|--------|-------------| -| `--bits 1\|2\|3\|4` | CQ quantization bits when downloading (default: 4) | -| `--platform auto\|cpu\|apple` | Target platform (default: auto — best for the host) | -| `--token ` | HuggingFace token (gated models) | -| `--prompt ` | Input prompt | -| `--input-ids ` | Comma-separated token IDs | -| `--image ` | Image file for multimodal bundles | -| `--audio ` | Audio file for speech / audio bundles | -| `--system ` | System prompt | -| `--thinking` | Enable thinking / reasoning mode | -| `--max-new-tokens ` | Max tokens to generate | -| `--result-json ` | Save bundle result payload to JSON | -| `--no-cloud-handoff` | Disable automatic cloud handoff | -| `--confidence-threshold ` | Confidence threshold below which to hand off to cloud | -| `--cloud-timeout-ms ` | Max wait time for cloud handoff | -| `--no-cloud-tele` | Disable cloud telemetry (write to cache only) | -| `--reconvert` | Force local rebuild from source | - ---- - -## Supported Model Families - -The transpiler works with any model that `torch.export` can capture, but has -tested adapters for: - -| Family | Tasks | Component Split | -|--------|-------|-----------------| -| Gemma 3/4 | text, multimodal | yes (vision + audio + LM + decoder) | -| Qwen 3/3.5 | text, vision | no | -| LFM2/2.5 | text, vision | no | -| Whisper | encoder | no | -| Parakeet CTC | encoder | no | -| Parakeet TDT | transcription | yes (encoder + decoder) | - -Adding a new model family means writing a small adapter in `model_adapters.py` that -wraps the HF forward into the canonical task interface. The rest of the pipeline -(capture, optimize, lower) is model-agnostic. - ---- - -## Design Decisions - -**ATen ops, not layer names.** After `torch.export`, the transpiler dispatches on -canonical ATen operation names (`aten.linear.default` -> `linear`). It never reads -HuggingFace module paths like `model.layers.0.self_attn`. This means the fusion -passes work on any model that produces the same computational pattern, regardless -of how the code is organized. - -**Weight binding through manifests.** The transpiler resolves weights exclusively -through `weights_manifest.json` written by `cactus convert`. It does not guess -filenames from layer names. This is intentional — it keeps the contract between -conversion and transpilation explicit and auditable. - -**Component splitting for multimodal.** Complex models like Gemma4 are split into -independent components (vision encoder, audio encoder, LM encoder, decoder) before -capture. Each component goes through the full optimize/lower pipeline independently. -This keeps each graph small enough for `torch.export` to handle and lets the runtime -load/unload components independently. - -**Saved IR as fallback.** Bundles include both the serialized `.cactus` graph and -the raw/optimized IR JSON. If the binary graph format changes between versions, the -runtime can re-lower from saved IR instead of requiring a full re-transpile. - -## See Also - -- [Cactus Engine API](/docs/cactus_engine.md) — C API for inference -- [Cactus Graph API](/docs/cactus_graph.md) — Computation graph reference -- [Fine-tuning Guide](/docs/finetuning.md) — Train LoRA fine-tunes and deploy to mobile -- [Transpiler source](https://github.com/cactus-compute/cactus/tree/main/python/cactus/transpile) — Full implementation diff --git a/docs/finetuning.md b/docs/finetuning.md index 539aea5a1..e38a86b90 100644 --- a/docs/finetuning.md +++ b/docs/finetuning.md @@ -67,13 +67,10 @@ cactus convert Qwen/Qwen3-0.6B ./my-qwen3-0.6b --lora username/my-lora-adapter ### 4. Run -Build the native library, then run (the convert above already produced the runtime -bundle alongside the CQ weights): - -```bash -cactus build -cactus run ./my-qwen3-0.6b -``` +The commands above currently produce CQ weights only. Local runtime bundle +generation is unavailable while the graph builder is being rewritten, so deploy the +fine-tuned model after the replacement graph builder lands, or run a prebuilt +Cactus-Compute bundle with `cactus run `. `cactus run` accepts the bundle path directly, or auto-builds from a HF model id when no bundle exists locally. diff --git a/list.txt b/list.txt new file mode 100644 index 000000000..c9bf1c52f --- /dev/null +++ b/list.txt @@ -0,0 +1,152 @@ +Current transpiler feature inventory + +- End-to-end PyTorch-to-Cactus pipeline: + canonicalize a model interface, capture with torch.export, import exported FX + graphs into IRGraph, canonicalize/optimize the IR, fuse recognized patterns, + lower to a Cactus Graph, and save runnable artifacts. + +- HuggingFace model loading: + task inference from config/model id, tokenizer/processor loading, optional + trust_remote_code/local-files-only handling, dtype selection, HF token support, + converted weight directory validation, and model-family aliases. + +- Supported task shapes: + causal_lm_logits, multimodal_causal_lm_logits, encoder_hidden_states, + ctc_logits, seq2seq_transcription, tdt_transcription, and text_embedding. + +- Tested/profiled model families: + Gemma 3/4, Qwen 3/3.5 including vision variants, LFM2/LFM2-VL/LFM2.5, + Whisper, Parakeet CTC, Parakeet TDT, Nomic text embedding, and Needle-style + encoder/decoder generation graphs. + +- Model interface adapters: + causal LM wrappers, multimodal causal LM wrappers, audio encoder wrappers, + CTC wrappers, seq2seq transcription wrappers, TDT transcription wrappers, + text embedding wrappers, cache-aware decoder wrappers, and family-specific + prompt/input normalization. + +- Representative input preparation: + prompt/input-id shaping, max-new-token context preallocation, audio WAV loading + and log-mel/native feature extraction, image resizing/limits, multimodal prompt + construction, system prompt support, thinking marker support, and default tiny + test media for shape capture. + +- Component pipeline support: + monolithic graph capture plus split component capture for models that need it, + including vision_encoder, audio_encoder, lm_encoder, decoder, decoder_step, + lm_encoder_step, decoder_prefill_chunk, decoder_embed_chunk, + decoder_media_step, source_encoder, decoder_cross_kv, text_embedding, and + model-specific component subsets. + +- Dynamic/cache behavior: + cached decoder-step graphs, internal KV cache metadata, shared KV cache routes, + cross-attention KV routes, configurable cache context length, dynamic-batch + Gemma4 decoder graph emission, max-slots slot pool sizing, FP16 cache component + hints, and reset support between generations. + +- JAX and Flax user graph support: + capture_jax_function, capture_jax_function_with_params, JaxGraphSpec, + build_jax_user_graph_bundle, build_jax_generation_graph_bundle, saved component + manifests/graphs/weights, loading saved JAX bundles, external weights dirs, + decoder-step attention rewrites, and encoder-decoder graph boundaries. + +- IR features: + IRValue, IRNode, IRGraph, graph metadata, lifted constants, graph verification, + shape/dtype metadata, constant tables, imported component annotations, opaque + node preservation for unknown ops, saved raw_ir.json and optimized_ir.json, and + re-lowering from saved IR. + +- Weight binding: + explicit weights_manifest.json resolution, source-name aliases, mmap-backed + constants, graph_bindings.json, per-component bound constant metadata, + materialized constant tensors, scalar constants, embedding binding reuse, + compatibility companion generation for INT8/CQ embeddings and legacy packed + INT4/CQ4 files, and debug logging for mmap decisions. + +- Canonicalization and cleanup: + target/op spelling normalization, dtype normalization, transpose/movedim to + permute rewrites, shape-only op cleanup, constant folding, materialization of + ones/zeros/arange, no-op cast/view/slice removal, precision legalization casts, + graph rebuild, dead-code elimination, unsupported-op counts, and repeated + cleanup passes. + +- Import-time semantic rewrites: + scaled_dot_product_attention to semantic attention, explicit linear tagging, + early RoPE recognition, import hints, nested exported subgraph inlining, and + ATen/op-based behavior instead of HuggingFace layer-name dispatch. + +- Imported ATen/export operations: + arange, identity/clone/alias/detach/dropout, add/subtract/multiply/divide, + floor_divide, comparisons, logical ops, where, masked_fill, masked_scatter, + abs, clamp, neg, exp, sqrt, reciprocal, log, rsqrt, cos, sin, reshape/view, + flatten, squeeze/unsqueeze, expand, repeat, one_hot, tril, unfold, transpose, + numpy_T, permute, movedim, matmul/mm/bmm, linear, addmm, relu, silu, gelu, + gelu_erf, sigmoid, glu, softplus, tanh, softmax, scaled dot-product attention, + sum, mean, variance, min/max/amin/amax, cat, split_with_sizes, chunk, unbind, + ones, pad, pow, diff, slice, select/index, gather, embedding, conv1d, conv2d, + layer_norm, rms_norm, group_norm, batch_norm, contiguous, and getitem. + +- Lowered Cactus graph operations: + elementwise math, scalar math, comparisons, logical ops, where, precision casts, + layout/view/reshape/flatten/squeeze/unsqueeze/expand/repeat/permute/transpose, + one_hot, tril, unfold, pad, matmul/linear/addmm, dense_mlp_tq_fused, reductions, + activations, softmax, LSTM cell, attention, scaled_dot_product_attention, + attention_block, self_attention_block, cumsum, cat/split/chunk, slice, index, + gather, embedding_from_tensor, conv_module, conv1d, conv2d, layer_norm, + rms_norm, rope, rel_pos_bias, group_norm, batch_norm, advanced_index, + masked_fill, and getitem. + +- Fusion/optimization passes: + RMSNorm, RMSNorm scale multiply, RoPE, attention, attention_block, + self_attention_block, dense MLP/gated MLP, conv module, LSTM cell, + relative-position bias, add_clipped, gated DeltaNet prefill/decode, linear + projection matching, gold-pattern annotation, unused-input pruning, and Gemma4 + decoder attention semantic normalization. + +- Special model-family behavior: + Gemma4 multimodal component splitting, Gemma4 checkpoint key repair, Gemma4 + prompt/token drift normalization, Gemma4 sliding-window/window-size attention + handling, Gemma4 decoder attention lowering fallback, Qwen/LFM multimodal + preprocessing, Parakeet TDT local runtime reconstruction and greedy decode, + Whisper prompt token handling, Nomic embedding component planning, and Needle + encoder/cross-KV/decoder-step planning. + +- Optional execution and validation: + skip-execute mode, execute-after-transpile mode, PyTorch reference comparison, + skip-reference-compare, next-token reporting for text/multimodal LMs, transcript + reporting for Parakeet TDT/CTC paths, generic audio output-shape reporting, and + result.json writing. + +- Saved artifact layout: + graph.cactus, graph_bindings.json, raw_ir.json, optimized_ir.json, result.json, + components/manifest.json, per-component graph.cactus files, per-component + bound_constants directories, logical input/output metadata, runtime node ids, + cache state node ids, and weight binding metadata. + +- CLI/API surface: + internal cactus.cli.transpile.run_transpile, python -m cactus.transpile.hf_model, + cactus convert graph-build forwarding, --task, --prompt, --system-prompt, + --enable-thinking, --input-ids, --image-file, --audio-file, --max-new-tokens, + --weights-dir, --artifact-dir, --graph-filename, --component-pipeline, + --components, --torch-dtype, --trust-remote-code, --local-files-only, + --allow-unconverted-weights, --dynamic-batch, --max-slots, --cache-context-length, + --skip-execute, --skip-reference-compare, --execute-after-transpile, and fusion + disable flags. + +- Apple NPU/CoreML emission: + --npu, --npu-quantize, --npu-audio-quantize, --npu-vision-quantize, audio encoder + mlpackage export, vision encoder mlpackage export, source encoder export, + CoreMLTools compatibility patches, per-component quantization defaults, and + integration into component bundle manifests. + +- Runtime compatibility helpers: + Python Graph wrapper monkey patches for missing current-runtime symbols, generic + matmul fallback, gather(axis=0), limited conv2d, not_equal, scalar math + legalization, Transformers/Gemma4 import support, torchvision/_lzma/nms import + patches, and flex_attention AuxRequest compatibility. + +- Runtime bundle execution support: + saved-bundle execution for causal_lm_logits, multimodal_causal_lm_logits, + encoder_hidden_states, and tdt_transcription; ctc_logits transpilation without a + saved-bundle executor; native run-binary handoff through components/manifest.json; + local bundle detection; and Cactus-Compute download fallback integration. diff --git a/mkdocs.yml b/mkdocs.yml index b7bd930cd..867ffc35d 100644 --- a/mkdocs.yml +++ b/mkdocs.yml @@ -98,7 +98,6 @@ nav: - Index API: docs/cactus_index.md - Cactus Quants: docs/cactus_quants.md - Hybrid Inference: docs/cactus_hybrid.md - - Transpiler: docs/cactus_transpiler.md - Guides: - Fine-tuning & Deployment: docs/finetuning.md - Runtime Compatibility: docs/compatibility.md diff --git a/python/README.md b/python/README.md index a259aa8ce..8a278f827 100644 --- a/python/README.md +++ b/python/README.md @@ -370,8 +370,8 @@ python python/test.py # compact output python python/test.py -v # verbose ``` -Tests are in `python/tests/` — bindings, CLI, server, graph, model, transpile, -and component-partition coverage. Add a new `test_*.py` to extend. +Tests are in `python/tests/` — bindings, CLI, server, graph, and model coverage. +Add a new `test_*.py` to extend. ## See Also diff --git a/python/cactus/cli/__init__.py b/python/cactus/cli/__init__.py index d5cb6c6bc..eea9a875d 100644 --- a/python/cactus/cli/__init__.py +++ b/python/cactus/cli/__init__.py @@ -23,7 +23,6 @@ from .transcribe import cmd_transcribe from .test import cmd_test, COMPONENTS from .convert import cmd_convert -from .upload import cmd_upload from .run import cmd_run from .list import cmd_list @@ -40,7 +39,7 @@ def _telemetry_parent(): def _build_parent(): - """Bundle-build flags shared by every command that prepares a model.""" + """Bundle-fetch flags shared by every command that prepares a model.""" p = argparse.ArgumentParser(add_help=False) p.add_argument("--bits", type=int, choices=[1, 2, 3, 4], default=4, help="CQ quantization (default: 4)") @@ -48,7 +47,7 @@ def _build_parent(): help=_PLATFORM_HELP) p.add_argument("--token", help="HuggingFace token") p.add_argument("--reconvert", action="store_true", - help="Force local rebuild from source") + help="Refresh the cached bundle or weights") return p @@ -117,7 +116,7 @@ def create_parser(): --prompt send prompt immediately --thinking enable thinking/reasoning mode --token HuggingFace token (gated models) - --reconvert force local convert fallback + --reconvert refresh cached bundle cactus transcribe [model] live microphone transcription with a model --file audio file to transcribe (WAV) @@ -125,19 +124,18 @@ def create_parser(): --bits 1|2|3|4 CQ quantization (default: 4) --platform {_PLATFORM_PIPE:<22} target platform (default: auto) --token HuggingFace token (gated models) - --reconvert force local rebuild from source + --reconvert refresh cached bundle - cactus download [model] fetch a prebuilt bundle, else build locally (default: {DEFAULT_MODEL_ID}) + cactus download [model] fetch a prebuilt bundle (default: {DEFAULT_MODEL_ID}) --bits 1|2|3|4 CQ quantization (default: 4) --platform {_PLATFORM_PIPE:<22} target platform (default: auto) --token HuggingFace token (gated models) - --reconvert force local rebuild from source + --reconvert refresh cached bundle - cactus convert [dir] build a runnable bundle locally (skips prebuilt fetch) + cactus convert [dir] convert HuggingFace weights to Cactus CQ --bits 1|2|3|4 CQ quantization (default: 4) - --platform {_PLATFORM_PIPE:<22} target platform (default: auto) --token HuggingFace token (gated models) - --reconvert force local rebuild from source + --reconvert force weight conversion from source --lora merge a LoRA adapter before converting cactus serve [model] OpenAI-compatible local HTTP server @@ -146,7 +144,7 @@ def create_parser(): --bits 1|2|3|4 CQ quantization (default: 4) --platform {_PLATFORM_PIPE:<22} target platform (default: auto) --token HuggingFace token (gated models) - --reconvert force local rebuild from source + --reconvert refresh cached bundle --no-cloud-handoff disable automatic cloud handoff --confidence-threshold <0..1> handoff to cloud below this confidence --cloud-timeout-ms max wait for cloud handoff before local fallback @@ -166,7 +164,7 @@ def create_parser(): --bits 1|2|3|4 CQ quantization (default: 4) --platform {_PLATFORM_PIPE:<22} target platform (default: auto) --token HuggingFace token (gated models) - --reconvert force local rebuild of test models + --reconvert refresh cached test bundles --suite run a single test suite from any component (kernels, graph, or engine) --list list components and suites @@ -180,49 +178,13 @@ def create_parser(): ----------------------------------------------------------------- - Advanced / build pipeline — convert runs automatically inside run, - serve, transcribe and download; reach for it only to control the - build (custom flags, LoRA, NPU, debugging): + Advanced: - cactus convert [dir] HuggingFace -> runnable cactus bundle - (quantizes weights to CQ, then builds - the runtime graph) + cactus convert [dir] HuggingFace -> Cactus CQ weights --bits 1|2|3|4 CQ quantization (default: 4) --token HuggingFace token (defaults to $HF_TOKEN) --reconvert force weight conversion from source --lora merge a LoRA adapter before converting - --weights-only stop after CQ weights (skip the graph) - --dynamic-batch emit a dynamic-batch decoder graph (Gemma4) - --max-slots KV-cache slot capacity for batched decode - --weights-dir path to CQ weights (default: weights/) - --task task (default: auto, inferred from config) - --artifact-dir write bundle here (default: weights/) - --prompt representative prompt for shape capture - --system-prompt system prompt for multimodal chat - --enable-thinking enable thinking markers when supported - --input-ids token ids for causal-LM shape capture - --image-file representative image (repeatable) - --audio-file representative audio file (WAV) - --max-new-tokens preallocate decode context for causal LM - --component-pipeline auto|on|off force component-pipeline graph build - --components subset of components to build - --torch-dtype float16 | float32 | bfloat16 - --trust-remote-code allow HF remote code during the build - --local-files-only require model/processor to be local - --allow-unconverted-weights debug-only: skip the CQ-weights check - --execute-after-transpile run a reference execution after building - --graph-filename override saved graph filename - --skip-reference-compare skip PyTorch comparison (with --execute-…) - --no-fuse-rms-norm disable RMSNorm fusion - --no-fuse-rope disable RoPE fusion - --no-fuse-attention disable attention fusion - --no-fuse-attention-block disable attention-block fusion - --no-fuse-add-clipped disable add-clipped fusion - --no-fuse-gated-deltanet disable gated DeltaNet fusion - --npu also emit CoreML .mlpackage(s) for NPU - --npu-quantize 0|4|8 force both NPU encoders to this quant - --npu-audio-quantize 0|4|8 audio encoder quant (default int8) - --npu-vision-quantize 0|4|8 vision encoder quant (default fp16) ----------------------------------------------------------------- """ @@ -354,77 +316,6 @@ def create_parser(): help="Path or HF id of a LoRA adapter to merge before converting (requires `peft`)") convert_parser.add_argument("--reconvert", action="store_true", help="Force conversion from source") - convert_parser.add_argument("--weights-only", action="store_true", - help="Only quantize weights; skip building the runtime graph bundle") - convert_parser.add_argument("--dynamic-batch", action="store_true", - help="Emit a dynamic-batch decoder graph for batched/continuous decode (Gemma4)") - convert_parser.add_argument("--max-slots", type=int, default=1, - help="KV-cache slot-pool capacity for batched decode; used with --dynamic-batch") - - convert_parser.add_argument("--weights-dir", - help="CQ weights directory (default: weights/)") - convert_parser.add_argument("--task", default="auto", - choices=["auto", "causal_lm_logits", "multimodal_causal_lm_logits", - "ctc_logits", "encoder_hidden_states", - "seq2seq_transcription", "tdt_transcription"], - help="Graph task (default: auto, from model config)") - convert_parser.add_argument("--prompt", - help="Prompt for causal/multimodal shape capture") - convert_parser.add_argument("--system-prompt", default=None, - help="System prompt for multimodal chat formats") - convert_parser.add_argument("--enable-thinking", action="store_true", - help="Enable thinking markers when the prompt supports them") - convert_parser.add_argument("--input-ids", default=None, - help="Comma-separated token ids for causal-LM shape capture") - convert_parser.add_argument("--image-file", action="append", default=[], - help="Image for multimodal shape capture (repeatable)") - convert_parser.add_argument("--audio-file", - help="Audio file (WAV) for audio/multimodal shape capture") - convert_parser.add_argument("--max-new-tokens", type=_positive_int, default=None, - help="Decode context to preallocate for causal LM (default: 32)") - convert_parser.add_argument("--component-pipeline", default="auto", choices=["auto", "on", "off"], - help="Split-component graph when supported (default: auto)") - convert_parser.add_argument("--components", - help="Comma-separated component subset (e.g. vision_encoder,decoder)") - convert_parser.add_argument("--torch-dtype", default=None, - choices=["float16", "float32", "bfloat16"], - help="Torch dtype for HF loading (default: float16)") - convert_parser.add_argument("--trust-remote-code", action="store_true", - help="Pass trust_remote_code=True to HF loaders") - convert_parser.add_argument("--local-files-only", action="store_true", - help="Require model/processor to already be local") - convert_parser.add_argument("--allow-unconverted-weights", action="store_true", - help="Debug: build the graph without CQ weights") - convert_parser.add_argument("--execute-after-transpile", action="store_true", - help="Run a reference execution after building the graph") - convert_parser.add_argument("--artifact-dir", - help="Graph bundle output directory (default: weights/)") - convert_parser.add_argument("--graph-filename", default=None, - help="Saved graph filename (default: graph.cactus)") - convert_parser.add_argument("--skip-reference-compare", action="store_true", - help="Skip PyTorch comparison (requires --execute-after-transpile)") - convert_parser.add_argument("--no-fuse-rms-norm", action="store_true", - help="Disable RMSNorm fusion") - convert_parser.add_argument("--no-fuse-rope", action="store_true", - help="Disable RoPE fusion") - convert_parser.add_argument("--no-fuse-attention", action="store_true", - help="Disable attention fusion") - convert_parser.add_argument("--no-fuse-attention-block", action="store_true", - help="Disable attention-block fusion") - convert_parser.add_argument("--no-fuse-add-clipped", action="store_true", - help="Disable add-clipped fusion") - convert_parser.add_argument("--no-fuse-gated-deltanet", action="store_true", - help="Disable gated DeltaNet fusion") - convert_parser.add_argument("--npu", action="store_true", - help="Also emit CoreML .mlpackage(s) for Apple NPU encoders") - convert_parser.add_argument("--npu-quantize", type=int, choices=[0, 4, 8], default=None, - help="Legacy: force both NPU encoders to same quant (0=fp16, 4=int4, 8=int8)") - convert_parser.add_argument("--npu-audio-quantize", type=int, choices=[0, 4, 8], default=None, - help="NPU audio encoder quant: 0=fp16, 4=int4, 8=int8 (default: 8)") - convert_parser.add_argument("--npu-vision-quantize", type=int, choices=[0, 4, 8], default=None, - help="NPU vision encoder quant: 0=fp16, 4=int4, 8=int8 (default: 0; int4 degrades Gemma4 vision)") - convert_parser.add_argument("--cache-context-length", default=None, - help="KV cache context length for cached decode graphs (default: model config)") return parser @@ -442,7 +333,6 @@ def create_parser(): "auth": cmd_auth, "clean": cmd_clean, "convert": cmd_convert, - "upload": cmd_upload, } diff --git a/python/cactus/cli/clean.py b/python/cactus/cli/clean.py index be21a9e21..9fe7311b2 100644 --- a/python/cactus/cli/clean.py +++ b/python/cactus/cli/clean.py @@ -17,7 +17,6 @@ def cmd_clean(args): destructive_targets = [ PROJECT_ROOT / "weights", - PROJECT_ROOT / "transpiled", PROJECT_ROOT / "venv", ] present = [t for t in destructive_targets if t.exists()] @@ -55,7 +54,6 @@ def remove_if_exists(path): remove_if_exists(PROJECT_ROOT / "venv") remove_if_exists(PROJECT_ROOT / "weights") - remove_if_exists(PROJECT_ROOT / "transpiled") print() print("Removing compiled libraries and frameworks...") diff --git a/python/cactus/cli/common.py b/python/cactus/cli/common.py index 71601e4cb..fac99ed54 100644 --- a/python/cactus/cli/common.py +++ b/python/cactus/cli/common.py @@ -27,12 +27,6 @@ def weights_root() -> Path: return Path.home() / ".cache" / "cactus" / "weights" -def transpiled_root() -> Path: - if is_repo_checkout(): - return PROJECT_ROOT / "transpiled" - return Path.home() / ".cache" / "cactus" / "transpiled" - - DEFAULT_MODEL_ID = "google/gemma-4-E2B-it" DEFAULT_TRANSCRIPTION_MODEL_ID = "nvidia/parakeet-tdt-0.6b-v3" DEFAULT_TEST_MODEL_ID = "LiquidAI/LFM2-VL-450M" diff --git a/python/cactus/cli/convert.py b/python/cactus/cli/convert.py index 24b22e2df..f86f8ab68 100644 --- a/python/cactus/cli/convert.py +++ b/python/cactus/cli/convert.py @@ -42,8 +42,10 @@ def _merge_lora_adapter(base_model_id, lora_path, token=None): def cmd_convert(args): - """Convert a HuggingFace model into a runnable cactus bundle: quantize weights to CQ, - then build the runtime graph. Pass --weights-only to stop after the weight conversion. + """Convert a HuggingFace model into Cactus CQ weights. + + Runtime graph/bundle generation used to happen after this step, but that graph + builder has been removed for a rewrite. """ from .model import ensure_weights @@ -66,101 +68,10 @@ def cmd_convert(args): reconvert=args.reconvert, output_dir=output_dir, ) - if getattr(args, "weights_only", False): - return 0 - args.weights_dir = args.weights_dir or output_dir - args.artifact_dir = args.artifact_dir or output_dir - return cmd_transpile(args) + return 0 except RuntimeError as e: print_color(RED, f"Conversion error: {e}") return 1 finally: if merged_dir: shutil.rmtree(merged_dir, ignore_errors=True) - - -def cmd_transpile(args): - """Build the runtime graph bundle from already-converted CQ weights.""" - from .transpile import run_transpile - from .model import _default_multimodal_assets - - extra_args = [] - if args.weights_dir: - extra_args.extend(["--weights-dir", args.weights_dir]) - if args.task and args.task != "auto": - extra_args.extend(["--task", args.task]) - if args.prompt is not None: - extra_args.extend(["--prompt", args.prompt]) - if args.system_prompt is not None: - extra_args.extend(["--system-prompt", args.system_prompt]) - if args.enable_thinking: - extra_args.append("--enable-thinking") - if args.input_ids is not None: - extra_args.extend(["--input-ids", args.input_ids]) - - image_files = list(args.image_file or []) - audio_file = args.audio_file - if not image_files or not audio_file: - default_images, default_audio = _default_multimodal_assets() - if not image_files: - image_files = default_images - if not audio_file and default_audio: - audio_file = default_audio - - for img in image_files: - extra_args.extend(["--image-file", img]) - if audio_file: - extra_args.extend(["--audio-file", audio_file]) - if args.max_new_tokens is not None: - extra_args.extend(["--max-new-tokens", str(args.max_new_tokens)]) - if args.component_pipeline and args.component_pipeline != "auto": - extra_args.extend(["--component-pipeline", args.component_pipeline]) - if args.components: - extra_args.extend(["--components", args.components]) - if args.torch_dtype: - extra_args.extend(["--torch-dtype", args.torch_dtype]) - if args.token: - extra_args.extend(["--token", args.token]) - if args.trust_remote_code: - extra_args.append("--trust-remote-code") - if args.local_files_only: - extra_args.append("--local-files-only") - if args.artifact_dir: - extra_args.extend(["--artifact-dir", args.artifact_dir]) - if args.graph_filename: - extra_args.extend(["--graph-filename", args.graph_filename]) - if args.skip_reference_compare: - extra_args.append("--skip-reference-compare") - if args.no_fuse_rms_norm: - extra_args.append("--no-fuse-rms-norm") - if args.no_fuse_rope: - extra_args.append("--no-fuse-rope") - if args.no_fuse_attention: - extra_args.append("--no-fuse-attention") - if args.no_fuse_attention_block: - extra_args.append("--no-fuse-attention-block") - if args.no_fuse_add_clipped: - extra_args.append("--no-fuse-add-clipped") - if args.no_fuse_gated_deltanet: - extra_args.append("--no-fuse-gated-deltanet") - if args.npu: - extra_args.append("--npu") - if args.npu_quantize is not None: - extra_args.extend(["--npu-quantize", str(args.npu_quantize)]) - if args.npu_audio_quantize is not None: - extra_args.extend(["--npu-audio-quantize", str(args.npu_audio_quantize)]) - if args.npu_vision_quantize is not None: - extra_args.extend(["--npu-vision-quantize", str(args.npu_vision_quantize)]) - if args.cache_context_length is not None: - extra_args.extend(["--cache-context-length", str(args.cache_context_length)]) - if getattr(args, "dynamic_batch", False): - extra_args.append("--dynamic-batch") - if getattr(args, "max_slots", 1) and int(args.max_slots) != 1: - extra_args.extend(["--max-slots", str(args.max_slots)]) - - return run_transpile( - args.model_id, - extra_args=extra_args, - execute_after_transpile=args.execute_after_transpile, - allow_unconverted_weights=args.allow_unconverted_weights, - ) diff --git a/python/cactus/cli/list.py b/python/cactus/cli/list.py index fafebbe76..6ddc42382 100644 --- a/python/cactus/cli/list.py +++ b/python/cactus/cli/list.py @@ -2,7 +2,7 @@ import stat as _stat import struct -from .common import CYAN, print_color, transpiled_root, weights_root +from .common import BLUE, CYAN, print_color, weights_root _PREC_TO_BITS = {3: 1, 4: 2, 5: 3, 6: 4} @@ -86,10 +86,34 @@ def _collect(roots): return models +def _is_runnable_bundle(path): + return (path / "components" / "manifest.json").is_file() + + +def _categorize(roots): + converted = [] + runnable = [] + for item in _collect(roots): + path = item[0] + if _is_runnable_bundle(path): + runnable.append(item) + else: + converted.append(item) + return converted, runnable + + +def _print_section(title, color, models): + print_color(color, title) + if not models: + print(" (none)") + return + for path, model_type, quant, size in models: + print(f" {path.name:<32} {model_type:<18} {quant:<4} {_human_size(size):>9} {path}") + + def cmd_list(_args): - roots = (weights_root(), transpiled_root()) - converted, transpiled = _categorize(roots) + converted, runnable = _categorize((weights_root(),)) _print_section("Converted weights (cactus convert)", BLUE, converted) print() - _print_section("Runnable bundles (cactus convert / download)", CYAN, transpiled) + _print_section("Runnable bundles (cactus download)", CYAN, runnable) return 0 diff --git a/python/cactus/cli/model.py b/python/cactus/cli/model.py index 5e59f27f0..db9e7bdfc 100644 --- a/python/cactus/cli/model.py +++ b/python/cactus/cli/model.py @@ -3,12 +3,9 @@ import os import shutil -from dataclasses import dataclass, replace from pathlib import Path -from .common import GREEN, PROJECT_ROOT, RED, YELLOW, print_color - - +from .common import GREEN, RED, YELLOW, print_color def _convert_from_source(model_id, *, bits, token, weights_dir): @@ -50,156 +47,32 @@ def ensure_weights(model_id, *, bits=4, platform=None, token=None, reconvert=Fal return _convert_from_source(model_id, bits=bits, token=token, weights_dir=weights_dir) - -_DEFAULT_MULTIMODAL_PROMPT = ( - "Respond with 2 lines. The first should be a description of the image, " - "and the second should be a transcription of the audio" -) -_DEFAULT_TEXT_PROMPT = "Hello" - - -@dataclass(frozen=True) -class _TranspileSpec: - task: str - components: tuple[str, ...] = () - default_max_new_tokens: int | None = None - needs_image: bool = False - needs_audio: bool = False - force_component_pipeline: bool = False - - -def _spec_from_plan(plan): - return _TranspileSpec( - task=plan.task, - components=tuple(plan.components or ()), - default_max_new_tokens=plan.default_max_new_tokens, - needs_image=bool(plan.needs_image), - needs_audio=bool(plan.needs_audio), - force_component_pipeline=bool(plan.force_component_pipeline), - ) - - -def _infer_transpile_spec(*, task, plan): - if task != "auto": - if plan is not None and task == plan.task: - return _spec_from_plan(plan) - return _TranspileSpec( - task=task, - needs_image=task == "multimodal_causal_lm_logits", - needs_audio=task in { - "tdt_transcription", "seq2seq_transcription", - "ctc_logits", "encoder_hidden_states", - "multimodal_causal_lm_logits", - }, - force_component_pipeline=task in { - "tdt_transcription", "seq2seq_transcription", - "multimodal_causal_lm_logits", - }, - ) - - if plan is None: - return _TranspileSpec(task="causal_lm_logits") - - return _spec_from_plan(plan) - - -def _default_max_new_tokens(spec): - if spec.default_max_new_tokens is not None: - return int(spec.default_max_new_tokens) - return { - "seq2seq_transcription": 128, - "multimodal_causal_lm_logits": 512, - "causal_lm_logits": 128, - }.get(spec.task, 32) - - -def _default_multimodal_assets(): - """Return bundled test image/audio paths for multimodal shape capture.""" - candidates = ( - Path(__file__).resolve().parent.parent / "assets", - PROJECT_ROOT / "cactus-engine" / "tests" / "assets", - ) - def _find(name): - return next((d / name for d in candidates if (d / name).exists()), None) - image = _find("test_monkey.png") - audio = _find("test.wav") - return ([str(image)] if image else []), (str(audio) if audio else None) - - -def _default_audio_asset(): - _, audio = _default_multimodal_assets() - return audio - - -def _remove_stale_transpile_artifacts(output_dir): - for relative in ( - "components", - "transpile_entrypoints.json", - "raw_ir.json", - "optimized_ir.json", - "graph.cactus", - "graph_bindings.json", - "result.json", - ): - path = output_dir / relative - if path.is_dir(): - shutil.rmtree(path) - elif path.exists(): - path.unlink() - for pattern in ("raw_ir_*.json", "optimized_ir_*.json"): - for path in output_dir.glob(pattern): - if path.is_file(): - path.unlink() - - -def _has_transpiled_bundle(path): +def _has_runnable_bundle(path): + path = Path(path) return (path / "components" / "manifest.json").exists() -_AUDIO_TASKS = frozenset({ - "tdt_transcription", "seq2seq_transcription", - "ctc_logits", "encoder_hidden_states", -}) - - - - def resolve_bundle_dir(model_id): path = Path(model_id).expanduser() - if path.is_dir() and _has_transpiled_bundle(path): + if path.is_dir() and _has_runnable_bundle(path): return path return None -@dataclass(frozen=True) -class TranspileOptions: - task: str = "auto" - prompt: str | None = None - image_files: list[str] | None = None - audio_file: str | None = None - max_new_tokens: int | None = None - component_pipeline: str = "auto" - components: str | None = None - system_prompt: str | None = None - trust_remote_code: bool = False - local_files_only: bool = False - npu: bool = False - npu_quantize: int | None = None - npu_audio_quantize: int | None = None - npu_vision_quantize: int | None = None - cache_context_length: str | int | None = None +def _local_build_unavailable(model_id) -> RuntimeError: + return RuntimeError( + f"No runnable prebuilt bundle is available for {model_id!r}. " + "Local bundle builds are unavailable because the graph builder has been removed for rewrite." + ) def ensure_runnable_bundle(model_id, *, bits=4, platform=None, token=None, - reconvert=False, prebuilt=True, output_dir=None, - transpile=None): - """Resolve a runnable bundle via the full fallback ladder, building if needed. + reconvert=False, prebuilt=True, output_dir=None): + """Resolve a runnable bundle from a local path, cache, or prebuilt download. - Rungs, in order: (1) a local bundle path, (2) a cached prior build, - (3) a prebuilt bundle on HuggingFace, (4) local convert + transpile. - With prebuilt=False, rung (3) is skipped and the bundle is always built - locally when not already present. - Raises RuntimeError if every rung fails. + The old local source-to-bundle fallback was removed. Until the replacement + compiler lands, commands that need runnable bundles can only use existing + local bundles or published Cactus-Compute bundles. """ from .download import download_bundle, get_bundle_dir @@ -211,24 +84,23 @@ def ensure_runnable_bundle(model_id, *, bits=4, platform=None, token=None, raise RuntimeError(f"path not found: {model_id}") cached = Path(output_dir) if output_dir else get_bundle_dir(model_id, bits=bits, platform=platform) - if _has_transpiled_bundle(cached) and not reconvert: + if reconvert and cached.exists(): + print_color(YELLOW, "Removing cached bundle before refresh...") + shutil.rmtree(cached) + elif _has_runnable_bundle(cached): return cached - if prebuilt and not reconvert: - try: - return download_bundle(model_id, bits=bits, platform=platform, - token=token, output_dir=cached) - except (RuntimeError, OSError) as exc: - print_color(YELLOW, f"No prebuilt bundle ({exc}); building locally") + if not prebuilt: + raise _local_build_unavailable(model_id) - opts = transpile or TranspileOptions() - if platform == "apple" and not opts.npu: - opts = replace(opts, npu=True) - return ensure_bundle(model_id, bits=bits, platform=platform, token=token, - reconvert=reconvert, output_dir=cached, transpile=opts) + try: + return download_bundle(model_id, bits=bits, platform=platform, + token=token, output_dir=cached) + except (RuntimeError, OSError) as exc: + raise _local_build_unavailable(model_id) from exc -def prepare_bundle(args, *, model_id=None, transpile=None, prebuilt=True, +def prepare_bundle(args, *, model_id=None, prebuilt=True, output_dir=None, fail_prefix="Model setup failed"): """Resolve the platform from args and return a runnable bundle, with uniform error handling shared by every model command. Returns the bundle Path, or @@ -243,7 +115,6 @@ def prepare_bundle(args, *, model_id=None, transpile=None, prebuilt=True, reconvert=getattr(args, "reconvert", False), prebuilt=prebuilt, output_dir=output_dir, - transpile=transpile, ) except (RuntimeError, OSError, ValueError) as exc: print_color(RED, f"{fail_prefix}: {exc}") @@ -251,134 +122,12 @@ def prepare_bundle(args, *, model_id=None, transpile=None, prebuilt=True, def ensure_bundle(model_id, *, bits=4, platform=None, token=None, - reconvert=False, output_dir=None, transpile=None): - from .download import get_bundle_dir - from .transpile import run_transpile - from cactus.transpile.component_plan import infer_component_plan_from_output - - opts = transpile or TranspileOptions() - - if output_dir is not None: - output_dir = Path(output_dir).expanduser().resolve() - else: - output_dir = get_bundle_dir(model_id, bits=bits, platform=platform) - - ensure_weights( - model_id, bits=bits, platform=platform, token=token, - reconvert=reconvert, output_dir=output_dir, + reconvert=False, output_dir=None): + return ensure_runnable_bundle( + model_id, + bits=bits, + platform=platform, + token=token, + reconvert=reconvert, + output_dir=output_dir, ) - - if _has_transpiled_bundle(output_dir): - return output_dir - - plan = infer_component_plan_from_output(str(output_dir), model_id=model_id) - spec = _infer_transpile_spec(task=opts.task, plan=plan) - _remove_stale_transpile_artifacts(output_dir) - - spec_prompt = opts.prompt - spec_image_files = list(opts.image_files or []) - spec_audio_file = opts.audio_file - - if spec_prompt is None and spec.task == "multimodal_causal_lm_logits": - spec_prompt = _DEFAULT_MULTIMODAL_PROMPT - elif spec_prompt is None and spec.task == "causal_lm_logits": - spec_prompt = _DEFAULT_TEXT_PROMPT - - effective_component_pipeline = opts.component_pipeline - effective_components = opts.components - - if spec.task == "multimodal_causal_lm_logits": - needs_image = spec.needs_image - needs_audio = spec.needs_audio - if not needs_image and not needs_audio: - needs_image = bool(spec_image_files) - needs_audio = bool(spec_audio_file) - if (needs_image and not spec_image_files) or (needs_audio and not spec_audio_file): - default_images, default_audio = _default_multimodal_assets() - if needs_image and not spec_image_files: - spec_image_files = default_images - if needs_audio and not spec_audio_file: - spec_audio_file = default_audio - print_color( - YELLOW, - "Multimodal transpile needs representative media shapes; " - "using bundled tiny test assets.", - ) - if needs_image and not spec_image_files: - raise RuntimeError("Building this multimodal model requires --image-file.") - if needs_audio and not spec_audio_file: - raise RuntimeError("Building this multimodal model requires --audio-file.") - - if effective_component_pipeline == "auto" and spec.force_component_pipeline: - effective_component_pipeline = "on" - if effective_components is None and spec.components: - effective_components = ",".join(spec.components) - - used_default_audio = False - if spec.task in _AUDIO_TASKS and not spec_audio_file: - spec_audio_file = _default_audio_asset() - used_default_audio = spec_audio_file is not None - if spec.task in _AUDIO_TASKS and used_default_audio: - print_color( - YELLOW, - f"{spec.task} transpile needs a representative audio shape; " - "using bundled tiny test audio asset.", - ) - elif spec.task in _AUDIO_TASKS and not spec_audio_file: - raise RuntimeError(f"Building a {spec.task} model requires --audio-file.") - - effective_max_new_tokens = opts.max_new_tokens or _default_max_new_tokens(spec) - - extra_args = [ - "--weights-dir", str(output_dir), - "--artifact-dir", str(output_dir), - "--task", spec.task, - "--max-new-tokens", str(effective_max_new_tokens), - "--component-pipeline", effective_component_pipeline, - ] - if spec_prompt is not None: - extra_args.extend(["--prompt", spec_prompt]) - if effective_components: - extra_args.extend(["--components", str(effective_components)]) - for img in spec_image_files: - extra_args.extend(["--image-file", img]) - if spec_audio_file: - extra_args.extend(["--audio-file", str(spec_audio_file)]) - if opts.system_prompt: - extra_args.extend(["--system-prompt", str(opts.system_prompt)]) - if token: - extra_args.extend(["--token", token]) - if opts.trust_remote_code or spec.task == "multimodal_causal_lm_logits": - extra_args.append("--trust-remote-code") - if opts.local_files_only: - extra_args.append("--local-files-only") - if opts.npu: - extra_args.append("--npu") - if opts.npu_quantize is not None: - extra_args.extend(["--npu-quantize", str(int(opts.npu_quantize))]) - if opts.npu_audio_quantize is not None: - extra_args.extend(["--npu-audio-quantize", str(int(opts.npu_audio_quantize))]) - if opts.npu_vision_quantize is not None: - extra_args.extend(["--npu-vision-quantize", str(int(opts.npu_vision_quantize))]) - if opts.cache_context_length is not None: - extra_args.extend(["--cache-context-length", str(opts.cache_context_length)]) - - rc = run_transpile(model_id, extra_args=extra_args) - if rc != 0: - raise RuntimeError(f"Build failed for {model_id}") - - try: - from cactus.convert.handoff_probe import ( - export_gemma4_handoff_probe, - export_parakeet_handoff_probe, - ) - - if export_gemma4_handoff_probe(output_dir, model_id=model_id): - print_color(GREEN, f"Gemma4 cloud handoff probe packaged into {output_dir}") - if export_parakeet_handoff_probe(output_dir, model_id=model_id): - print_color(GREEN, f"Parakeet cloud handoff probe packaged into {output_dir}") - except Exception as e: - print_color(YELLOW, f"Warning: failed to package cloud handoff probe: {e}") - - print_color(GREEN, f"Model built at {output_dir}") - return output_dir diff --git a/python/cactus/cli/run.py b/python/cactus/cli/run.py index e7bd3814a..eea8e127e 100644 --- a/python/cactus/cli/run.py +++ b/python/cactus/cli/run.py @@ -19,12 +19,9 @@ def cmd_run(args) -> int: if args.result_json: args.result_json = str(Path(args.result_json).expanduser()) - from .model import TranspileOptions, prepare_bundle + from .model import prepare_bundle - bundle_dir = prepare_bundle(args, transpile=TranspileOptions( - image_files=[args.image] if args.image else None, - audio_file=args.audio, - )) + bundle_dir = prepare_bundle(args) if bundle_dir is None: return 1 diff --git a/python/cactus/cli/serve.py b/python/cactus/cli/serve.py index d5dc0ab7e..005203c79 100644 --- a/python/cactus/cli/serve.py +++ b/python/cactus/cli/serve.py @@ -57,7 +57,7 @@ def cmd_serve(args): ) except RuntimeError as exc: print_color(RED, f"Error: {exc}") - print("Prepare a bundle first with `cactus run ` (or `cactus convert `).") + print("Prepare a bundle first with `cactus run ` or `cactus download `.") return 1 models = sorted(application.state.registry.models) diff --git a/python/cactus/cli/transcribe.py b/python/cactus/cli/transcribe.py index a4d00acea..9b555a14c 100644 --- a/python/cactus/cli/transcribe.py +++ b/python/cactus/cli/transcribe.py @@ -4,7 +4,7 @@ def cmd_transcribe(args): - from .model import prepare_bundle, TranspileOptions + from .model import prepare_bundle if args.audio_file: audio_path = Path(args.audio_file).expanduser() @@ -15,7 +15,7 @@ def cmd_transcribe(args): apply_runtime_env(args) - bundle_dir = prepare_bundle(args, transpile=TranspileOptions(audio_file=args.audio_file)) + bundle_dir = prepare_bundle(args) if bundle_dir is None: return 1 diff --git a/python/cactus/cli/transpile.py b/python/cactus/cli/transpile.py deleted file mode 100644 index fdeb1104d..000000000 --- a/python/cactus/cli/transpile.py +++ /dev/null @@ -1,66 +0,0 @@ -from __future__ import annotations - -import os -import subprocess -import sys -from pathlib import Path - -from .common import PROJECT_ROOT, RED, YELLOW, print_color -from .download import get_weights_dir - - -def _weights_dir_looks_transpile_ready(weights_dir): - root = Path(weights_dir).expanduser().resolve() - if not root.is_dir(): - return False - if (root / "weights_manifest.json").exists(): - return True - if any(root.glob("*.cq[1-4].weights")): - return True - return (root / "config.txt").exists() and any(root.glob("*.weights")) - - -def _extra_args_has_option(extra_args, option): - prefix = f"{option}=" - return any(arg == option or arg.startswith(prefix) for arg in extra_args) - - -def _prepend_python_path(env): - python_root = str(PROJECT_ROOT / "python") - existing = env.get("PYTHONPATH") - env["PYTHONPATH"] = python_root if not existing else f"{python_root}{os.pathsep}{existing}" - - -def run_transpile(model_id, *, extra_args, execute_after_transpile=False, - allow_unconverted_weights=False): - from .runtime import ensure_python_runtime_library - - extra_args = list(extra_args or []) - command = [sys.executable, "-m", "cactus.transpile.hf_model", "--model-id", model_id] - default_weights_dir = get_weights_dir(model_id) - if not default_weights_dir.name: - print_color(RED, f"Error: cannot derive a unique output dir from model_id {model_id!r}.") - print_color(YELLOW, "Pass --weights-dir and --artifact-dir explicitly.") - return 1 - if not _extra_args_has_option(extra_args, "--weights-dir"): - if _weights_dir_looks_transpile_ready(default_weights_dir): - command.extend(["--weights-dir", str(default_weights_dir)]) - if not _extra_args_has_option(extra_args, "--artifact-dir"): - command.extend(["--artifact-dir", str(default_weights_dir)]) - - if allow_unconverted_weights: - command.append("--allow-unconverted-weights") - if not execute_after_transpile and "--skip-execute" not in extra_args: - command.append("--skip-execute") - command.extend(extra_args) - - try: - transpile_lib = ensure_python_runtime_library() - except RuntimeError as exc: - print_color(RED, f"Error: {exc}") - return 1 - - env = os.environ.copy() - env["CACTUS_LIB_PATH"] = str(transpile_lib) - _prepend_python_path(env) - return subprocess.run(command, cwd=PROJECT_ROOT, env=env).returncode diff --git a/python/cactus/cli/upload.py b/python/cactus/cli/upload.py deleted file mode 100644 index b0277b737..000000000 --- a/python/cactus/cli/upload.py +++ /dev/null @@ -1,53 +0,0 @@ -import os -import tempfile -import zipfile -from pathlib import Path - -from .common import BLUE, GREEN, PROJECT_ROOT, print_color -from .download import get_model_dir_name, resolve_platform -from .model import prepare_bundle -from .utils import suggested_cq_repo, variant_suffix - - -def cmd_upload(args): - """Build a runnable bundle locally and upload it to Cactus-Compute on HuggingFace.""" - token = args.token or os.getenv("HF_TOKEN") - if not token: - raise SystemExit("missing HuggingFace token: pass --token or set HF_TOKEN") - - platform = resolve_platform(args.platform) - repo_id = suggested_cq_repo(args.model_id) - stem = f"{get_model_dir_name(args.model_id)}-{variant_suffix(args.bits, platform)}" - archive_name = f"{stem}.zip" - tag = (PROJECT_ROOT / "CACTUS_VERSION").read_text(encoding="utf-8").strip() - - bundle_dir = prepare_bundle(args, prebuilt=False, fail_prefix=f"Build failed for {args.model_id}") - if bundle_dir is None: - return 1 - - from huggingface_hub import HfApi - - api = HfApi(token=token) - api.create_repo(repo_id=repo_id, repo_type="model", exist_ok=True) - - with tempfile.TemporaryDirectory(prefix="cactus-upload-") as tmp: - archive_path = Path(tmp) / archive_name - with zipfile.ZipFile(archive_path, "w", compression=zipfile.ZIP_DEFLATED) as zf: - for path in sorted(p for p in bundle_dir.rglob("*") if p.is_file()): - zf.write(path, f"{stem}/{path.relative_to(bundle_dir).as_posix()}") - print_color(BLUE, f"Uploading {archive_name} ({archive_path.stat().st_size / 1e6:.1f} MB) -> {repo_id}") - api.upload_file( - path_or_fileobj=str(archive_path), - path_in_repo=archive_name, - repo_id=repo_id, - repo_type="model", - commit_message=f"Upload {archive_name} ({tag})", - ) - - try: - api.delete_tag(repo_id=repo_id, tag=tag, repo_type="model") - except Exception: - pass - api.create_tag(repo_id=repo_id, tag=tag, repo_type="model", revision="main") - print_color(GREEN, f"Done: {repo_id} @ {tag}") - return 0 diff --git a/python/cactus/convert/compat.py b/python/cactus/convert/compat.py index 99d4fb5ec..4353c3382 100644 --- a/python/cactus/convert/compat.py +++ b/python/cactus/convert/compat.py @@ -1,29 +1,174 @@ from __future__ import annotations +import builtins +import importlib.util +import sys +from functools import lru_cache +from pathlib import Path -def patch_transformers_import_compat() -> list[str]: - """Keep Transformers processor/model imports working in local dev envs. +try: + import torch +except Exception: # pragma: no cover + torch = None - Some pyenv Python builds are missing the stdlib `_lzma` extension. Importing - `AutoProcessor` can still reach torchvision datasets, which import `lzma` - even when conversion only needs processor metadata. The transpiler already - has the battle-tested patch, so reuse it before converter-side HF imports. - """ - notes: list[str] = [] +_TORCHVISION_COMPAT_LIBRARIES: list[object] = [] +_GEMMA4_MODULE = "transformers.models.gemma4.modeling_gemma4" + + +def _transformers_supports_model_module(module_name: str) -> bool: + try: + return importlib.util.find_spec(module_name) is not None + except Exception: + return False + + +def _candidate_external_site_packages() -> list[Path]: + candidates: list[Path] = [] + major_minor = f"python{sys.version_info.major}.{sys.version_info.minor}" + pyenv_versions = Path.home() / ".pyenv" / "versions" + if not pyenv_versions.exists(): + return candidates + for version_dir in sorted(pyenv_versions.iterdir(), reverse=True): + site_packages = version_dir / "lib" / major_minor / "site-packages" + if site_packages.exists(): + candidates.append(site_packages) + return candidates + + +def _ensure_transformers_supports_gemma4() -> str | None: + if _transformers_supports_model_module(_GEMMA4_MODULE): + return None + + for site_packages in _candidate_external_site_packages(): + candidate = site_packages / Path(*_GEMMA4_MODULE.split(".")) + if not candidate.with_suffix(".py").exists() and not candidate.is_dir(): + continue + sys.path.insert(0, str(site_packages)) + for module_name in list(sys.modules): + root_name = module_name.split(".", 1)[0] + if root_name in {"transformers", "huggingface_hub", "tokenizers"}: + del sys.modules[module_name] + if _transformers_supports_model_module(_GEMMA4_MODULE): + return str(site_packages) + try: + sys.path.remove(str(site_packages)) + except ValueError: + pass + return None + + +def _patch_torchvision_missing_nms_op() -> str | None: + """Keep mismatched torch/torchvision installs importable for processors.""" + if torch is None: + return None + + try: + import torchvision # type: ignore # noqa: F401 + return None + except RuntimeError as exc: + if "operator torchvision::nms does not exist" not in str(exc): + return None + except Exception: + return None + try: - from cactus.transpile.runtime_support import ( - ensure_transformers_supports_gemma4, - patch_torch_flex_attention_compat, - patch_transformers_torchvision_probe, - ) + import torchvision.extension as tv_extension # type: ignore + + if bool(getattr(tv_extension, "_HAS_OPS", False)): + return None + except Exception: + pass + + try: + library = torch.library.Library("torchvision", "DEF") + library.define("nms(Tensor dets, Tensor scores, float iou_threshold) -> Tensor") + _TORCHVISION_COMPAT_LIBRARIES.append(library) + return "defined missing torchvision::nms operator for torchvision import compatibility" except Exception: - return notes + return None + + +def _patch_transformers_torchvision_probe() -> str | None: + has_torchvision = importlib.util.find_spec("torchvision") is not None + has_lzma = importlib.util.find_spec("_lzma") is not None + if not has_torchvision: + return None + + if has_lzma: + return _patch_torchvision_missing_nms_op() + + base_note: str | None = None + try: + import backports.lzma as backports_lzma # type: ignore + + sys.modules.setdefault("lzma", backports_lzma) + base_note = "using backports.lzma because this Python build is missing _lzma" + nms_patch_note = _patch_torchvision_missing_nms_op() + if nms_patch_note: + return f"{base_note}; {nms_patch_note}" + return base_note + except Exception: + pass + + class _InterpolationModeStub: + NEAREST = "nearest" + BILINEAR = "bilinear" + BICUBIC = "bicubic" + LANCZOS = "lanczos" + + class _TorchvisionFunctionalStub: + InterpolationMode = _InterpolationModeStub + + def __getattr__(self, name: str): + raise RuntimeError( + "torchvision functionality is unavailable because this Python build is missing _lzma; " + f"attempted to access torchvision.transforms.functional.{name}" + ) + + builtins.F = _TorchvisionFunctionalStub() + builtins.tvF = builtins.F + + import transformers.utils as tf_utils # type: ignore + import transformers.utils.import_utils as tf_import_utils # type: ignore + + @lru_cache + def _disabled() -> bool: + return False + + tf_import_utils.is_torchvision_available = _disabled + tf_import_utils.is_torchvision_v2_available = _disabled + tf_utils.is_torchvision_available = _disabled + tf_utils.is_torchvision_v2_available = _disabled + return "disabled torchvision import checks because this Python build is missing _lzma" + + +def _patch_torch_flex_attention_compat() -> str | None: + try: + import torch.nn.attention.flex_attention as flex_attention # type: ignore + except Exception: + return None + + if hasattr(flex_attention, "AuxRequest"): + return None + + class _AuxRequest: + def __init__(self, **kwargs): + for key, value in kwargs.items(): + setattr(self, key, value) + + flex_attention.AuxRequest = _AuxRequest # type: ignore[attr-defined] + return "installed torch flex_attention AuxRequest compatibility stub" + + +def patch_transformers_import_compat() -> list[str]: + """Keep Transformers processor/model imports working in local dev envs.""" + notes: list[str] = [] for patch in ( - patch_transformers_torchvision_probe, - patch_torch_flex_attention_compat, - ensure_transformers_supports_gemma4, + _patch_transformers_torchvision_probe, + _patch_torch_flex_attention_compat, + _ensure_transformers_supports_gemma4, ): try: note = patch() diff --git a/python/cactus/convert/tests/test_naming_qdq.py b/python/cactus/convert/tests/test_naming_qdq.py index 860c8ffd4..106556b19 100644 --- a/python/cactus/convert/tests/test_naming_qdq.py +++ b/python/cactus/convert/tests/test_naming_qdq.py @@ -320,10 +320,10 @@ def test_nomic_norm2_weight_uses_runtime_name(): def test_nomic_keeps_qkv_and_moe_experts_fused(): - # The v2 transpile path binds graph weights by their HF parameter name, so the + # Runtime graph bundles bind these weights by their HF parameter name, so the # converter emits one fused tensor per HF parameter (no q/k/v or per-expert split). # w2 is stored transposed so the second expert matmul can consume it as a direct - # linear weight in the transpiled graph. + # linear weight in the graph bundle. adapter = adapter_for_family("nomic") adapter.num_experts = 8 diff --git a/python/cactus/server.py b/python/cactus/server.py index 18d23dc4b..6c50a051b 100644 --- a/python/cactus/server.py +++ b/python/cactus/server.py @@ -147,7 +147,7 @@ def default_llm(self, preferred: str | None = None) -> ModelInfo: for info in sorted(self.models.values(), key=lambda x: x.id): if info.model_type in LLM_MODEL_TYPES: return info - raise RuntimeError("No valid LLM bundles found. Prepare a transpiled bundle before running `cactus serve`.") + raise RuntimeError("No valid LLM bundles found. Download a runnable bundle before running `cactus serve`.") def default_stt(self) -> ModelInfo | None: for info in sorted(self.models.values(), key=lambda x: x.id): diff --git a/python/cactus/transpile/README.md b/python/cactus/transpile/README.md deleted file mode 100644 index 650e0686b..000000000 --- a/python/cactus/transpile/README.md +++ /dev/null @@ -1,1024 +0,0 @@ -# Cactus Transpiler - -This folder contains the Python side of the Cactus transpiler: the code that takes -a Hugging Face or PyTorch model, captures it with `torch.export`, converts that -export into Cactus' intermediate representation (`IRGraph`), canonicalizes and -optimizes the IR, lowers it into a runtime `Graph`, and optionally saves or runs -the result. - -The transpiler has two main output shapes: - -- A single lowered graph for simple tasks such as text logits or encoder-only audio. -- A component bundle for models that are easier to split into staged graphs, most - notably Gemma4 multimodal and Parakeet TDT. - -This README is a code-first walkthrough of the actual flow in the repo today. - -## Related Entry Points - -These files are outside `python/cactus/transpile`, but they are the top-level -entry points for the transpiler (driven by `cactus convert` for end users; the -`cactus transpile` command is hidden from the CLI menu and used internally): - -- `python/cactus/cli/convert.py` (`cmd_convert` → `cmd_transpile`) -- `python/cactus/cli/transpile.py` (`run_transpile`) -- `python/cactus/transpile/hf_model.py` - -There is no separate `cactus transpile` subcommand. `cactus convert ` -quantizes the weights to CQ and then builds the runtime graph in one step; the graph -build is `run_transpile`, a thin wrapper around `python/cactus/transpile/hf_model.py`. -The packaged module is the real top-level transpile program. - -## CLI Entry Flow - -### `cactus convert` - -`cmd_convert` (`python/cactus/cli/convert.py`) quantizes the weights via `ensure_weights`, -then (unless `--weights-only`) calls `cmd_transpile`, which forwards options to -`run_transpile`: - -```python -def run_transpile(model_id, *, extra_args, execute_after_transpile=False, ...): - transpile_lib = _ensure_python_runtime_library() - model_id = resolve_model_id_alias(model_id) - command = [sys.executable, "-m", "cactus.transpile.hf_model", "--model-id", model_id] - if "--weights-dir" not in extra_args: - default_weights_dir = get_weights_dir(model_id) - if _weights_dir_looks_transpile_ready(default_weights_dir): - command.extend(["--weights-dir", str(default_weights_dir)]) - else: - return 1 - if not execute_after_transpile and "--skip-execute" not in extra_args: - command.append("--skip-execute") - command.extend(extra_args) -``` - -Two important details: - -- The graph build saves artifacts and skips execution by default - (pass `--execute-after-transpile` to run a reference execution). -- Building the runtime graph requires converted Cactus CQ weights. `cactus convert` - produces them before this stage; or pass `--weights-dir` to an existing converted folder. - -### `cactus run` - -`cactus run ` is the single user-facing entry. Resolution -order (`python/cactus/cli/model.py:ensure_runnable_bundle`): - -1. `` is an existing path with `components/manifest.json` → run as-is. -2. Cached download at `weights/-cq[-]/` → run. -3. Fresh download from `huggingface.co/Cactus-Compute` via `download_bundle`. -4. Fallback: local `ensure_bundle` (convert + transpile for custom models not on HF). - -Bundle handoff is to the native `run` binary (`python/cactus/bin/run`), which -loads the bundle via `cactus_init()`. - -## End-To-End Timeline - -## 1. Parse Args, Infer Task, Load Model Bundle - -The real transpile program lives in `python/cactus/transpile/hf_model.py`. -Its `main()` function decides what kind of model/task it is dealing with, validates -optional converted weights, and loads the HF model plus tokenizer/processor: - -```python -if args.task == "auto": - inferred_task = _infer_task_from_config(args.model_id) - ... - if image_files or (args.audio_file and has_multimodal_config and inferred_task == "causal_lm_logits"): - task = "multimodal_causal_lm_logits" - else: - task = inferred_task -else: - task = args.task - -validated_weights_dir = _validate_weights_dir(args.weights_dir.strip() or None, model_id=args.model_id) -weights_dir = str(validated_weights_dir) if validated_weights_dir is not None else None - -model_source, processor_or_tokenizer, model, model_config = _load_transformers_bundle( - model_id=args.model_id, - task=task, - torch_dtype=torch_dtype, - token=args.token, - trust_remote_code=args.trust_remote_code, - local_files_only=args.local_files_only, -) -``` - -### What decides the task? - -`_infer_task_from_config()` looks at: - -- `config.json` architectures -- model type -- known special cases like Parakeet TDT and Whisper -- model id substrings as a last fallback - -Today the main task shapes are: - -- `causal_lm_logits` -- `multimodal_causal_lm_logits` -- `ctc_logits` -- `encoder_hidden_states` -- `tdt_transcription` - -## 2. Prepare Example Inputs - -Before capture, the transpiler builds one concrete example input batch. That -example batch drives both `torch.export` capture and later optional execution. - -The high-level branches are: - -- `_prepare_text_inputs()` for text-only causal LM capture -- `_prepare_audio_inputs()` for speech/audio tasks -- `_prepare_gemma4_multimodal_inputs()` for Gemma4 multimodal - -Example: - -```python -if task == "causal_lm_logits": - prepared = _prepare_text_inputs(...) -elif task == "multimodal_causal_lm_logits": - prepared = _prepare_gemma4_multimodal_inputs(...) -elif task == "tdt_transcription": - prepared = _prepare_audio_inputs(...) -else: - prepared = _prepare_audio_inputs(...) -``` - -The `PreparedInputs` object carries: - -- `names`: logical input names like `input_ids` or `input_features` -- `tensors`: the actual example tensors -- `metadata`: prompt/audio/image/sample-rate info that later gets written into - manifests - -## 3. Canonicalize The Model Interface - -The transpiler does not try to capture every model's native HF forward signature -as-is. It first wraps the model in an adapter that exposes a smaller, more stable -task-specific interface. - -That happens in `model_adapters.py`: - -```python -def canonicalize_model_interface(model, task="causal_lm_logits", *, input_names=None, weights_dir=None): - family = _family_key(model) - ... - if task == "causal_lm_logits": - ... - elif task == "multimodal_causal_lm_logits": - adapter_factory = lambda inner_model: Gemma4MultimodalCausalLMLogitsAdapter( - inner_model, - input_names=resolved_input_names, - weights_dir=weights_dir, - ) - elif task == "ctc_logits": - adapter_factory = lambda inner_model: CTCLogitsAdapter(...) - elif task == "encoder_hidden_states": - adapter_factory = lambda inner_model: EncoderHiddenStatesAdapter(...) -``` - -This step is where family-specific knowledge lives: - -- Gemma/Gemma3/Gemma4/Qwen causal LM wrappers -- Gemma4 multimodal wrappers -- generic CTC and encoder wrappers - -The adapter is also where graph metadata and import hints get attached. - -## 3b. Import PyTorch Ops, Not Model Layer Names - -The compiler boundary is the `torch.export` graph. After capture, Cactus imports -`call_function` nodes by their canonical ATen operation name, not by Hugging Face -module names. For example, `torch.ops.aten.linear.default`, `aten.linear.default`, -and equivalent exported spellings normalize through `aten_ops.py` to the same -canonical IR op: - -```python -aten.linear.default -> linear -aten.silu.default -> silu -aten.conv1d.default -> conv1d -aten.rms_norm.default -> rms_norm -``` - -This is the main generality boundary: model wrappers may still prepare inputs or -split components, but actual compute lowering is ATen op / pattern based. Layer -names are allowed as structural hints for component boundaries and exact -converted-weight manifest aliases; they should not be required to decide how a -PyTorch op maps to a Cactus op. - -## 4. Decide Between Monolithic Capture And Component Capture - -After input prep, the script asks `build_component_module_specs()` whether the -current model supports a split component pipeline: - -```python -component_specs = build_component_module_specs( - model, - task=task, - named_tensors=_named_tensor_store(prepared), - weights_dir=weights_dir, -) -``` - -Today that returns component specs for: - -- Gemma4 multimodal: `vision_encoder`, `audio_encoder`, `lm_encoder`, `decoder` -- Parakeet TDT: `audio_encoder`, `decoder` - -If component specs exist and `--component-pipeline` is `auto` or `on`, the script -takes the component path. Otherwise it captures a single wrapped model graph. - -## 5A. Component Pipeline Path - -The component path is driven by `_run_component_pipeline_transpile()` in -`python/cactus/transpile/hf_model.py`: - -```python -for spec in component_specs: - captured_components[spec.component] = capture_component_spec(spec, fusion_config=fusion_config) -``` - -Each component spec is captured by `component_pipeline.capture_component_spec()`: - -```python -captured = capture_model(wrapped, spec.example_inputs) -raw_ir_graph = copy.deepcopy(captured.ir_graph) -optimized_ir_graph = copy.deepcopy(captured.ir_graph) -canonicalize_exported_graph(optimized_ir_graph) -optimize_graph(optimized_ir_graph, config=fusion_config) -transpiled_graph = transpile_preoptimized_ir(copy.deepcopy(optimized_ir_graph)) -``` - -So each component goes through the same capture/import/canonicalize/optimize/lower -pipeline independently. - -For Gemma4 multimodal, `model_adapters.py` builds the specs by first computing -example image/audio features, then generating example inputs for the downstream -`lm_encoder` and `decoder` stages. - -## 5B. Monolithic Capture Path - -The non-component path wraps the canonical adapter in `TranspileWrapper` and runs -the full graph through capture: - -```python -wrapper = TranspileWrapper(canonical.module, weights_dir=weights_dir).eval() - -captured = capture_model(wrapper, prepared.tensors) -raw_ir_graph = copy.deepcopy(captured.ir_graph) - -canonicalize_exported_graph(captured.ir_graph) -optimize_graph(captured.ir_graph, config=fusion_config) -tg = _lower_preoptimized_ir(captured.ir_graph) -``` - -`TranspileWrapper` mainly exists to pass `weights_dir` into the graph metadata: - -```python -def get_transpile_metadata(self) -> dict[str, object]: - ... - if self.weights_dir: - graph_meta["weights_dir"] = self.weights_dir -``` - -## 6. Capture With `torch.export` - -`capture_pytorch.py` owns the `torch.export` boundary: - -```python -normalized_kwargs = _inject_export_safe_kwargs(model, _normalize_kwargs(kwargs)) -... -transpile_metadata = _collect_transpile_metadata(model, example_args, example_kwargs) - -ep = export(model, args=example_args, kwargs=example_kwargs, strict=strict) -... -ir_graph = import_captured_to_ir(raw_captured, strict=strict) -``` - -Important behavior here: - -- HF-style `use_cache=False` and `return_dict=False` are injected automatically - when the model signature supports them. -- Metadata providers are collected from any module that defines - `get_transpile_metadata()` or `transpile_metadata()`. -- The captured export is immediately imported into the repo's own IR. - -The capture output is a `CapturedModel` dataclass holding: - -- the exported program -- the FX graph -- the lifted state dict -- the imported `IRGraph` -- the example inputs -- the collected metadata - -## 7. Import Exported FX Into `IRGraph` - -`import_ir.py` converts the exported FX/Export graph into the repo's intermediate -representation: - -```python -ir = IRGraph(values={}, nodes={}, order=[], inputs=[], outputs=[], constants={}, meta=dict(graph_meta)) -ctx = ImportContext(strict=strict, transpile_metadata=transpile_metadata) -weights_dir = resolve_transpile_weights_dir(ir.meta) - -for node in captured.graph.nodes: - if node.op == "placeholder": - ... - if node.op == "get_attr": - ... - if node.op == "call_function": - import_call_function(...) - if node.op == "output": - import_output(ir, node, ctx) - -apply_import_semantics(ir) -annotate_ir_components(ir) -verify_ir(ir) -``` - -### What gets imported? - -- Real graph inputs become `IRGraph.inputs` -- lifted parameters/buffers/constants become `IRGraph.constants` -- FX `call_function` nodes become `IRNode`s via `importers.py` - -### How weight binding metadata enters the IR - -When a captured constant corresponds to a known converted weight file, -`import_captured_to_ir()` resolves that mapping early: - -```python -binding = resolve_weight_binding(weights_dir=weights_dir, source_name=target) -... -ir.values[value_id_str].meta.update( - { - "path": binding.path, - "kind": binding.kind, - "source_name": binding.source_name, - } -) -``` - -That metadata is what later lets lowering use `g.mmap_weights()` instead of -embedding the constant payload into the graph. Embedding lookups still lower to -`embedding_from_tensor()`, but the backing table is represented as a normal mmap -weight tensor so tied output heads can reuse the same constant safely. -By default this resolver only trusts explicit `weights_manifest.json` entries -written by `cactus convert` plus wrapper-local aliases derived from those entries. -It does not guess old model-specific filenames from layer names. - -## 8. Importers: FX Target To Canonical IR Node - -`importers.py` is the large target-dispatch layer. It normalizes FX/PyTorch targets -to canonical op names and creates `IRNode`s with stable attributes. - -The dispatch point is: - -```python -def import_call_function(ir, node, ctx, *, shape, dtype, torch_op): - op = normalize_target(torch_op) - importer = OP_IMPORTERS.get(op) - if importer is None: - import_opaque_call_function(...) - else: - importer(...) -``` - -Import is intentionally ATen/op based. It does not apply module-path hints such -as `model.layers.N.self_attn`; model-specific behavior should live in explicit -preprocessing/component boundaries or in topology-based fusion passes. - -Some notable importer behaviors: - -- `normalize.py` reduces many `aten.*` spellings into one canonical name. -- Unknown ops are imported as `kind="opaque"` IR nodes rather than being silently - dropped. -- Import hints can inject attrs/meta into matching nodes after import. -- Some nested exported subgraphs such as `wrap_with_set_grad_enabled` are inlined - explicitly in `import_ir.py`. - -## 9. Import-Time Semantic Rewrites - -Immediately after raw import, `import_semantics.py` performs a first semantic pass: - -```python -def apply_import_semantics(graph: IRGraph) -> IRGraph: - _rewrite_explicit_attention(graph) - _tag_explicit_linear(graph) - _rewrite_early_rope(graph) - rebuild_graph(graph) - verify_ir(graph) - return graph -``` - -This step is intentionally early. It rewrites export patterns into higher-level -meaning before the main optimizer sees them. - -Examples: - -- `scaled_dot_product_attention` becomes semantic `attention` -- explicit linears get tagged -- RoPE patterns can become semantic `rope` before later fusion passes - -## 10. Canonical Cleanup - -`canonicalize/cleanup.py` turns imported IR into a stricter canonical form: - -```python -def canonicalize_exported_graph(graph: IRGraph, *, max_passes: int = 8) -> IRGraph: - for _ in range(max_passes): - for node_id in list(graph.order): - _canonicalize_node(graph, node) - rebuild_graph(graph) - - for node_id in list(graph.order): - _simplify_node(graph, node) - rebuild_graph(graph) - - if _legalize_precisions(graph): - rebuild_graph(graph) - - dce(graph) -``` - -What this pass is responsible for: - -- renaming ops into canonical spellings -- turning shape-only ops into a stable `view` -- rewriting `transpose` and `movedim` into `permute` -- constant-folding trivial expressions -- materializing `ones`, `zeros`, and `arange` constants -- removing no-op casts/views/slices -- inserting FP16 legalizing casts where the runtime requires them -- dead-code elimination - -It also records unsupported canonical ops in `graph.meta["canonical_unsupported_op_counts"]`. - -## 11. Fusion And Graph Optimization - -`optimize_graph.py` runs the semantic fusion passes: - -```python -def optimize_graph(graph: IRGraph, *, max_passes: int = 8, config: FusionConfig | None = None) -> IRGraph: - canonicalize_exported_graph(graph) - ... - if config.enable_rms_norm: - fuse_rms_norm(graph) - fuse_rms_norm_scale_multiply(graph) - if config.enable_rope: - fuse_rope(graph) - if config.enable_attention: - fuse_attention(graph) - ... - normalize_gemma4_decoder_attention_semantics(graph) - fuse_dense_mlp_tq(graph) - ... - annotate_gold_patterns(graph) - _prune_unused_inputs(graph) -``` - -This is where recognizable subgraphs collapse into semantic or fused nodes such as: - -- `rms_norm` -- `rope` -- `attention` -- `attention_block` -- `self_attention_block` -- `conv_module` -- `lstm_cell` -- `gated_deltanet_prefill` -- `gated_deltanet_decode` - -### Gemma4-specific optimization logic - -Gemma4 gets extra handling in this pass: - -- import hints and graph metadata carry layer types and sliding window info -- exported masks may be elided and rewritten into `window_size` -- some decoder attention behavior is normalized specifically for the saved runtime - -## 12. Lower IR Into A Cactus `Graph` - -`lower.py` is the last compiler stage. It maps `IRValue`s and `IRNode`s into the -runtime `Graph` object from `cactus.bindings.cactus`. - -The high-level structure is: - -```python -def transpile_preoptimized_ir(ir: IRGraph) -> TranspiledGraph: - g = Graph() - env = {} - - for value_id in ir.inputs: - tensor = _lower_input_value(g, ir.values[value_id]) - env[value_id] = tensor - - for value_id, const in ir.constants.items(): - binding = _lookup_weight_binding(value) - if binding is not None: - binding = ensure_binding_compatible(binding, source_tensor=const) - lowered_const = _lower_constant_value(g, value, const, binding=binding) - env[value_id] = lowered_const - - for node_id in ir.order: - outputs = _lower_ir_node(g, ir.nodes[node_id], env, ir) -``` - -### Constant lowering - -This is the critical constant binding behavior: - -```python -def _lower_constant_value(g, value, const, *, binding=None): - if binding is not None: - return g.mmap_weights(binding.path) - - if tensor_value.numel() == 1: - return tensor_value.item() - - return _materialize_constant_tensor(g, tensor_value) -``` - -So constants can become: - -- `g.mmap_weights(path)` -- an inline scalar -- a materialized tensor node baked into the graph - -### Where weight compatibility is repaired - -`weight_compat.py` sits between IR metadata and actual lowering. It can rewrite a -binding to a more executable companion file, for example: - -- convert INT8 embeddings to a cached CQ embedding format -- materialize CQ4 companion files for legacy packed INT4 tensors - -### What lowering implements - -`_lower_ir_node()` is the actual op dispatcher. It covers: - -- elementwise math and comparisons -- views, reshape, expand, permute, transpose -- matmul and linear -- embedding and gather/index paths -- reductions and softmax -- conv1d/conv2d -- norms -- fused attention-family nodes -- fused LSTM and Conformer-style blocks - -If an IR op survives all prior passes and still has no lowering, `lower.py` raises: - -```python -raise NotImplementedError(f"unsupported IR op in lowering: {op}") -``` - -### Gemma4 decoder attention fallback - -One especially important special case is that Gemma4 decoder attention can bypass -the runtime attention kernel and be lowered manually as matmul + mask + softmax: - -```python -def _should_lower_gemma4_decoder_attention_without_kernel(ir, node): - family = str(ir.meta.get("adapter_family") or ir.meta.get("family") or "").strip().lower() - component = str(ir.meta.get("component", "") or "").strip().lower() - if family != "gemma4" or component != "decoder": - return False - return node.op in {"attention", "scaled_dot_product_attention"} -``` - -That path exists because the transpiled Gemma4 decoder has some runtime/kernel -compatibility constraints that are easier to control in lowered graph code. - -## 13. Save Artifacts - -After lowering, `hf_model.py` writes the IR and bundle artifacts. - -For the monolithic path it saves: - -- `raw_ir.json` -- `optimized_ir.json` -- `graph.cactus` -- `graph_bindings.json` -- component subgraphs under `components/` - -For the component path it saves: - -- `raw_ir.json` -- `optimized_ir.json` -- per-component IRs -- `components/manifest.json` -- one lowered graph per component - -The component bundle writer is `_write_component_bundle()`: - -```python -manifest_path = bundle_dir / "manifest.json" -_write_json( - manifest_path, - { - "model_id": model_id, - "model_source": model_source, - "task": task, - "family": family, - "component_order": component_order, - "inputs": _serialize_json_compatible(inputs_metadata), - "components": manifest_components, - }, -) -``` - -Why both a `.cactus` graph and a binding manifest? - -Because `Graph.save()` stores the graph structure, but the file-backed weight -bindings need to be reattached later. The manifests carry: - -- logical input/output names -- runtime node ids -- paths to bound weights or saved Cactus tensor constants - -## 14. Optional Execution And Reference Compare - -If execution is enabled, the transpiler module will run the lowered graph immediately. - -By task: - -- text LM: run graph, report next token -- Gemma4 multimodal: run pipeline, report next token -- Parakeet TDT: run component encoder + custom decode loop, report transcript -- generic audio: run encoder graph and report output shapes - -Many paths also compare the transpiled output against a PyTorch reference forward -unless `--skip-reference-compare` is set. - -## 15. Run A Saved Bundle Later - -Saved bundles are executed by the native C++ `run` binary, not by Python. The -Python CLI just resolves the bundle path and exec's the binary: - -``` -cactus run - → python/cactus/cli/run.py:cmd_run - → subprocess.run([run_bin, bundle_dir, ...]) - → run invokes cactus_init(bundle_dir) from libcactus_engine -``` - -The runtime reads `components/manifest.json` from the bundle directory, mmaps -the `.weights` files, and loads the serialized `.cactus` graphs. Per-platform -`.mlpackage` files (Apple CoreML / future vendor platforms) are dispatched -to the matching hardware when present. - -## Major Situations And How They Are Handled - -### Hugging Face models that normally return caches or `ModelOutput` - -`capture_pytorch.py` forces export-safe kwargs: - -- `use_cache=False` -- `return_dict=False` - -This keeps the capture boundary tensor-only. - -### Weights directory exists but has no manifest - -`hf_model.py` fails fast: - -```python -if not manifest_path.exists(): - raise RuntimeError("weights_dir is missing weights_manifest.json") -``` - -This is intentional. Converted CQ/Cactus weights are the source of truth for -transpiled bundles; raw filename guessing is not part of the runtime contract. - -### Weights directory resolves, but no bindings match - -The transpile script fails hard and suggests re-running conversion: - -- component path: if component graphs resolve `0` bindings -- monolithic path: if full optimized IR resolves `0` bindings - -### Local snapshot exists but is incomplete - -`_load_model_source()` will reject incomplete local snapshots when -`--local-files-only` is set. - -### Gemma4 transformers support is missing locally - -`runtime_support.py` and `hf_model.py` search alternate site-packages and -patch import compatibility for: - -- missing `transformers.models.gemma4` -- missing `_lzma` -- torchvision import probes -- missing `torchvision::nms` -- missing `flex_attention.AuxRequest` - -### Gemma4 local checkpoint key mismatches - -`_repair_gemma4_checkpoint_weights()` remaps local checkpoint keys before loading -the model. - -### Gemma4 multimodal prompt drift - -The transpiler normalizes prompt and tokenization details to match native Cactus: - -- inserts missing image/audio placeholder tokens -- builds a Cactus-style chat prompt when requested -- expands HF newline-merge tokens `108` and `109` back into repeated newline token - `107` -- can replace processor-generated image/audio tensors with native-like versions - -### Parakeet TDT does not run as a single generic forward - -`tdt_runtime.py` provides: - -- a local PyTorch reconstruction of the model -- audio preprocessing -- a greedy TDT decode loop -- component specs for an encoder graph and a recurrent decoder-step graph - -### Some runtime symbols are not present in the current shared library - -`runtime_compat.py` monkey-patches the Python `Graph` wrapper so the transpiler can -still build graphs against the current runtime. It provides compatibility behavior -for items like: - -- generic `matmul` -- `gather(axis=0)` -- limited `conv2d` -- `not_equal` -- scalar math legalization - -### Opaque import is allowed, but execution may still fail later - -If `importers.py` does not know a target, it creates an opaque IR node instead of -dropping it. That preserves information, but lowering still needs a real lowering -case. Opaque nodes are therefore a sign that later lowering may stop with -`unsupported IR op`. - -## Artifact Layout - -A typical transpile artifact directory looks like this: - -```text -transpiled// - raw_ir.json - optimized_ir.json - graph.cactus - graph_bindings.json - result.json - components/ - manifest.json - vision_encoder/ - graph.cactus - bound_constants/ - ... - audio_encoder/ - graph.cactus - lm_encoder/ - graph.cactus - decoder/ - graph.cactus -``` - -Notes: - -- `graph.cactus` is the monolithic lowered graph, when one exists. -- `components/manifest.json` is the portable bundle manifest. -- `graph_bindings.json` and per-component `bound_constant_bindings` tell the loader - how to reattach mmap-backed weights and saved Cactus tensor constants. - -## File-By-File Map - -### Core transpiler files - -| File | Responsibility | -| --- | --- | -| `audio_preprocess.py` | WAV loading, resampling, generic log-mel extraction, Cactus audio frontend wrappers, Parakeet-native features, Gemma4-native audio features. | -| `capture_pytorch.py` | `torch.export` capture wrapper, export-safe kwargs injection, metadata collection, shape propagation, and `CapturedModel` creation. | -| `component_partition.py` | Heuristic component labeling of IR nodes/values and subgraph extraction into `audio_encoder` / `vision_encoder` / `lm_encoder` / `decoder`. | -| `component_pipeline.py` | Capture/optimize/lower loop for one component spec, plus runtime execution of a staged pipeline of components. | -| `graph_ir.py` | Definitions for `IRValue`, `IRNode`, `IRGraph`, deep-copy behavior, and IR consistency verification. | -| `import_ir.py` | Top-level importer from exported FX graph into `IRGraph`; resolves lifted constants, weight binding metadata, and inlined subgraphs. | -| `import_semantics.py` | Early semantic rewrites done immediately after import, especially attention and RoPE recognition. | -| `importers.py` | Large target-to-IR importer registry; converts normalized PyTorch ops into canonical `IRNode`s. | -| `lower.py` | Final lowering from optimized IR into runtime `Graph`, including op lowering, constant binding, fused op handling, and execution wrapper. | -| `model_adapters.py` | Family- and task-specific forward adapters, Gemma4 multimodal pipeline adapters, task canonicalization, and component spec builders. | -| `model_patterns.py` | Topology/op catalog used to annotate recognizable transformer structures after optimization. | -| `model_profiles.py` | Compact model-family profile constants, including compatibility modules, graph context defaults, and weight-name aliases used by fallback runtimes. | -| `multimodal_runtime.py` | Shared multimodal runtime helpers used by bundle execution: prompt normalization, image/audio limiting, and native-like multimodal preprocessing. | -| `normalize.py` | Target-name normalization (`aten.*` to canonical names) and dtype-name normalization into IR spellings. | -| `ops.py` | Canonical op registry and alias table for IR ops, including backend op names and attr schemas. | -| `optimize_graph.py` | Main fusion/optimization pass driver, pattern fusion, Gemma4 attention normalization, and gold-pattern annotation. | -| `runtime_support.py` | Shared import/runtime compatibility helpers for optional Transformers modules, torchvision probes, and flex-attention shims. | -| `tdt_runtime.py` | TDT fallback runtime: local recurrent decoder loop, config loader, and encoder/decoder component-spec generation. | -| `runtime_compat.py` | Runtime wrapper patches that make the current Python `Graph` API usable for transpiled graphs even when some C symbols are absent. | -| `weight_binding.py` | Resolves IR constants to converted `.weights` files through `weights_manifest.json`; also resolves default weights dirs. | -| `weight_compat.py` | Makes bound weights executable with the current runtime by materializing compatibility companion files when needed. | - -### Canonicalization subdirectory - -| File | Responsibility | -| --- | --- | -| `canonicalize/cleanup.py` | Multi-pass canonical cleanup, constant folding, no-op removal, attr normalization, and precision legalization. | -| `canonicalize/utils.py` | Shared graph rebuild, DCE, node removal, dtype helpers, and constant materialization utilities. | - -### Fusion subdirectory - -| File | Responsibility | -| --- | --- | -| `fusion/__init__.py` | Re-exports the fusion matchers used by `optimize_graph.py`. | -| `fusion/common.py` | Shared fusion helper types and graph traversal helpers. | -| `fusion/linear.py` | Matcher for linear projection patterns. | -| `fusion/rms_norm.py` | Matcher for RMSNorm computational subgraphs. | -| `fusion/rope.py` | Matcher for rotary embedding patterns. | -| `fusion/attention.py` | Matchers for attention, attention blocks, self-attention blocks, and GQA-style projection patterns. | -| `fusion/conv.py` | Matcher for Conformer-like convolution modules. | -| `fusion/lstm.py` | Matcher for LSTM cell subgraphs. | -| `fusion/mlp.py` | Matcher for gated MLP branches such as GELU/SiLU gate-up-down patterns. | -| `fusion/rel_pos_bias.py` | Matcher for relative-position-bias subgraphs. | -| `fusion/deltanet.py` | Matcher for fused gated DeltaNet patterns. | - -## CLI Usage - -## 1. Basic Convert - -The simplest invocation (internal command, not in the CLI menu) is: - -```bash -cactus convert -``` - -This quantizes the weights to CQ and then builds the runtime graph. The graph stage: - -- requires the converted Cactus CQ weights (produced by the same command, or passed - with `--weights-dir`; default lookup is `weights/`) -- resolves the runtime library -- launches `python/cactus/transpile/hf_model.py` -- saves artifacts under `weights//` (alongside the converted CQ weights; override with `--artifact-dir`) -- adds `--skip-execute`, so it does not run the graph unless you ask - -To stop after the CQ weights (skip the graph), pass `--weights-only`. To also execute -a reference run right after building the graph, pass `--execute-after-transpile`. - -## 2. Graph Options - -The graph-build options are parsed by `python/cactus/transpile/hf_model.py` and exposed -on `cactus convert`. - -Common ones: - -- `--task` -- `--prompt` -- `--input-ids` -- `--audio-file` -- `--image-file` -- `--weights-dir` -- `--artifact-dir` -- `--component-pipeline` -- `--skip-reference-compare` -- `--graph-filename` -- `--no-fuse-rms-norm` -- `--no-fuse-rope` -- `--no-fuse-attention` -- `--allow-unconverted-weights` for compiler-only debugging - -Examples: - -### Text causal LM - -```bash -cactus convert /path/to/converted_weights --bits 4 \ - --prompt "The capital of France is" \ - --artifact-dir ./transpiled/text_demo -``` - -### Gemma4 multimodal - -```bash -cactus convert \ - --task multimodal_causal_lm_logits \ - --image-file /path/to/image.jpg \ - --audio-file /path/to/audio.wav \ - --prompt "Describe what is happening in the image and audio." -``` - -### Whisper-style encoder capture - -```bash -cactus convert /path/to/converted_weights --bits 4 \ - --audio-file /path/to/sample.wav \ - --task encoder_hidden_states -``` - -### Parakeet TDT - -```bash -cactus convert \ - --audio-file /path/to/sample.wav \ - --task tdt_transcription -``` - -If you want the full direct script help: - -```bash -python -m cactus.transpile.hf_model --help -``` - -## 3. Run A Saved Transpiled Bundle - -Use the bundle root directory: - -```bash -cactus run /path/to/transpiled/model -``` - -### Text causal LM bundle - -```bash -cactus run /path/to/transpiled/model \ - --prompt "The capital of France is" -``` - -Or pass token ids directly: - -```bash -cactus run /path/to/transpiled/model \ - --input-ids 2,13,42 -``` - -### Multimodal Gemma4 bundle - -```bash -cactus run /path/to/transpiled/model \ - --image /path/to/image.jpg \ - --audio /path/to/audio.wav \ - --prompt "What do you observe?" \ - --system "Answer concisely." -``` - -### Audio encoder / transcription bundle - -```bash -cactus run /path/to/transpiled/model \ - --audio /path/to/sample.wav -``` - -### Save result payload - -```bash -cactus run /path/to/transpiled/model \ - --prompt "Hello" \ - --result-json ./result.json -``` - -Internally, `cli/model.py:ensure_runnable_bundle` checks whether the positional -argument is a local path with `components/manifest.json` and dispatches the -native `run` binary against it. - -## Current Runtime Limitations - -As of the current code: - -- Saved bundle execution is implemented for `causal_lm_logits`, - `multimodal_causal_lm_logits`, `encoder_hidden_states`, and - `tdt_transcription`. -- `ctc_logits` can be transpiled, but the runtime does not currently implement - a saved-bundle executor for that task. -- The component pipeline is model-family-specific; it is not a universal split-IR - path. - -## Useful Environment Variables - -- `CACTUS_TRANSPILER_WEIGHTS_DIR`: - fallback weights directory used by `weight_binding.py` -- `CACTUS_TRANSPILER_WEIGHTS_DIR_`: - family-specific weights directory override -- `CACTUS_TRANSPILER_DEBUG_MMAP=1`: - log constant binding behavior during lowering -- `CACTUS_GEMMA4_CAPTURE_FP32=1`: - temporarily promote some Gemma4 text modules to FP32 during capture on CPU - -## Practical Summary - -If you only want the high-level lifecycle, it is: - -1. `cactus convert` quantizes the weights, then launches `python/cactus/transpile/hf_model.py`. -2. The script infers a task, loads the HF model, and creates one example input batch. -3. `model_adapters.py` wraps the model into a smaller canonical task interface. -4. `capture_pytorch.py` runs `torch.export`. -5. `import_ir.py` + `importers.py` convert the export to `IRGraph`. -6. `import_semantics.py` adds early semantic rewrites. -7. `canonicalize/cleanup.py` normalizes and simplifies the IR. -8. `optimize_graph.py` fuses patterns into higher-level semantic ops. -9. `lower.py` builds a runtime `Graph`, mmaps bound weights when possible, and - materializes everything else. -10. The script writes `raw_ir.json`, `optimized_ir.json`, `graph.cactus`, and - `components/manifest.json`. -11. `cactus run` later resolves the bundle path and invokes the native `run` - binary, which mmaps the weights, loads the serialized graph, and runs it - via `cactus_init()` + `cactus_complete()` (see `cactus-engine`). diff --git a/python/cactus/transpile/__init__.py b/python/cactus/transpile/__init__.py deleted file mode 100644 index 47c8ee364..000000000 --- a/python/cactus/transpile/__init__.py +++ /dev/null @@ -1 +0,0 @@ -"""Cactus graph transpiler package.""" diff --git a/python/cactus/transpile/aten_ops.py b/python/cactus/transpile/aten_ops.py deleted file mode 100644 index 90cb587ec..000000000 --- a/python/cactus/transpile/aten_ops.py +++ /dev/null @@ -1,184 +0,0 @@ -from __future__ import annotations - -from dataclasses import dataclass -from typing import Any - - -@dataclass(frozen=True) -class AtenOpRule: - torch_prefix: str - ir_op: str - - -def canonical_torch_op(target: Any) -> str: - """Return a stable torch/export op spelling independent of Python repr noise.""" - - if isinstance(target, str): - value = target - else: - rendered = str(target).strip() - if rendered.startswith("aten.") or rendered.startswith("": "getitem", - "aten.view.default": "reshape", - "aten.reshape.default": "reshape", - "aten._unsafe_view.default": "reshape", - "aten.flatten.using_ints": "flatten", - "aten.t.default": "transpose", - "aten.gelu.erf": "gelu_erf", - # Keep the legacy importer key until the importer itself is renamed. - "aten.diff.default": "aten.diff.default", -} - - -_PREFIX_RULES: tuple[AtenOpRule, ...] = ( - AtenOpRule("aten.mul_", "multiply_inplace"), - AtenOpRule("aten._assert_tensor_metadata", "identity"), - AtenOpRule("aten.lift_fresh_copy", "identity"), - AtenOpRule("aten.clone", "identity"), - AtenOpRule("aten.alias", "identity"), - AtenOpRule("aten.detach", "identity"), - AtenOpRule("aten.dropout", "identity"), - AtenOpRule("aten.native_dropout", "identity"), - AtenOpRule("aten.arange.start", "arange"), - AtenOpRule("aten.arange", "arange"), - AtenOpRule("aten.add", "add"), - AtenOpRule("aten.sub", "subtract"), - AtenOpRule("aten.mul", "multiply"), - AtenOpRule("aten.floor_divide", "floor_divide"), - AtenOpRule("aten.div", "divide"), - AtenOpRule("aten.to.dtype", "precision_cast"), - AtenOpRule("aten.to.device", "identity"), - AtenOpRule("aten.type_as", "type_as"), - AtenOpRule("aten.neg", "negate"), - AtenOpRule("aten.ne.", "not_equal"), - AtenOpRule("aten.eq.", "equal"), - AtenOpRule("aten.gt.", "greater"), - AtenOpRule("aten.ge.", "greater_equal"), - AtenOpRule("aten.lt.", "less"), - AtenOpRule("aten.le.", "less_equal"), - AtenOpRule("aten.logical_or", "logical_or"), - AtenOpRule("aten.__or__", "logical_or"), - AtenOpRule("aten.logical_and", "logical_and"), - AtenOpRule("aten.__and__", "logical_and"), - AtenOpRule("aten.logical_not", "logical_not"), - AtenOpRule("aten.bitwise_not", "logical_not"), - AtenOpRule("aten.abs", "abs"), - AtenOpRule("aten.clamp", "clamp"), - AtenOpRule("aten.where", "where"), - AtenOpRule("aten.masked_scatter", "masked_scatter"), - AtenOpRule("aten.masked_fill", "masked_fill"), - AtenOpRule("aten.expand", "expand"), - AtenOpRule("aten.repeat", "repeat"), - AtenOpRule("aten.exp", "scalar_exp"), - AtenOpRule("aten.sqrt", "scalar_sqrt"), - AtenOpRule("aten.reciprocal", "reciprocal"), - AtenOpRule("aten.log", "scalar_log"), - AtenOpRule("aten.rsqrt", "rsqrt"), - AtenOpRule("aten.cos", "cos"), - AtenOpRule("aten.sin", "sin"), - AtenOpRule("aten.squeeze", "squeeze"), - AtenOpRule("aten.unsqueeze", "unsqueeze"), - AtenOpRule("aten.transpose", "transpose"), - AtenOpRule("aten.permute", "permute"), - AtenOpRule("aten.numpy_T", "numpy_T"), - AtenOpRule("aten.movedim", "movedim"), - AtenOpRule("aten.contiguous", "contiguous"), - AtenOpRule("aten.mm", "matmul"), - AtenOpRule("aten.matmul", "matmul"), - AtenOpRule("aten.bmm", "matmul"), - AtenOpRule("aten.linear", "linear"), - AtenOpRule("aten.addmm", "addmm"), - AtenOpRule("aten.relu", "relu"), - AtenOpRule("aten.silu", "silu"), - AtenOpRule("aten.gelu", "gelu"), - AtenOpRule("aten.sigmoid", "sigmoid"), - AtenOpRule("aten.glu", "glu"), - AtenOpRule("aten.softplus", "softplus"), - AtenOpRule("aten.tanh", "tanh"), - AtenOpRule("aten.softmax", "softmax"), - AtenOpRule("aten.scaled_dot_product_attention", "scaled_dot_product_attention"), - AtenOpRule("aten.sum", "sum"), - AtenOpRule("aten.mean", "mean"), - AtenOpRule("aten.all", "min"), - AtenOpRule("aten.var", "variance"), - AtenOpRule("aten.minimum", "minimum"), - AtenOpRule("aten.maximum", "maximum"), - AtenOpRule("aten.min", "min"), - AtenOpRule("aten.amin", "min"), - AtenOpRule("aten.max", "max"), - AtenOpRule("aten.amax", "max"), - AtenOpRule("aten.cat", "cat"), - AtenOpRule("aten.split_with_sizes", "split_with_sizes"), - AtenOpRule("aten.chunk", "chunk"), - AtenOpRule("aten.ones", "ones"), - AtenOpRule("aten.slice_scatter", "slice_scatter"), - AtenOpRule("aten.select_scatter", "select_scatter"), - AtenOpRule("aten.slice", "slice"), - AtenOpRule("aten.select", "index"), - AtenOpRule("aten.gather", "gather"), - AtenOpRule("aten.pad", "pad"), - AtenOpRule("aten.one_hot", "one_hot"), - AtenOpRule("aten.tril", "tril"), - AtenOpRule("aten.unfold", "unfold"), - AtenOpRule("aten.embedding", "embedding"), - AtenOpRule("aten.conv1d", "conv1d"), - AtenOpRule("aten.conv2d", "conv2d"), - AtenOpRule("aten.pow", "pow"), - AtenOpRule("aten.layer_norm", "layer_norm"), - AtenOpRule("aten.native_layer_norm", "layer_norm"), - AtenOpRule("aten.group_norm", "group_norm"), - AtenOpRule("aten.batch_norm", "batch_norm"), - AtenOpRule("aten._native_batch_norm_legit_no_training", "batch_norm"), - AtenOpRule("aten.rms_norm", "rms_norm"), -) - - -_PREFIX_RULES_BY_LENGTH: tuple[AtenOpRule, ...] = tuple( - sorted(_PREFIX_RULES, key=lambda rule: len(rule.torch_prefix), reverse=True) -) - - -def canonical_ir_op(target: Any) -> str: - torch_op = canonical_torch_op(target) - exact = _EXACT_RULES.get(torch_op) - if exact is not None: - return exact - if torch_op in {"aten.ne", "aten.eq", "aten.gt", "aten.ge", "aten.lt", "aten.le"}: - return { - "aten.ne": "not_equal", - "aten.eq": "equal", - "aten.gt": "greater", - "aten.ge": "greater_equal", - "aten.lt": "less", - "aten.le": "less_equal", - }[torch_op] - for rule in _PREFIX_RULES_BY_LENGTH: - if torch_op.startswith(rule.torch_prefix): - return rule.ir_op - return torch_op - - -def is_aten_op(target: Any) -> bool: - return canonical_torch_op(target).startswith("aten.") diff --git a/python/cactus/transpile/audio_preprocess.py b/python/cactus/transpile/audio_preprocess.py deleted file mode 100644 index 5e6d0aeea..000000000 --- a/python/cactus/transpile/audio_preprocess.py +++ /dev/null @@ -1,400 +0,0 @@ -from __future__ import annotations - -import os -import tempfile -from collections.abc import Iterator -from contextlib import contextmanager -from functools import lru_cache -from pathlib import Path - -import numpy as np -import torch -from scipy.io import wavfile -from scipy.signal import resample_poly - - -_PARAKEET_SAMPLE_RATE = 16000 -_PARAKEET_NUM_MELS = 128 -_PARAKEET_N_FFT = 512 -_PARAKEET_FRAME_LENGTH = 400 -_PARAKEET_HOP_LENGTH = 160 -_PARAKEET_PREEMPHASIS = 0.97 -_PARAKEET_LOG_FLOOR = np.float32(2**-24) -_DEFAULT_MAX_AUDIO_SECONDS = 30.0 - - -def audio_duration_limit_seconds() -> float: - raw = os.environ.get("CACTUS_TRANSPILER_MAX_AUDIO_SECONDS", str(_DEFAULT_MAX_AUDIO_SECONDS)) - try: - value = float(raw) - except (TypeError, ValueError): - return _DEFAULT_MAX_AUDIO_SECONDS - return max(0.0, value) - - -def normalize_audio_samples(samples: np.ndarray) -> np.ndarray: - if samples.dtype == np.uint8: - return ((samples.astype(np.float32) - 128.0) / 128.0).astype(np.float32) - if np.issubdtype(samples.dtype, np.integer): - info = np.iinfo(samples.dtype) - scale = float(max(abs(info.min), info.max)) - if scale <= 0: - scale = 1.0 - return (samples.astype(np.float32) / scale).astype(np.float32) - return samples.astype(np.float32) - - -def _limit_waveform_duration( - waveform: np.ndarray, - *, - sample_rate: int, - max_seconds: float | None, -) -> np.ndarray: - if max_seconds is None or max_seconds <= 0.0: - return waveform - max_samples = int(round(float(sample_rate) * float(max_seconds))) - if max_samples <= 0 or int(waveform.shape[0]) <= max_samples: - return waveform - return np.ascontiguousarray(waveform[:max_samples]) - - -def load_audio_waveform( - audio_file: str | Path, - *, - target_sample_rate: int, - max_seconds: float | None = None, -) -> np.ndarray: - if max_seconds is None: - max_seconds = audio_duration_limit_seconds() - sample_rate, waveform = wavfile.read(str(audio_file)) - if waveform.ndim > 1: - waveform = waveform.mean(axis=1) - waveform = normalize_audio_samples(np.asarray(waveform)) - if int(sample_rate) != int(target_sample_rate): - gcd = int(np.gcd(int(sample_rate), int(target_sample_rate))) - waveform = resample_poly( - waveform, - int(target_sample_rate) // gcd, - int(sample_rate) // gcd, - ).astype(np.float32) - sample_rate = int(target_sample_rate) - waveform = _limit_waveform_duration( - waveform.astype(np.float32), - sample_rate=int(target_sample_rate), - max_seconds=max_seconds, - ) - return waveform.astype(np.float32) - - -@contextmanager -def limited_audio_file( - audio_file: str | Path, - *, - target_sample_rate: int = _PARAKEET_SAMPLE_RATE, - max_seconds: float | None = None, -) -> Iterator[Path]: - if max_seconds is None: - max_seconds = audio_duration_limit_seconds() - if max_seconds <= 0.0: - yield Path(audio_file).expanduser().resolve() - return - - waveform = load_audio_waveform( - audio_file, - target_sample_rate=target_sample_rate, - max_seconds=max_seconds, - ) - temp_audio_path: Path | None = None - try: - with tempfile.NamedTemporaryFile(suffix=".wav", delete=False) as temp_file: - temp_audio_path = Path(temp_file.name) - wavfile.write( - str(temp_audio_path), - int(target_sample_rate), - (np.clip(waveform, -1.0, 1.0) * np.float32(32767.0)).astype(np.int16), - ) - yield temp_audio_path - finally: - if temp_audio_path is not None: - try: - temp_audio_path.unlink() - except OSError: - pass - - -def _hz_to_mel(hz: np.ndarray) -> np.ndarray: - return 2595.0 * np.log10(1.0 + hz / 700.0) - - -def _mel_to_hz(mel: np.ndarray) -> np.ndarray: - return 700.0 * (np.power(10.0, mel / 2595.0) - 1.0) - - -@lru_cache(maxsize=16) -def _mel_filter_bank(num_mels: int, sample_rate: int, n_fft: int) -> np.ndarray: - try: - from transformers.audio_utils import mel_filter_bank as transformers_mel_filter_bank # type: ignore - - return transformers_mel_filter_bank( - num_frequency_bins=n_fft // 2 + 1, - num_mel_filters=num_mels, - min_frequency=0.0, - max_frequency=sample_rate / 2.0, - sampling_rate=sample_rate, - norm="slaney", - mel_scale="slaney", - ).astype(np.float32) - except Exception: - pass - - num_freq_bins = n_fft // 2 + 1 - mel_points = np.linspace( - _hz_to_mel(np.array([0.0], dtype=np.float32))[0], - _hz_to_mel(np.array([sample_rate / 2.0], dtype=np.float32))[0], - num_mels + 2, - ) - hz_points = _mel_to_hz(mel_points) - bins = np.floor((n_fft + 1) * hz_points / sample_rate).astype(np.int64) - - filters = np.zeros((num_mels, num_freq_bins), dtype=np.float32) - for index in range(1, num_mels + 1): - left = max(0, min(num_freq_bins - 1, int(bins[index - 1]))) - center = max(left + 1, min(num_freq_bins - 1, int(bins[index]))) - right = max(center + 1, min(num_freq_bins, int(bins[index + 1]))) - for freq in range(left, center): - filters[index - 1, freq] = float(freq - left) / float(max(center - left, 1)) - for freq in range(center, right): - filters[index - 1, freq] = float(right - freq) / float(max(right - center, 1)) - return filters - - -def generic_log_mel_features( - waveform: np.ndarray, - *, - sample_rate: int, - num_mels: int, - n_fft: int, - hop_length: int, - frame_length: int, - preemphasis: float | None = None, - mel_floor: np.float32 = _PARAKEET_LOG_FLOOR, - normalize_active_frames_only: bool = True, -) -> tuple[np.ndarray, int]: - waveform_tensor = torch.from_numpy(waveform).to(torch.float32).unsqueeze(0) - if preemphasis is not None and abs(preemphasis) > 0.0: - waveform_tensor = torch.cat( - [ - waveform_tensor[:, :1], - waveform_tensor[:, 1:] - float(preemphasis) * waveform_tensor[:, :-1], - ], - dim=1, - ) - - window = torch.hann_window(frame_length, periodic=False, dtype=torch.float32) - stft = torch.stft( - waveform_tensor, - n_fft=n_fft, - hop_length=hop_length, - win_length=frame_length, - window=window, - return_complex=True, - center=True, - pad_mode="constant", - ) - magnitudes = torch.view_as_real(stft) - magnitudes = torch.sqrt(magnitudes.pow(2).sum(-1)) - power = magnitudes.pow(2) - - feature_length = max(1, int(waveform.shape[0] // hop_length)) - mel_filters = torch.from_numpy(_mel_filter_bank(num_mels, sample_rate, n_fft)).to(torch.float32) - mel_spec = power.transpose(1, 2) @ mel_filters - mel_spec = torch.log(torch.clamp(mel_spec, min=float(mel_floor))) - - if normalize_active_frames_only: - attention_mask = ( - torch.arange(mel_spec.shape[1], dtype=torch.long, device=mel_spec.device).unsqueeze(0) - < feature_length - ) - mask = attention_mask.unsqueeze(-1).to(torch.float32) - masked = mel_spec * mask - mean = masked.sum(dim=1) / float(feature_length) - centered = masked - mean.unsqueeze(1) - denom = float(max(1, feature_length - 1)) - variance = (centered.pow(2) * mask).sum(dim=1) / denom - mel_spec = (mel_spec - mean.unsqueeze(1)) / torch.sqrt(variance.unsqueeze(1) + np.float32(1e-5)) - mel_spec = mel_spec * mask - else: - # Native Parakeet normalizes across the full centered-STFT frame count, - # then trims to waveform_samples // hop_length. Keep that order for parity. - frame_count = int(mel_spec.shape[1]) - mean = mel_spec.mean(dim=1) - centered = mel_spec - mean.unsqueeze(1) - denom = float(max(1, frame_count - 1)) - variance = centered.pow(2).sum(dim=1) / denom - mel_spec = centered / torch.sqrt(variance.unsqueeze(1) + np.float32(1e-5)) - return mel_spec[0].cpu().numpy().astype(np.float32), feature_length - - -def prepare_cactus_audio_features( - audio_file: str | Path, - *, - model_type: str, - expected_frames: int | None, - expected_mels: int, - torch_dtype: torch.dtype, - layout: str = "frames_mels", -) -> tuple[torch.Tensor, int]: - from cactus.bindings.cactus import cactus_preprocess_audio_features - - capacity_frames = int(expected_frames) if isinstance(expected_frames, int) and expected_frames > 0 else 0 - if capacity_frames <= 0: - lowered_model_type = model_type.lower() - if "whisper" in lowered_model_type: - capacity_frames = 3000 - else: - sample_rate, waveform = wavfile.read(str(audio_file)) - samples = int(waveform.shape[0]) if hasattr(waveform, "shape") else len(waveform) - if sample_rate <= 0: - sample_rate = _PARAKEET_SAMPLE_RATE - capacity_frames = max( - 1, - int(np.ceil(samples * (_PARAKEET_SAMPLE_RATE / float(sample_rate)) / _PARAKEET_HOP_LENGTH)) + 8, - ) - if "whisper" in model_type.lower(): - capacity_frames = max(capacity_frames, 3000) - capacity = max(1, capacity_frames * int(expected_mels)) - with limited_audio_file(audio_file, target_sample_rate=_PARAKEET_SAMPLE_RATE) as frontend_audio_file: - values, mel_bins, frames = cactus_preprocess_audio_features( - str(frontend_audio_file), - model_type, - int(expected_mels), - capacity, - ) - if int(mel_bins) != int(expected_mels): - raise ValueError(f"Cactus audio frontend returned {mel_bins} mel bins, expected {expected_mels}") - feature_array = np.asarray(values, dtype=np.float32).reshape(int(mel_bins), int(frames)) - active_frames = int(frames) - if isinstance(expected_frames, int) and expected_frames > 0: - active_frames = min(active_frames, expected_frames) - if layout == "frames_mels": - feature_array = feature_array[:, :active_frames].T - if isinstance(expected_frames, int) and expected_frames > active_frames: - feature_array = np.pad(feature_array, ((0, expected_frames - active_frames), (0, 0)), mode="constant") - elif layout == "mels_frames": - feature_array = feature_array[:, :active_frames] - if isinstance(expected_frames, int) and expected_frames > active_frames: - feature_array = np.pad(feature_array, ((0, 0), (0, expected_frames - active_frames)), mode="constant") - else: - raise ValueError(f"unsupported Cactus audio feature layout: {layout!r}") - tensor = torch.from_numpy(np.ascontiguousarray(feature_array)).unsqueeze(0).to(dtype=torch_dtype) - return tensor, active_frames - - -def prepare_native_parakeet_audio_features( - audio_file: str | Path, - *, - expected_frames: int | None, - expected_mels: int, - torch_dtype: torch.dtype, -) -> tuple[torch.Tensor, int]: - try: - return prepare_cactus_audio_features( - audio_file, - model_type="parakeet", - expected_frames=expected_frames, - expected_mels=expected_mels, - torch_dtype=torch_dtype, - ) - except Exception: - pass - - waveform = load_audio_waveform(audio_file, target_sample_rate=_PARAKEET_SAMPLE_RATE) - features, feature_length = generic_log_mel_features( - waveform, - sample_rate=_PARAKEET_SAMPLE_RATE, - num_mels=expected_mels or _PARAKEET_NUM_MELS, - n_fft=_PARAKEET_N_FFT, - hop_length=_PARAKEET_HOP_LENGTH, - frame_length=_PARAKEET_FRAME_LENGTH, - preemphasis=_PARAKEET_PREEMPHASIS, - mel_floor=_PARAKEET_LOG_FLOOR, - normalize_active_frames_only=False, - ) - active_frames = feature_length - if isinstance(expected_frames, int) and expected_frames > 0: - active_frames = min(active_frames, expected_frames) - features = features[:active_frames, :] - if features.shape[1] != expected_mels: - raise ValueError( - f"feature mel dimension mismatch: expected {expected_mels}, got {features.shape[1]}" - ) - if isinstance(expected_frames, int) and expected_frames > active_frames: - features = np.pad(features, ((0, expected_frames - active_frames), (0, 0)), mode="constant") - tensor = torch.from_numpy(np.ascontiguousarray(features)).unsqueeze(0).to(dtype=torch_dtype) - return tensor, active_frames - - -def prepare_native_gemma4_audio_features( - audio_file: str | Path, - *, - expected_mels: int, - torch_dtype: torch.dtype, - max_seconds: float | None = None, - pad_to_max_seconds: bool = True, -) -> tuple[torch.Tensor, torch.Tensor, int]: - if max_seconds is None: - max_seconds = audio_duration_limit_seconds() - waveform = load_audio_waveform( - audio_file, - target_sample_rate=_PARAKEET_SAMPLE_RATE, - max_seconds=max_seconds, - ) - # Native Gemma4 pads to the nearest 320 samples, adds 160 samples of - # semicausal left padding, and then runs a 321-point analysis frame. - estimated_frames = int(np.ceil((len(waveform) + 640) / 160.0)) + 8 - capacity = max(1, estimated_frames * int(expected_mels)) - values, mel_bins, frames = prepare_cactus_audio_features_raw( - audio_file, - model_type="gemma4", - mel_bins=int(expected_mels), - capacity=capacity, - max_seconds=max_seconds, - ) - if int(mel_bins) != int(expected_mels): - raise ValueError(f"Cactus Gemma4 audio frontend returned {mel_bins} mel bins, expected {expected_mels}") - max_frames = estimated_frames - if max_seconds is not None and max_seconds > 0.0: - max_samples = int(round(float(_PARAKEET_SAMPLE_RATE) * float(max_seconds))) - max_frames = int(np.ceil((max_samples + 640) / 160.0)) + 8 - active_frames = min(int(frames), int(max_frames)) - features = np.asarray(values, dtype=np.float32).reshape(int(frames), int(mel_bins))[:active_frames, :] - output_frames = max(active_frames, int(max_frames)) if pad_to_max_seconds else active_frames - if output_frames > active_frames: - features = np.pad(features, ((0, output_frames - active_frames), (0, 0)), mode="constant") - feature_tensor = torch.from_numpy(np.ascontiguousarray(features)).unsqueeze(0).to(dtype=torch_dtype) - mask = torch.zeros((1, int(output_frames)), dtype=torch.bool) - mask[:, :active_frames] = True - return feature_tensor, mask, int(active_frames) - - -def prepare_cactus_audio_features_raw( - audio_file: str | Path, - *, - model_type: str, - mel_bins: int, - capacity: int, - max_seconds: float | None = None, -) -> tuple[list[float], int, int]: - from cactus.bindings.cactus import cactus_preprocess_audio_features - - with limited_audio_file( - audio_file, - target_sample_rate=_PARAKEET_SAMPLE_RATE, - max_seconds=max_seconds, - ) as frontend_audio_file: - return cactus_preprocess_audio_features( - str(frontend_audio_file), - model_type, - int(mel_bins), - int(capacity), - ) diff --git a/python/cactus/transpile/canonicalize/cleanup.py b/python/cactus/transpile/canonicalize/cleanup.py deleted file mode 100644 index b7f4de2aa..000000000 --- a/python/cactus/transpile/canonicalize/cleanup.py +++ /dev/null @@ -1,704 +0,0 @@ -from __future__ import annotations - -from collections import Counter - -import torch - -from cactus.transpile.graph_ir import IRGraph -from cactus.transpile.graph_ir import IRNode -from cactus.transpile.graph_ir import IRValue -from cactus.transpile.graph_ir import verify_ir -from cactus.transpile.ops import canonicalize_op -from cactus.transpile.ops import get_op -from cactus.transpile.ops import has_op -from cactus.transpile.canonicalize.utils import bypass_unary_node -from cactus.transpile.canonicalize.utils import dce -from cactus.transpile.canonicalize.utils import ir_dtype_to_torch -from cactus.transpile.canonicalize.utils import materialize_constant_output -from cactus.transpile.canonicalize.utils import normalize_dim -from cactus.transpile.canonicalize.utils import normalize_dtype_name -from cactus.transpile.canonicalize.utils import numel -from cactus.transpile.canonicalize.utils import rebuild_graph - - -COMPILER_SUPPORTED_OPS = { - "arange", - "ones", -} - -CANONICAL_UNSUPPORTED_RENAMES = { - "aten.__and__.Tensor": "logical_and", - "aten.__or__.Tensor": "logical_or", - "aten.cumsum.default": "cumsum", - "aten.eq.Scalar": "equal", - "aten.eq.Tensor": "equal", - "aten.ge.Scalar": "greater_equal", - "aten.ge.Tensor": "greater_equal", - "aten.gt.Scalar": "greater", - "aten.gt.Tensor": "greater", - "aten.index.Tensor": "advanced_index", - "aten.le.Scalar": "less_equal", - "aten.le.Tensor": "less_equal", - "aten.lt.Scalar": "less", - "aten.lt.Tensor": "less", -} - -FP32_SUPPORTED_ALL_INPUT_OPS = { - "add", - "subtract", - "multiply", - "divide", - "precision_cast", - "view", - "permute", - "transpose", - "slice", - "index", - "advanced_index", - "gather", - "matmul", - "linear", - "addmm", - "layer_norm", - "batch_norm", - "glu", - "sum", - "mean", - "variance", - "min", - "max", - "scalar_add", - "scalar_subtract", - "scalar_subtract_reverse", - "scalar_multiply", - "scalar_divide", - "scalar_floor_divide", - "scalar_not_equal", - "scalar_equal", - "scalar_greater", - "scalar_greater_equal", - "scalar_less", - "scalar_less_equal", - "abs", - "negate", - "pow", - "scalar_exp", - "scalar_sqrt", - "scalar_log", - "scalar_cos", - "scalar_sin", -} - -FP32_SUPPORTED_INPUT_INDICES = { - # The embedding op dequantizes its weight (input 0) directly (CQ or FP16 when - # bound), and indices (input 1) stay integral — neither needs fp16 legalization. - "embedding": {0, 1}, -} - -FP16_ONLY_OUTPUT_OPS = { - "scalar_divide_reverse", - "not_equal", - "equal", - "greater", - "greater_equal", - "less", - "less_equal", - "logical_and", - "logical_or", - "logical_not", - "one_hot", - "scalar_not_equal", - "scalar_equal", - "scalar_greater", - "scalar_greater_equal", - "scalar_less", - "scalar_less_equal", - "cat", - "softmax", - "rms_norm", - "embedding", - "group_norm", - "lstm_cell", - "relu", - "silu", - "gelu", - "gelu_erf", - "sigmoid", - "tanh", - "attention", - "attention_block", - "conv_module", - "dense_mlp_tq_fused", - "rope", -} - - -def canonicalize_exported_graph(graph: IRGraph, *, max_passes: int = 8) -> IRGraph: - verify_ir(graph) - - for _ in range(max_passes): - changed = False - - for node_id in list(graph.order): - node = graph.nodes.get(node_id) - if node is None: - continue - if _canonicalize_node(graph, node): - changed = True - - rebuild_graph(graph) - - for node_id in list(graph.order): - node = graph.nodes.get(node_id) - if node is None: - continue - if _simplify_node(graph, node): - changed = True - - rebuild_graph(graph) - - if _legalize_precisions(graph): - changed = True - rebuild_graph(graph) - - dce(graph) - - if not changed: - break - - unsupported_counts = summarize_unsupported_ops(graph) - graph.meta = dict(graph.meta) - graph.meta["canonical_unsupported_ops"] = tuple(sorted(unsupported_counts)) - graph.meta["canonical_unsupported_op_counts"] = unsupported_counts - verify_ir(graph) - return graph - - -def summarize_unsupported_ops(graph: IRGraph) -> dict[str, int]: - counts: Counter[str] = Counter() - for node_id in graph.order: - node = graph.nodes[node_id] - if not _is_supported_without_backend_kernel(node.op): - counts[node.op] += 1 - return dict(sorted(counts.items())) - - -def _canonicalize_node(graph: IRGraph, node: IRNode) -> bool: - original_op = node.op - - if original_op in {"aten.new_ones.default", "aten.ones.default"}: - return _materialize_new_ones_constant(graph, node) - - if original_op in {"aten.new_zeros.default", "aten.zeros.default", "new_zeros"}: - return _materialize_new_zeros_constant(graph, node) - - if original_op == "arange": - return _materialize_arange_constant(graph, node) - - if original_op in {"identity", "contiguous"}: - return bypass_unary_node(graph, node) - - if original_op in CANONICAL_UNSUPPORTED_RENAMES: - node.op = CANONICAL_UNSUPPORTED_RENAMES[original_op] - return True - - if original_op == "multiply_inplace": - node.op = "multiply" - return True - - if original_op == "view" and "shape" not in node.attrs: - return _rewrite_view_from_output_shape(graph, node) - - if original_op in {"reshape", "flatten", "unsqueeze", "squeeze"}: - return _rewrite_view_from_output_shape(graph, node) - - if original_op == "expand": - return _rewrite_expand_if_view_equivalent(graph, node) - - if original_op == "transpose": - return _rewrite_transpose_as_permute(graph, node) - - if original_op == "movedim": - return _rewrite_movedim_as_permute(graph, node) - - if original_op == "type_as": - return _rewrite_type_as(graph, node) - - if original_op == "concat": - node.op = "cat" - if "dim" in node.attrs and "axis" not in node.attrs: - node.attrs = {**node.attrs, "axis": int(node.attrs["dim"])} - return True - - if original_op == "scaled_dot_product_attention": - node.op = "attention" - node.attrs = { - "scale": float(node.attrs.get("scale", 0.0)), - "is_causal": bool(node.attrs.get("is_causal", False)), - "window_size": int(node.attrs.get("window_size", 0)), - **({"mask": node.attrs["mask"]} if "mask" in node.attrs else {}), - **({"additive_mask": bool(node.attrs.get("additive_mask", False))} if "additive_mask" in node.attrs else {}), - } - return True - - canonical_op = canonicalize_op(original_op) - if canonical_op != original_op: - node.op = canonical_op - _normalize_attr_spellings(node) - if has_op(node.op): - _filter_attrs_to_schema(node) - return True - - if has_op(node.op): - _normalize_attr_spellings(node) - _filter_attrs_to_schema(node) - return True - - return False - - -def _simplify_node(graph: IRGraph, node: IRNode) -> bool: - if _fold_constant_node(graph, node): - return True - - if node.op == "precision_cast": - return _remove_noop_precision_cast(graph, node) - - if node.op == "view": - if _remove_noop_view(graph, node): - return True - return _collapse_view_chain(graph, node) - - if node.op == "slice": - return _remove_noop_slice(graph, node) - - if node.op == "scalar_add": - return _remove_noop_scalar(graph, node, neutral=0.0) - - if node.op == "scalar_subtract": - return _remove_noop_scalar(graph, node, neutral=0.0) - - if node.op == "scalar_multiply": - return _remove_noop_scalar(graph, node, neutral=1.0) - - if node.op == "scalar_divide": - return _remove_noop_scalar(graph, node, neutral=1.0) - - return False - - -def _fold_constant_node(graph: IRGraph, node: IRNode) -> bool: - if len(node.outputs) != 1: - return False - - output_value = graph.values.get(node.outputs[0]) - if output_value is None: - return False - - if node.op in FP16_ONLY_OUTPUT_OPS: - output_value.dtype = "fp16" - node.meta["dtype"] = "fp16" - - output_torch_dtype = ir_dtype_to_torch(output_value.dtype) if output_value.dtype is not None else None - - constant_inputs = [graph.constants.get(value_id) for value_id in node.inputs] - if any(_is_external_weight_constant(graph, value_id) for value_id in node.inputs): - return False - if node.op == "view" and len(constant_inputs) == 1 and isinstance(constant_inputs[0], torch.Tensor): - if output_value.shape is None: - return False - return materialize_constant_output( - graph, - node, - constant_inputs[0].reshape(tuple(int(dim) for dim in output_value.shape)).to(dtype=output_torch_dtype), - ) - - if node.op == "expand" and len(constant_inputs) == 1 and isinstance(constant_inputs[0], torch.Tensor): - if output_value.shape is None: - return False - return materialize_constant_output( - graph, - node, - constant_inputs[0].expand(tuple(int(dim) for dim in output_value.shape)).clone().to(dtype=output_torch_dtype), - ) - - if node.op == "precision_cast" and len(constant_inputs) == 1 and isinstance(constant_inputs[0], torch.Tensor): - dtype = normalize_dtype_name(node.attrs.get("dtype")) - if dtype is None: - return False - return materialize_constant_output(graph, node, constant_inputs[0].to(dtype=ir_dtype_to_torch(dtype))) - - if node.op in {"cos", "scalar_cos"} and len(constant_inputs) == 1 and isinstance(constant_inputs[0], torch.Tensor): - return materialize_constant_output(graph, node, torch.cos(constant_inputs[0])) - - if node.op in {"sin", "scalar_sin"} and len(constant_inputs) == 1 and isinstance(constant_inputs[0], torch.Tensor): - return materialize_constant_output(graph, node, torch.sin(constant_inputs[0])) - - if node.op in {"scalar_add", "scalar_subtract", "scalar_multiply", "scalar_divide"}: - if len(constant_inputs) != 1 or not isinstance(constant_inputs[0], torch.Tensor): - return False - scalar = float(node.attrs["value"]) - if node.op == "scalar_add": - result = constant_inputs[0] + scalar - elif node.op == "scalar_subtract": - result = constant_inputs[0] - scalar - elif node.op == "scalar_multiply": - result = constant_inputs[0] * scalar - else: - result = constant_inputs[0] / scalar - return materialize_constant_output(graph, node, result.to(dtype=output_torch_dtype)) - - if node.op in {"scalar_not_equal", "scalar_equal", "scalar_greater", "scalar_greater_equal", "scalar_less", "scalar_less_equal"}: - if len(constant_inputs) != 1 or not isinstance(constant_inputs[0], torch.Tensor): - return False - scalar = float(node.attrs["value"]) - if node.op == "scalar_not_equal": - result = constant_inputs[0] != scalar - elif node.op == "scalar_equal": - result = constant_inputs[0] == scalar - elif node.op == "scalar_greater": - result = constant_inputs[0] > scalar - elif node.op == "scalar_greater_equal": - result = constant_inputs[0] >= scalar - elif node.op == "scalar_less": - result = constant_inputs[0] < scalar - else: - result = constant_inputs[0] <= scalar - return materialize_constant_output(graph, node, result.to(dtype=output_torch_dtype)) - - if node.op in {"not_equal", "equal", "greater", "greater_equal", "less", "less_equal", "logical_and", "logical_or"}: - if len(constant_inputs) != 2 or not all(isinstance(value, torch.Tensor) for value in constant_inputs): - return False - lhs, rhs = constant_inputs - if node.op == "not_equal": - result = lhs != rhs - elif node.op == "equal": - result = lhs == rhs - elif node.op == "greater": - result = lhs > rhs - elif node.op == "greater_equal": - result = lhs >= rhs - elif node.op == "less": - result = lhs < rhs - elif node.op == "less_equal": - result = lhs <= rhs - elif node.op == "logical_and": - result = torch.logical_and(lhs != 0, rhs != 0) - else: - result = torch.logical_or(lhs != 0, rhs != 0) - return materialize_constant_output(graph, node, result.to(dtype=output_torch_dtype)) - - if node.op == "logical_not": - if len(constant_inputs) != 1 or not isinstance(constant_inputs[0], torch.Tensor): - return False - result = torch.logical_not(constant_inputs[0] != 0) - return materialize_constant_output(graph, node, result.to(dtype=output_torch_dtype)) - - return False - - -def _is_external_weight_constant(graph: IRGraph, value_id: str) -> bool: - if value_id not in graph.constants: - return False - value = graph.values.get(value_id) - if value is None or not isinstance(value.meta, dict): - return False - meta = value.meta - return ( - str(meta.get("kind", "") or "") == "weight" - or bool(str(meta.get("source_name", "") or "")) - or bool(str(meta.get("path", "") or "")) - ) - - -def _legalize_precisions(graph: IRGraph) -> bool: - changed = False - counter = 0 - - for node_id in list(graph.order): - node = graph.nodes.get(node_id) - if node is None: - continue - - if node.op in FP32_SUPPORTED_ALL_INPUT_OPS: - continue - - inserted: list[IRNode] = [] - allowed_indices = FP32_SUPPORTED_INPUT_INDICES.get(node.op, set()) - for input_index, value_id in enumerate(list(node.inputs)): - if input_index in allowed_indices: - continue - value = graph.values.get(value_id) - if value is None or normalize_dtype_name(value.dtype) != "fp32": - continue - - counter += 1 - cast_node_id = _unique_ir_id( - f"{node.id}_legalize_fp16_{counter}", - graph.nodes.keys(), - graph.order, - ) - cast_value_id = _unique_ir_id( - f"{value_id}_legalize_fp16_{counter}", - graph.values.keys(), - graph.constants.keys(), - ) - cast_node = IRNode( - id=cast_node_id, - op="precision_cast", - inputs=[value_id], - outputs=[cast_value_id], - attrs={"dtype": "fp16"}, - meta={"inserted_by": "canonicalize_exported_graph"}, - ) - graph.values[cast_value_id] = IRValue( - id=cast_value_id, - shape=value.shape, - dtype="fp16", - producer=cast_node_id, - users=[], - ) - node.inputs[input_index] = cast_value_id - inserted.append(cast_node) - changed = True - - if inserted: - insert_at = graph.order.index(node_id) - for offset, cast_node in enumerate(inserted): - graph.nodes[cast_node.id] = cast_node - graph.order.insert(insert_at + offset, cast_node.id) - - if node.op in FP16_ONLY_OUTPUT_OPS: - for output_id in node.outputs: - output_value = graph.values.get(output_id) - if output_value is not None: - output_value.dtype = "fp16" - node.meta["dtype"] = "fp16" - - return changed - - -def _unique_ir_id(base: str, *existing_groups) -> str: - existing: set[str] = set() - for group in existing_groups: - existing.update(str(item) for item in group) - if base not in existing: - return base - suffix = 1 - while True: - candidate = f"{base}_{suffix}" - if candidate not in existing: - return candidate - suffix += 1 - - -def _rewrite_view_from_output_shape(graph: IRGraph, node: IRNode) -> bool: - if not node.outputs: - return False - output_value = graph.values.get(node.outputs[0]) - if output_value is None or output_value.shape is None: - return False - node.op = "view" - node.attrs = {"shape": tuple(int(dim) for dim in output_value.shape)} - return True - - -def _rewrite_expand_if_view_equivalent(graph: IRGraph, node: IRNode) -> bool: - if not node.inputs or not node.outputs: - return False - input_value = graph.values.get(node.inputs[0]) - output_value = graph.values.get(node.outputs[0]) - if input_value is None or output_value is None: - return False - if input_value.shape is None or output_value.shape is None: - return False - if numel(input_value.shape) != numel(output_value.shape): - node.meta["canonicalization_blocked"] = "expand_requires_broadcast" - return False - node.op = "view" - node.attrs = {"shape": tuple(int(dim) for dim in output_value.shape)} - return True - - -def _rewrite_transpose_as_permute(graph: IRGraph, node: IRNode) -> bool: - if not node.inputs: - return False - input_value = graph.values.get(node.inputs[0]) - if input_value is None or input_value.shape is None: - return False - rank = len(input_value.shape) - dim0 = normalize_dim(int(node.attrs["dim0"]), rank) - dim1 = normalize_dim(int(node.attrs["dim1"]), rank) - permutation = list(range(rank)) - permutation[dim0], permutation[dim1] = permutation[dim1], permutation[dim0] - node.op = "permute" - node.attrs = {"permutation": tuple(permutation)} - return True - - -def _rewrite_movedim_as_permute(graph: IRGraph, node: IRNode) -> bool: - if not node.inputs: - return False - input_value = graph.values.get(node.inputs[0]) - if input_value is None or input_value.shape is None: - return False - rank = len(input_value.shape) - source = [normalize_dim(int(dim), rank) for dim in node.attrs["source"]] - destination = [normalize_dim(int(dim), rank) for dim in node.attrs["destination"]] - if len(source) != len(destination): - return False - - remaining = [dim for dim in range(rank) if dim not in source] - permutation = list(remaining) - for dst, src in sorted(zip(destination, source)): - permutation.insert(dst, src) - - node.op = "permute" - node.attrs = {"permutation": tuple(permutation)} - return True - - -def _rewrite_type_as(graph: IRGraph, node: IRNode) -> bool: - if not node.inputs: - return False - target_dtype = None - if node.outputs: - output_value = graph.values.get(node.outputs[0]) - if output_value is not None: - target_dtype = output_value.dtype - if target_dtype is None and len(node.inputs) > 1: - target_value = graph.values.get(node.inputs[1]) - if target_value is not None: - target_dtype = target_value.dtype - if target_dtype is None: - return False - node.op = "precision_cast" - node.inputs = [node.inputs[0]] - node.attrs = {"dtype": normalize_dtype_name(target_dtype)} - return True - - -def _normalize_attr_spellings(node: IRNode) -> None: - attrs = dict(node.attrs) - if "dim" in attrs and "axis" not in attrs and node.op in {"softmax", "sum", "mean", "variance", "min", "max", "cat", "slice", "index"}: - attrs["axis"] = attrs.pop("dim") - if node.op == "permute" and "dims" in attrs and "permutation" not in attrs: - attrs["permutation"] = tuple(int(dim) for dim in attrs.pop("dims")) - if node.op == "slice": - attrs.setdefault("step", 1) - if node.op == "precision_cast" and "dtype" in attrs: - attrs["dtype"] = normalize_dtype_name(attrs["dtype"]) - node.attrs = attrs - - -def _filter_attrs_to_schema(node: IRNode) -> None: - schema = get_op(node.op) - if not schema.attrs: - node.attrs = {} - return - node.attrs = {name: node.attrs[name] for name in schema.attrs if name in node.attrs} - - -def _remove_noop_precision_cast(graph: IRGraph, node: IRNode) -> bool: - if len(node.inputs) != 1 or len(node.outputs) != 1: - return False - input_value = graph.values.get(node.inputs[0]) - output_value = graph.values.get(node.outputs[0]) - target_dtype = normalize_dtype_name(node.attrs.get("dtype")) - input_dtype = None if input_value is None else normalize_dtype_name(input_value.dtype) - output_dtype = None if output_value is None else normalize_dtype_name(output_value.dtype) - if target_dtype is None or target_dtype == input_dtype or (input_dtype is not None and output_dtype == input_dtype): - return bypass_unary_node(graph, node) - return False - - -def _remove_noop_view(graph: IRGraph, node: IRNode) -> bool: - if len(node.inputs) != 1 or len(node.outputs) != 1: - return False - input_value = graph.values.get(node.inputs[0]) - output_value = graph.values.get(node.outputs[0]) - if input_value is None or output_value is None: - return False - if input_value.shape is None or output_value.shape is None: - return False - if tuple(input_value.shape) == tuple(output_value.shape): - return bypass_unary_node(graph, node) - return False - - -def _collapse_view_chain(graph: IRGraph, node: IRNode) -> bool: - if len(node.inputs) != 1: - return False - input_value = graph.values.get(node.inputs[0]) - if input_value is None or input_value.producer is None: - return False - producer = graph.nodes.get(input_value.producer) - if producer is None or producer.op != "view" or len(producer.inputs) != 1: - return False - node.inputs = [producer.inputs[0]] - return True - - -def _remove_noop_slice(graph: IRGraph, node: IRNode) -> bool: - if len(node.inputs) != 1 or len(node.outputs) != 1: - return False - input_value = graph.values.get(node.inputs[0]) - output_value = graph.values.get(node.outputs[0]) - if input_value is None or output_value is None: - return False - if input_value.shape is None or output_value.shape is None: - return False - if tuple(input_value.shape) == tuple(output_value.shape): - return bypass_unary_node(graph, node) - return False - - -def _remove_noop_scalar(graph: IRGraph, node: IRNode, *, neutral: float) -> bool: - value = node.attrs.get("value") - if value is None: - return False - if float(value) != float(neutral): - return False - return bypass_unary_node(graph, node) - - -def _materialize_new_ones_constant(graph: IRGraph, node: IRNode) -> bool: - if len(node.outputs) != 1: - return False - output_id = node.outputs[0] - output_value = graph.values.get(output_id) - if output_value is None or output_value.shape is None or output_value.dtype is None: - return False - torch_dtype = ir_dtype_to_torch(output_value.dtype) - const = torch.ones(tuple(int(dim) for dim in output_value.shape), dtype=torch_dtype) - return materialize_constant_output(graph, node, const) - - -def _materialize_new_zeros_constant(graph: IRGraph, node: IRNode) -> bool: - if len(node.outputs) != 1: - return False - output_id = node.outputs[0] - output_value = graph.values.get(output_id) - if output_value is None or output_value.shape is None or output_value.dtype is None: - return False - torch_dtype = ir_dtype_to_torch(output_value.dtype) - const = torch.zeros(tuple(int(dim) for dim in output_value.shape), dtype=torch_dtype) - return materialize_constant_output(graph, node, const) - - -def _materialize_arange_constant(graph: IRGraph, node: IRNode) -> bool: - if len(node.outputs) != 1: - return False - start = int(node.attrs.get("start", 0)) - end = node.attrs.get("end") - if end is None: - return False - step = int(node.attrs.get("step", 1)) - dtype = normalize_dtype_name(node.attrs.get("dtype", "fp32")) - torch_dtype = ir_dtype_to_torch(dtype) - const = torch.arange(start, int(end), step=step, dtype=torch_dtype) - return materialize_constant_output(graph, node, const) - - -def _is_supported_without_backend_kernel(op_name: str) -> bool: - return has_op(op_name) or op_name in COMPILER_SUPPORTED_OPS diff --git a/python/cactus/transpile/canonicalize/utils.py b/python/cactus/transpile/canonicalize/utils.py deleted file mode 100644 index 7ee113b9a..000000000 --- a/python/cactus/transpile/canonicalize/utils.py +++ /dev/null @@ -1,277 +0,0 @@ -from __future__ import annotations - -from typing import Any - -import torch - -from cactus.transpile.graph_ir import IRGraph -from cactus.transpile.graph_ir import IRNode -from cactus.transpile.graph_ir import IRValue -from cactus.transpile.graph_ir import verify_ir -from cactus.transpile.model_profiles import profile_for_family - - -def normalize_dtype_name(dtype: Any) -> str | None: - if dtype is None: - return None - if not isinstance(dtype, str): - return str(dtype) - if dtype.startswith("torch."): - torch_name = dtype.split(".", 1)[1] - mapping = { - "bool": "bool", - "int8": "int8", - "int16": "int16", - "int32": "int32", - "int64": "int64", - "float16": "fp16", - "bfloat16": "bf16", - "float32": "fp32", - "float64": "fp64", - } - return mapping.get(torch_name, dtype) - return dtype - - -def ir_dtype_to_torch(dtype: str) -> torch.dtype: - normalized = normalize_dtype_name(dtype) - if normalized == "bool": - return torch.bool - if normalized == "int8": - return torch.int8 - if normalized == "int16": - return torch.int16 - if normalized == "int32": - return torch.int32 - if normalized == "int64": - return torch.int64 - if normalized == "fp16": - return torch.float16 - if normalized == "bf16": - return torch.bfloat16 - if normalized == "fp32": - return torch.float32 - if normalized == "fp64": - return torch.float64 - raise NotImplementedError(f"unsupported IR dtype for constant materialization: {dtype}") - - -def normalize_dim(dim: int, rank: int) -> int: - if dim < 0: - dim += rank - return dim - - -def numel(shape: tuple[int, ...]) -> int: - total = 1 - for dim in shape: - total *= int(dim) - return total - - -def refresh_value_users(graph: IRGraph) -> None: - for value in graph.values.values(): - value.users = [] - for node_id in graph.order: - node = graph.nodes[node_id] - for input_id in node.inputs: - value = graph.values.get(input_id) - if value is not None: - value.users.append(node_id) - - -def rebuild_graph(graph: IRGraph) -> None: - new_values: dict[str, IRValue] = {} - - for value_id in graph.inputs: - old = graph.values.get(value_id, IRValue(id=value_id)) - new_values[value_id] = IRValue( - id=value_id, - shape=old.shape, - dtype=old.dtype, - producer=None, - users=[], - meta=dict(old.meta), - ) - - for value_id in graph.constants.keys(): - old = graph.values.get(value_id, IRValue(id=value_id)) - new_values[value_id] = IRValue( - id=value_id, - shape=old.shape, - dtype=old.dtype, - producer=None, - users=[], - meta=dict(old.meta), - ) - - for node_id in graph.order: - node = graph.nodes[node_id] - for output_id in node.outputs: - old = graph.values.get(output_id, IRValue(id=output_id)) - new_values[output_id] = IRValue( - id=output_id, - shape=old.shape, - dtype=old.dtype, - producer=node_id, - users=[], - meta=dict(old.meta), - ) - - for node_id in graph.order: - node = graph.nodes[node_id] - for input_id in node.inputs: - if input_id not in new_values: - old = graph.values.get(input_id, IRValue(id=input_id)) - new_values[input_id] = IRValue( - id=input_id, - shape=old.shape, - dtype=old.dtype, - producer=old.producer, - users=[], - meta=dict(old.meta), - ) - new_values[input_id].users.append(node_id) - - for output_id in graph.outputs: - if output_id not in new_values: - old = graph.values.get(output_id, IRValue(id=output_id)) - new_values[output_id] = IRValue( - id=output_id, - shape=old.shape, - dtype=old.dtype, - producer=old.producer, - users=[], - meta=dict(old.meta), - ) - - graph.values = new_values - verify_ir(graph) - - -def dce(graph: IRGraph) -> None: - live_values = set(graph.outputs) - live_nodes: set[str] = set() - - changed = True - while changed: - changed = False - for node_id in reversed(graph.order): - node = graph.nodes.get(node_id) - if node is None: - continue - if not any(output in live_values for output in node.outputs): - continue - if node_id not in live_nodes: - live_nodes.add(node_id) - changed = True - for input_id in node.inputs: - if input_id not in live_values: - live_values.add(input_id) - changed = True - - graph.order = [node_id for node_id in graph.order if node_id in live_nodes] - graph.nodes = {node_id: node for node_id, node in graph.nodes.items() if node_id in live_nodes} - - live_value_ids = set(graph.inputs) | set(graph.outputs) - for node in graph.nodes.values(): - live_value_ids.update(node.inputs) - live_value_ids.update(node.outputs) - - graph.values = {value_id: value for value_id, value in graph.values.items() if value_id in live_value_ids} - graph.constants = {value_id: value for value_id, value in graph.constants.items() if value_id in live_value_ids} - rebuild_graph(graph) - - -def remove_node(graph: IRGraph, node_id: str) -> None: - if node_id in graph.nodes: - del graph.nodes[node_id] - graph.order = [existing_id for existing_id in graph.order if existing_id != node_id] - - -def drop_value_if_unused(graph: IRGraph, value_id: str) -> None: - if value_id in graph.inputs or value_id in graph.outputs or value_id in graph.constants: - return - value = graph.values.get(value_id) - if value is None: - return - if value.producer is None and not value.users: - del graph.values[value_id] - - -def bypass_unary_node(graph: IRGraph, node: IRNode) -> bool: - if len(node.inputs) != 1 or len(node.outputs) != 1: - return False - input_id = node.inputs[0] - output_id = node.outputs[0] - - for other in graph.nodes.values(): - other.inputs = [input_id if value_id == output_id else value_id for value_id in other.inputs] - graph.outputs = [input_id if value_id == output_id else value_id for value_id in graph.outputs] - - remove_node(graph, node.id) - if output_id in graph.values: - del graph.values[output_id] - drop_value_if_unused(graph, input_id) - return True - - -def _is_gemma4_graph(graph: IRGraph) -> bool: - family = str(graph.meta.get("adapter_family") or graph.meta.get("family") or "").lower() - profile = profile_for_family(family) - return profile is not None and profile.family == "gemma4" - - -def _weight_binding_meta(value: IRValue | None) -> dict[str, object] | None: - if value is None: - return None - path = value.meta.get("path") - kind = value.meta.get("kind") - source_name = value.meta.get("source_name") - if isinstance(path, str) and isinstance(kind, str) and isinstance(source_name, str): - return dict(value.meta) - return None - - -def _gemma4_materialized_binding_meta(graph: IRGraph, node: IRNode) -> dict[str, object] | None: - if not _is_gemma4_graph(graph) or len(node.inputs) != 1: - return None - if node.op == "precision_cast": - pass - elif node.op == "scalar_add" and float(node.attrs.get("value", 0.0)) in (0.0, 1.0): - pass - else: - return None - - value_id = node.inputs[0] - for _ in range(4): - value = graph.values.get(value_id) - binding_meta = _weight_binding_meta(value) - if binding_meta is not None: - binding_meta["materialized_from_value_id"] = value_id - binding_meta["materialized_by_op"] = node.op - return binding_meta - producer_id = value.producer if value is not None else None - producer_node = graph.nodes.get(producer_id) if producer_id is not None else None - if producer_node is None or producer_node.op not in {"precision_cast", "type_as"}: - break - if len(producer_node.inputs) != 1: - break - value_id = producer_node.inputs[0] - return None - - -def materialize_constant_output(graph: IRGraph, node: IRNode, value: Any) -> bool: - if len(node.outputs) != 1: - return False - output_id = node.outputs[0] - output_value = graph.values.get(output_id) - if output_value is None: - return False - binding_meta = _gemma4_materialized_binding_meta(graph, node) - if binding_meta is not None: - output_value.meta.update(binding_meta) - graph.constants[output_id] = value - output_value.producer = None - remove_node(graph, node.id) - return True diff --git a/python/cactus/transpile/capture_jax.py b/python/cactus/transpile/capture_jax.py deleted file mode 100644 index bef4b50e3..000000000 --- a/python/cactus/transpile/capture_jax.py +++ /dev/null @@ -1,1663 +0,0 @@ -from __future__ import annotations - -from collections.abc import Callable, Sequence -import copy -from dataclasses import dataclass -import math -from typing import Any - -import numpy as np -import torch - -from cactus.transpile.graph_ir import IRGraph -from cactus.transpile.graph_ir import IRNode -from cactus.transpile.graph_ir import IRValue -from cactus.transpile.graph_ir import verify_ir -from cactus.transpile.jax_semantic_rewrites import apply_jax_semantic_rewrites -from cactus.transpile.lower import TranspiledGraph -from cactus.transpile.lower import transpile_ir -from cactus.transpile.weight_binding import resolve_weight_binding - - -_BINARY_OPS = { - "add": "add", - "sub": "subtract", - "mul": "multiply", - "div": "divide", -} - -_BINARY_EVALUATORS = { - "add": np.add, - "sub": np.subtract, - "mul": np.multiply, - "div": np.divide, -} - -_UNARY_OPS = { - "logistic": "sigmoid", - "tanh": "tanh", - "sqrt": "scalar_sqrt", - "exp": "scalar_exp", - "log": "scalar_log", - "cos": "scalar_cos", - "sin": "scalar_sin", - "neg": "negate", - "abs": "abs", -} - -_UNARY_EVALUATORS = { - "sqrt": np.sqrt, - "exp": np.exp, - "log": np.log, - "cos": np.cos, - "sin": np.sin, - "neg": np.negative, - "logistic": lambda value: 1.0 / (1.0 + np.exp(-value)), - "tanh": np.tanh, -} - -_COMPARE_OPS = { - "eq": "equal", - "ne": "not_equal", - "lt": "less", - "le": "less_equal", - "gt": "greater", - "ge": "greater_equal", -} - -_RHS_SCALAR_COMPARE_OPS = { - "eq": "scalar_equal", - "ne": "scalar_not_equal", - "lt": "scalar_less", - "le": "scalar_less_equal", - "gt": "scalar_greater", - "ge": "scalar_greater_equal", -} - -_LHS_SCALAR_COMPARE_OPS = { - "eq": "scalar_equal", - "ne": "scalar_not_equal", - "lt": "scalar_greater", - "le": "scalar_greater_equal", - "gt": "scalar_less", - "ge": "scalar_less_equal", -} - - -def _dtype_to_ir(dtype: Any) -> str: - dtype = np.dtype(dtype) - if dtype.name == "bfloat16": - return "fp32" - if dtype == np.dtype(np.float16): - return "fp16" - if dtype == np.dtype(np.float32): - return "fp32" - if dtype == np.dtype(np.float64): - return "fp32" - if dtype == np.dtype(np.int8): - return "int8" - if dtype in { - np.dtype(np.int16), - np.dtype(np.int32), - np.dtype(np.int64), - np.dtype(np.uint8), - np.dtype(np.uint16), - np.dtype(np.uint32), - np.dtype(np.uint64), - np.dtype(np.bool_), - }: - return "fp32" - raise NotImplementedError(f"unsupported JAX dtype: {dtype}") - - -def _shape_dtype_from_aval(aval: Any) -> tuple[tuple[int, ...] | None, str | None]: - shape = getattr(aval, "shape", None) - dtype = getattr(aval, "dtype", None) - ir_shape = None if shape is None else tuple(int(dim) for dim in shape) - if ir_shape == (): - ir_shape = (1,) - ir_dtype = None if dtype is None else _dtype_to_ir(dtype) - return ir_shape, ir_dtype - - -@dataclass(frozen=True) -class _SyntheticAval: - shape: tuple[int, ...] - dtype: Any - - -def _constant_to_torch(value: Any) -> torch.Tensor: - array = np.asarray(value) - if array.dtype.name == "bfloat16": - array = array.astype(np.float32) - if array.dtype == np.dtype(np.float64): - array = array.astype(np.float32) - if np.issubdtype(array.dtype, np.floating): - compare_array = array.astype(np.float32, copy=False) - array = np.where(compare_array < -1.0e30, -65504.0, array) - array = np.where(compare_array > 1.0e30, 65504.0, array) - elif array.dtype in { - np.dtype(np.int16), - np.dtype(np.int32), - np.dtype(np.int64), - np.dtype(np.uint8), - np.dtype(np.uint16), - np.dtype(np.uint32), - np.dtype(np.uint64), - np.dtype(np.bool_), - }: - array = array.astype(np.float32) - return torch.from_numpy(np.array(array)) - - -def _tree_path_name(path: Sequence[Any]) -> str: - parts: list[str] = [] - for entry in path: - if hasattr(entry, "key"): - part = str(entry.key) - elif hasattr(entry, "name"): - part = str(entry.name) - elif hasattr(entry, "idx"): - part = str(entry.idx) - else: - part = str(entry) - parts.append(part) - return ".".join(parts) or "param" - - -def _flatten_named_leaves(tree_util: Any, params: Any) -> list[tuple[str, np.ndarray]]: - leaves: list[tuple[str, np.ndarray]] = [] - for path, value in tree_util.tree_flatten_with_path(params)[0]: - leaves.append((_tree_path_name(path), np.asarray(value))) - return leaves - - -def _weight_binding_fields(meta: dict[str, object]) -> dict[str, object] | None: - path = meta.get("path") - kind = meta.get("kind") - source_name = meta.get("source_name") - if isinstance(path, str) and isinstance(kind, str) and isinstance(source_name, str): - return {"path": path, "kind": kind, "source_name": source_name} - return None - - -def _propagate_weight_binding_meta(graph: IRGraph) -> None: - changed = True - while changed: - changed = False - for value_id in list(graph.constants): - value = graph.values[value_id] - if _weight_binding_fields(value.meta) is not None: - continue - if value.meta.get("derived_by_op") != "convert_element_type": - continue - derived_from = value.meta.get("derived_from_value_ids") - if not isinstance(derived_from, (tuple, list)): - continue - for source_id in derived_from: - source_value = graph.values.get(str(source_id)) - if source_value is None: - continue - binding_meta = _weight_binding_fields(source_value.meta) - if binding_meta is None: - continue - value.meta.update(binding_meta) - value.meta["materialized_from_value_id"] = str(source_id) - graph.meta.setdefault("weight_bindings", {})[value_id] = dict(value.meta) - changed = True - break - - -def _binding_meta( - *, - name: str, - weights_dir: str | None, - explicit: dict[str, dict[str, str]], -) -> dict[str, str]: - if name in explicit: - return dict(explicit[name]) - binding = resolve_weight_binding(weights_dir=weights_dir, source_name=name) - if binding is None: - return {} - return { - "path": binding.path, - "kind": binding.kind, - "source_name": binding.source_name, - } - - -def _freeze_leading_param_inputs( - graph: IRGraph, - named_leaves: Sequence[tuple[str, np.ndarray]], - *, - weights_dir: str | None, - explicit: dict[str, dict[str, str]], -) -> None: - if not named_leaves: - return - if len(graph.inputs) < len(named_leaves): - raise ValueError( - f"JAX graph has {len(graph.inputs)} inputs, cannot freeze {len(named_leaves)} parameter leaves" - ) - param_input_ids = list(graph.inputs[: len(named_leaves)]) - graph.inputs = graph.inputs[len(named_leaves) :] - graph.meta["weight_bindings"] = {} - for value_id, (name, leaf_array) in zip(param_input_ids, named_leaves, strict=True): - value = graph.values[value_id] - tensor = _constant_to_torch(leaf_array) - value.shape = tuple(int(dim) for dim in tensor.shape) - value.dtype = _dtype_to_ir(tensor.numpy().dtype) - value.meta.update( - { - "source_name": name, - "jax_param_input": True, - **_binding_meta(name=name, weights_dir=weights_dir, explicit=explicit), - } - ) - graph.constants[value_id] = tensor - if "path" in value.meta: - graph.meta["weight_bindings"][value_id] = dict(value.meta) - - -def _literal_value(literal: Any) -> Any: - value = getattr(literal, "val", literal) - if isinstance(value, np.ndarray): - return value.item() if value.ndim == 0 else value - if hasattr(value, "item"): - try: - return value.item() - except Exception: - return value - return value - - -class _JaxImportContext: - def __init__(self) -> None: - self.var_ids: dict[int, str] = {} - self.gather_index_aliases: dict[str, str] = {} - self.next_value_id = 0 - self.next_node_id = 0 - - def value_id(self, var: Any) -> str: - key = id(var) - existing = self.var_ids.get(key) - if existing is not None: - return existing - name = f"v{self.next_value_id}" - self.next_value_id += 1 - self.var_ids[key] = name - return name - - def bind_value(self, var: Any, value_id: str) -> None: - self.var_ids[id(var)] = value_id - - def node_id(self, op: str) -> str: - name = f"n{self.next_node_id}_{op}" - self.next_node_id += 1 - return name - - -@dataclass(frozen=True) -class JaxGraphSpec: - name: str - fn: Callable[..., Any] - example_args: Sequence[Any] - role: str = "generic" - input_names: Sequence[str] | None = None - output_names: Sequence[str] | None = None - graph_meta: dict[str, object] | None = None - - -@dataclass -class CapturedJaxGraph: - spec: JaxGraphSpec - raw_ir_graph: IRGraph - ir_graph: IRGraph - graph: TranspiledGraph - - def execute(self, *args: Any) -> list[Any]: - if len(args) != len(self.spec.example_args): - raise ValueError( - f"graph {self.spec.name!r} expected {len(self.spec.example_args)} inputs, got {len(args)}" - ) - self.graph.set_inputs( - [ - _coerce_runtime_input(arg, example) - for arg, example in zip(args, self.spec.example_args, strict=True) - ] - ) - return self.graph.execute() - - -@dataclass -class CapturedJaxGraphBundle: - graphs: dict[str, CapturedJaxGraph] - params: Any - weights_dir: str | None = None - - def execute(self, graph_name: str, *args: Any) -> list[Any]: - if graph_name not in self.graphs: - available = ", ".join(sorted(self.graphs)) or "" - raise ValueError(f"unknown JAX graph {graph_name!r}; available graphs: {available}") - return self.graphs[graph_name].execute(*args) - - -def _add_value(graph: IRGraph, value_id: str, aval: Any, *, producer: str | None = None) -> None: - shape, dtype = _shape_dtype_from_aval(aval) - graph.add_value(IRValue(id=value_id, shape=shape, dtype=dtype, producer=producer)) - - -def _coerce_runtime_input(value: Any, example: Any) -> np.ndarray: - array = np.asarray(value) - example_array = np.asarray(example) - target_dtype = example_array.dtype - if target_dtype.name == "bfloat16": - target_dtype = np.dtype(np.float32) - if array.dtype != target_dtype: - array = array.astype(target_dtype, copy=False) - if array.shape == (): - array = array.reshape((1,)) - return array - - -def _add_constant( - graph: IRGraph, - ctx: _JaxImportContext, - var: Any, - value: Any, - *, - source_name: str, - meta: dict[str, object] | None = None, -) -> str: - value_id = ctx.value_id(var) - if value_id in graph.values: - return value_id - tensor = _constant_to_torch(value) - graph.add_value( - IRValue( - id=value_id, - shape=tuple(int(dim) for dim in tensor.shape), - dtype=_dtype_to_ir(tensor.numpy().dtype), - producer=None, - meta={"source_name": source_name, **dict(meta or {})}, - ) - ) - graph.constants[value_id] = tensor - return value_id - - -def _register_node( - graph: IRGraph, - node: IRNode, - *, - out_avals: Sequence[Any], -) -> None: - graph.add_node(node) - graph.order.append(node.id) - for output_id, aval in zip(node.outputs, out_avals, strict=True): - shape, dtype = _shape_dtype_from_aval(aval) - graph.values[output_id].shape = shape - graph.values[output_id].dtype = dtype - - -def _primitive_name(eqn: Any) -> str: - primitive = getattr(eqn, "primitive", None) - return str(getattr(primitive, "name", primitive)) - - -def _is_literal(var: Any) -> bool: - return type(var).__name__ == "Literal" - - -def _ensure_literal_constant(graph: IRGraph, ctx: _JaxImportContext, literal: Any) -> str: - value_id = ctx.value_id(literal) - if value_id in graph.values: - return value_id - value = _literal_value(literal) - return _add_constant(graph, ctx, literal, value, source_name=f"literal:{value!r}") - - -def _derived_meta(input_ids: Sequence[str], *, op: str) -> dict[str, object]: - if not input_ids: - return {"derived_by_op": op} - return {"derived_from_value_ids": tuple(input_ids), "derived_by_op": op} - - -def _input_ids(graph: IRGraph, ctx: _JaxImportContext, invars: Sequence[Any]) -> list[str]: - result: list[str] = [] - for var in invars: - if _is_literal(var): - result.append(_ensure_literal_constant(graph, ctx, var)) - else: - result.append(ctx.value_id(var)) - return result - - -def _out_avals(eqn: Any) -> tuple[Any, ...]: - return tuple(getattr(var, "aval", None) for var in eqn.outvars) - - -def _literal_number(var: Any) -> float | None: - if not _is_literal(var): - return None - value = _literal_value(var) - if isinstance(value, (bool, np.bool_)): - return float(value) - if isinstance(value, (int, float, np.number)): - return float(value) - return None - - -def _constant_scalar(graph: IRGraph, value_id: str) -> float | None: - value = graph.constants.get(value_id) - if value is None: - return None - if graph.values[value_id].meta.get("jax_closed_constant"): - return None - array = np.asarray(value) - if array.shape != (): - return None - return float(array.item()) - - -def _constant_singleton_scalar(graph: IRGraph, value_id: str) -> float | None: - value = graph.constants.get(value_id) - if value is None: - return None - array = np.asarray(value) - if array.size != 1: - return None - return float(array.reshape(()).item()) - - -def _constant_scalar_or_singleton(graph: IRGraph, value_id: str) -> float | None: - scalar = _constant_scalar(graph, value_id) - if scalar is not None: - return scalar - singleton = _constant_singleton_scalar(graph, value_id) - if singleton is not None: - return singleton - value = graph.constants.get(value_id) - if value is None: - return None - array = np.asarray(value) - if array.size == 0: - return None - first = float(array.reshape(-1)[0].item()) - if np.allclose(array, first, rtol=0.0, atol=0.0): - return first - return None - - -def _replace_scalar_constant_with_broadcast( - graph: IRGraph, - value_id: str, - output_var: Any, - *, - source_op: str, -) -> str: - scalar = _constant_singleton_scalar(graph, value_id) - if scalar is None: - return value_id - shape = tuple(int(dim) for dim in getattr(output_var.aval, "shape", ())) - if not shape or _product(shape) <= 1: - return value_id - dtype = getattr(output_var.aval, "dtype", np.float32) - replacement_id = f"{value_id}__{source_op}_broadcast" - if replacement_id in graph.values: - return replacement_id - tensor = _constant_to_torch(np.full(shape, scalar, dtype=np.dtype(dtype))) - graph.add_value( - IRValue( - id=replacement_id, - shape=shape, - dtype=_dtype_to_ir(tensor.numpy().dtype), - producer=None, - meta={ - "source_name": f"broadcasted_scalar:{source_op}", - **_derived_meta((value_id,), op=f"{source_op}:scalar_broadcast"), - }, - ) - ) - graph.constants[replacement_id] = tensor - return replacement_id - - -def _where_scalar_value(value: float) -> float: - if value < -1.0e30: - return -65504.0 - if value > 1.0e30: - return 65504.0 - return value - - -def _constant_array(graph: IRGraph, value_id: str) -> np.ndarray | None: - value = graph.constants.get(value_id) - if value is None: - return None - if graph.values[value_id].meta.get("jax_closed_constant"): - return None - return np.asarray(value) - - -def _try_fold_constant_unary( - graph: IRGraph, - ctx: _JaxImportContext, - *, - prim: str, - input_id: str, - outvar: Any, -) -> bool: - evaluator = _UNARY_EVALUATORS.get(prim) - array = _constant_array(graph, input_id) - if evaluator is None or array is None: - return False - _add_constant( - graph, - ctx, - outvar, - evaluator(array), - source_name=f"folded:{prim}", - meta=_derived_meta((input_id,), op=prim), - ) - return True - - -def _try_fold_constant_binary( - graph: IRGraph, - ctx: _JaxImportContext, - *, - prim: str, - input_ids: Sequence[str], - outvar: Any, -) -> bool: - evaluator = _BINARY_EVALUATORS.get(prim) - if evaluator is None: - return False - lhs_array = _constant_array(graph, input_ids[0]) - rhs_array = _constant_array(graph, input_ids[1]) - if lhs_array is None or rhs_array is None: - return False - _add_constant( - graph, - ctx, - outvar, - evaluator(lhs_array, rhs_array), - source_name=f"folded:{prim}", - meta=_derived_meta(input_ids, op=prim), - ) - return True - - -def _alias_output(ctx: _JaxImportContext, outvar: Any, source_id: str) -> None: - ctx.bind_value(outvar, source_id) - - -def _product(values: Sequence[int]) -> int: - result = 1 - for value in values: - result *= int(value) - return result - - -def _conv_static_int(values: Sequence[int], *, name: str) -> int: - ints = tuple(int(value) for value in values) - if not ints: - return 1 - if any(value != ints[0] for value in ints): - raise NotImplementedError(f"JAX conv_general_dilated requires uniform {name}, got {ints}") - return ints[0] - - -def _conv_padding_int(padding: Sequence[Sequence[int]]) -> int: - pairs = tuple((int(left), int(right)) for left, right in padding) - if not pairs: - return 0 - if any(left != right for left, right in pairs): - raise NotImplementedError(f"JAX conv_general_dilated requires symmetric padding, got {pairs}") - return _conv_static_int(tuple(left for left, _ in pairs), name="padding") - - -def _conv_output_permutation(out_spec: Sequence[int]) -> tuple[int, ...]: - spec = tuple(int(axis) for axis in out_spec) - return tuple(spec.index(final_axis) for final_axis in range(len(spec))) - - -def _jaxpr_invars(jaxpr_like: Any) -> Sequence[Any]: - return tuple(getattr(getattr(jaxpr_like, "jaxpr", jaxpr_like), "invars", ())) - - -def _jaxpr_outvars(jaxpr_like: Any) -> Sequence[Any]: - return tuple(getattr(getattr(jaxpr_like, "jaxpr", jaxpr_like), "outvars", ())) - - -def _jaxpr_eqns(jaxpr_like: Any) -> Sequence[Any]: - return tuple(getattr(getattr(jaxpr_like, "jaxpr", jaxpr_like), "eqns", ())) - - -def _jaxpr_constvars(jaxpr_like: Any) -> Sequence[Any]: - return tuple(getattr(getattr(jaxpr_like, "jaxpr", jaxpr_like), "constvars", ())) - - -def _jaxpr_consts(jaxpr_like: Any) -> Sequence[Any]: - return tuple(getattr(jaxpr_like, "consts", ()) or ()) - - -def _inline_jaxpr( - graph: IRGraph, - ctx: _JaxImportContext, - jaxpr_like: Any, - input_ids: Sequence[str], - outvars: Sequence[Any], -) -> None: - inner_output_ids = _inline_jaxpr_outputs(graph, ctx, jaxpr_like, input_ids) - if len(outvars) != len(inner_output_ids): - raise NotImplementedError("nested JAXPR output arity mismatch") - for outer_var, inner_output_id in zip(outvars, inner_output_ids, strict=True): - _alias_output(ctx, outer_var, inner_output_id) - - -def _inline_jaxpr_outputs( - graph: IRGraph, - ctx: _JaxImportContext, - jaxpr_like: Any, - input_ids: Sequence[str], -) -> list[str]: - invars = _jaxpr_invars(jaxpr_like) - if len(input_ids) != len(invars): - raise NotImplementedError("nested JAXPR input arity mismatch") - invar_keys = {id(var) for var in invars} - for inner_eqn in _jaxpr_eqns(jaxpr_like): - for outvar in inner_eqn.outvars: - if id(outvar) not in invar_keys: - ctx.var_ids.pop(id(outvar), None) - for inner_var, input_id in zip(invars, input_ids, strict=True): - ctx.bind_value(inner_var, input_id) - for index, (constvar, const) in enumerate(zip(_jaxpr_constvars(jaxpr_like), _jaxpr_consts(jaxpr_like), strict=True)): - _add_constant(graph, ctx, constvar, const, source_name=f"nested_const_{index}") - for inner_eqn in _jaxpr_eqns(jaxpr_like): - _import_eqn(graph, ctx, inner_eqn) - return [ctx.value_id(inner_var) for inner_var in _jaxpr_outvars(jaxpr_like)] - - -def _generated_value( - graph: IRGraph, - ctx: _JaxImportContext, - *, - stem: str, - shape: tuple[int, ...] | None, - dtype: str | None, - producer: str, - meta: dict[str, object] | None = None, -) -> str: - value_id = f"v{ctx.next_value_id}_{stem}" - ctx.next_value_id += 1 - if value_id in graph.values: - raise ValueError(f"duplicate generated IR value id: {value_id}") - graph.values[value_id] = IRValue( - id=value_id, - shape=shape, - dtype=dtype, - producer=producer, - meta=dict(meta or {}), - ) - return value_id - - -def _register_generated_node( - graph: IRGraph, - node: IRNode, - *, - output_shapes: Sequence[tuple[int, ...] | None], - output_dtypes: Sequence[str | None], -) -> None: - if node.id in graph.nodes: - raise ValueError(f"duplicate IR node id: {node.id}") - graph.nodes[node.id] = node - graph.order.append(node.id) - for output_id, shape, dtype in zip(node.outputs, output_shapes, output_dtypes, strict=True): - graph.values[output_id].shape = shape - graph.values[output_id].dtype = dtype - - -def _scan_slice_input( - graph: IRGraph, - ctx: _JaxImportContext, - value_id: str, - *, - iteration: int, - input_index: int, -) -> str: - source_value = graph.values[value_id] - if source_value.shape is None or not source_value.shape: - return value_id - slice_node_id = ctx.node_id("scan_slice") - sliced_shape = (1, *tuple(int(dim) for dim in source_value.shape[1:])) - sliced_id = _generated_value( - graph, - ctx, - stem=f"scan{iteration}_{input_index}", - shape=sliced_shape, - dtype=source_value.dtype, - producer=slice_node_id, - meta={ - "derived_from_value_ids": (value_id,), - "derived_by_op": "scan_slice", - }, - ) - slice_node = IRNode( - slice_node_id, - "slice", - [value_id], - [sliced_id], - attrs={"axis": 0, "start": iteration, "end": iteration + 1, "step": 1}, - meta={"jax_generated": "scan_unroll"}, - ) - _register_generated_node(graph, slice_node, output_shapes=(sliced_shape,), output_dtypes=(source_value.dtype,)) - view_node_id = ctx.node_id("scan_slice_view") - output_shape = tuple(int(dim) for dim in source_value.shape[1:]) - if output_shape == () and sliced_shape == (1,): - return sliced_id - output_id = _generated_value( - graph, - ctx, - stem=f"scan{iteration}_{input_index}_view", - shape=output_shape, - dtype=source_value.dtype, - producer=view_node_id, - meta={ - "derived_from_value_ids": (sliced_id,), - "derived_by_op": "scan_slice_view", - }, - ) - view_node = IRNode( - view_node_id, - "view", - [sliced_id], - [output_id], - attrs={"shape": output_shape}, - meta={"jax_generated": "scan_unroll"}, - ) - _register_generated_node(graph, view_node, output_shapes=(output_shape,), output_dtypes=(source_value.dtype,)) - return output_id - - -def _add_iota_constant(graph: IRGraph, ctx: _JaxImportContext, eqn: Any) -> None: - params = dict(getattr(eqn, "params", {}) or {}) - shape = tuple(int(dim) for dim in params["shape"]) - dimension = int(params["dimension"]) - dtype = params.get("dtype", getattr(eqn.outvars[0].aval, "dtype", np.float32)) - values = np.arange(shape[dimension], dtype=np.dtype(dtype)) - reshape = [1] * len(shape) - reshape[dimension] = shape[dimension] - values = np.broadcast_to(values.reshape(reshape), shape) - _add_constant(graph, ctx, eqn.outvars[0], values, source_name=f"iota:{shape}:{dimension}") - - -def _import_eqn(graph: IRGraph, ctx: _JaxImportContext, eqn: Any) -> None: - prim = _primitive_name(eqn) - params = dict(getattr(eqn, "params", {}) or {}) - - if prim in {"jit", "remat2"}: - _inline_jaxpr(graph, ctx, params["jaxpr"], _input_ids(graph, ctx, eqn.invars), eqn.outvars) - return - - if prim in {"custom_jvp_call", "custom_jvp_call_jaxpr"}: - call_jaxpr = params.get("call_jaxpr") or params.get("fun_jaxpr") or params.get("jaxpr") - if call_jaxpr is None: - raise NotImplementedError("JAX custom_jvp_call import requires a primal call_jaxpr") - _inline_jaxpr(graph, ctx, call_jaxpr, _input_ids(graph, ctx, eqn.invars), eqn.outvars) - return - - if prim == "scan": - length = int(params.get("length", 1)) - num_consts = int(params.get("num_consts", 0)) - num_carry = int(params.get("num_carry", len(eqn.outvars))) - input_ids = _input_ids(graph, ctx, eqn.invars) - const_ids = input_ids[:num_consts] - carry_ids = input_ids[num_consts : num_consts + num_carry] - xs_ids = input_ids[num_consts + num_carry :] - if len(eqn.outvars) != num_carry: - raise NotImplementedError("JAX scan import currently supports carry-only scan outputs") - for iteration in range(length): - sliced_xs = [ - _scan_slice_input(graph, ctx, value_id, iteration=iteration, input_index=index) - for index, value_id in enumerate(xs_ids) - ] - body_outputs = _inline_jaxpr_outputs(graph, ctx, params["jaxpr"], [*const_ids, *carry_ids, *sliced_xs]) - if len(body_outputs) != num_carry: - raise NotImplementedError("JAX scan body output arity mismatch") - carry_ids = body_outputs - for outvar, carry_id in zip(eqn.outvars, carry_ids, strict=True): - _alias_output(ctx, outvar, carry_id) - return - - if prim == "stop_gradient": - _alias_output(ctx, eqn.outvars[0], _input_ids(graph, ctx, eqn.invars)[0]) - return - - if prim == "iota": - _add_iota_constant(graph, ctx, eqn) - return - - node = _node_for_eqn(graph, ctx, eqn) - if node is None: - return - _register_node(graph, node, out_avals=_out_avals(eqn)) - - -def _node_for_eqn(graph: IRGraph, ctx: _JaxImportContext, eqn: Any) -> IRNode | None: - prim = _primitive_name(eqn) - inputs = _input_ids(graph, ctx, eqn.invars) - outputs = [ctx.value_id(var) for var in eqn.outvars] - node_id = ctx.node_id(prim) - params = dict(getattr(eqn, "params", {}) or {}) - meta = {"jax_primitive": prim} - - if prim in _BINARY_OPS: - if _try_fold_constant_binary(graph, ctx, prim=prim, input_ids=inputs, outvar=eqn.outvars[0]): - return None - if prim == "div": - lhs_literal = _literal_number(eqn.invars[0]) - if lhs_literal is not None: - reciprocal_id = f"{outputs[0]}__reciprocal" - reciprocal_node = IRNode( - node_id, - "pow", - [inputs[1]], - [reciprocal_id], - attrs={"exponent": -1.0}, - meta=meta, - ) - _register_node(graph, reciprocal_node, out_avals=(eqn.outvars[0].aval,)) - if lhs_literal == 1.0: - _alias_output(ctx, eqn.outvars[0], reciprocal_id) - return None - return IRNode( - ctx.node_id("scalar_div_mul"), - "scalar_multiply", - [reciprocal_id], - outputs, - attrs={"value": lhs_literal}, - meta=meta, - ) - if prim in _BINARY_OPS: - if len(inputs) == 2 and len(outputs) == 1: - inputs = [ - _replace_scalar_constant_with_broadcast(graph, input_id, eqn.outvars[0], source_op=prim) - for input_id in inputs - ] - return IRNode(node_id, _BINARY_OPS[prim], inputs, outputs, meta=meta) - if prim in _UNARY_OPS: - if _try_fold_constant_unary(graph, ctx, prim=prim, input_id=inputs[0], outvar=eqn.outvars[0]): - return None - return IRNode(node_id, _UNARY_OPS[prim], inputs, outputs, meta=meta) - if prim == "log1p": - add_id = f"{outputs[0]}__log1p_add" - _register_node( - graph, - IRNode( - node_id, - "scalar_add", - inputs, - [add_id], - attrs={"value": 1.0}, - meta=meta, - ), - out_avals=(eqn.outvars[0].aval,), - ) - return IRNode(ctx.node_id("log1p_log"), "scalar_log", [add_id], outputs, meta=meta) - if prim == "expm1": - exp_id = f"{outputs[0]}__expm1_exp" - _register_node( - graph, - IRNode(node_id, "scalar_exp", inputs, [exp_id], meta=meta), - out_avals=(eqn.outvars[0].aval,), - ) - return IRNode(ctx.node_id("expm1_sub"), "scalar_add", [exp_id], outputs, attrs={"value": -1.0}, meta=meta) - if prim in _COMPARE_OPS: - lhs_scalar = _constant_scalar(graph, inputs[0]) - rhs_scalar = _constant_scalar(graph, inputs[1]) - if rhs_scalar is not None: - return IRNode(node_id, _RHS_SCALAR_COMPARE_OPS[prim], [inputs[0]], outputs, attrs={"value": rhs_scalar}, meta=meta) - if lhs_scalar is not None: - return IRNode(node_id, _LHS_SCALAR_COMPARE_OPS[prim], [inputs[1]], outputs, attrs={"value": lhs_scalar}, meta=meta) - return IRNode(node_id, _COMPARE_OPS[prim], inputs, outputs, meta=meta) - if prim == "and": - return IRNode(node_id, "logical_and", inputs, outputs, meta=meta) - if prim == "or": - return IRNode(node_id, "logical_or", inputs, outputs, meta=meta) - if prim == "square": - return IRNode(node_id, "multiply", [inputs[0], inputs[0]], outputs, meta=meta) - if prim == "rsqrt": - return IRNode(node_id, "pow", inputs, outputs, attrs={"exponent": -0.5}, meta=meta) - if prim == "integer_pow": - exponent = int(params["y"]) - if exponent == 2: - return IRNode(node_id, "multiply", [inputs[0], inputs[0]], outputs, meta=meta) - return IRNode(node_id, "pow", inputs, outputs, attrs={"exponent": float(exponent)}, meta=meta) - if prim == "pow": - rhs_literal = _literal_number(eqn.invars[1]) - if rhs_literal is not None: - return IRNode(node_id, "pow", [inputs[0]], outputs, attrs={"exponent": float(rhs_literal)}, meta=meta) - lhs_literal = _literal_number(eqn.invars[0]) - if lhs_literal is not None and lhs_literal > 0.0: - intermediate = f"{outputs[0]}__logmul" - scale_node = IRNode( - node_id, - "scalar_multiply", - [inputs[1]], - [intermediate], - attrs={"value": math.log(lhs_literal)}, - meta=meta, - ) - _register_node(graph, scale_node, out_avals=(eqn.outvars[0].aval,)) - return IRNode(ctx.node_id("pow_exp"), "scalar_exp", [intermediate], outputs, meta=meta) - raise NotImplementedError("JAX pow import only supports positive scalar base") - if prim == "dot_general": - dimension_numbers = params.get("dimension_numbers") - if dimension_numbers is not None: - ((lhs_contract, rhs_contract), (lhs_batch, rhs_batch)) = dimension_numbers - lhs_shape = tuple(getattr(eqn.invars[0].aval, "shape", ())) - rhs_shape = tuple(getattr(eqn.invars[1].aval, "shape", ())) - lhs_contract = tuple(int(dim) for dim in lhs_contract) - rhs_contract = tuple(int(dim) for dim in rhs_contract) - lhs_batch = tuple(int(dim) for dim in lhs_batch) - rhs_batch = tuple(int(dim) for dim in rhs_batch) - if lhs_batch or rhs_batch: - if tuple(lhs_shape[dim] for dim in lhs_batch) != tuple(rhs_shape[dim] for dim in rhs_batch): - raise NotImplementedError("JAX dot_general batch dimensions must have matching sizes") - expected_rhs_contract = (len(rhs_batch),) - if ( - tuple(lhs_contract) != (len(lhs_shape) - 1,) - or tuple(rhs_contract) != expected_rhs_contract - or len(rhs_shape) > 2 - ): - batch_shape = tuple(lhs_shape[dim] for dim in lhs_batch) - lhs_contract_shape = tuple(lhs_shape[dim] for dim in lhs_contract) - rhs_contract_shape = tuple(rhs_shape[dim] for dim in rhs_contract) - if lhs_contract_shape != rhs_contract_shape: - raise NotImplementedError("JAX dot_general contraction dimensions must have matching sizes") - lhs_non_contract = tuple( - dim for dim in range(len(lhs_shape)) if dim not in set(lhs_batch) | set(lhs_contract) - ) - rhs_non_contract = tuple( - dim for dim in range(len(rhs_shape)) if dim not in set(rhs_batch) | set(rhs_contract) - ) - lhs_order = lhs_batch + lhs_non_contract + lhs_contract - rhs_order = rhs_batch + rhs_contract + rhs_non_contract - lhs_non_shape = tuple(lhs_shape[dim] for dim in lhs_non_contract) - rhs_non_shape = tuple(rhs_shape[dim] for dim in rhs_non_contract) - lhs_matrix_shape = batch_shape + (_product(lhs_non_shape), _product(lhs_contract_shape)) - rhs_matrix_shape = batch_shape + (_product(rhs_contract_shape), _product(rhs_non_shape)) - output_shape = batch_shape + lhs_non_shape + rhs_non_shape - dtype = getattr(eqn.outvars[0].aval, "dtype", getattr(eqn.invars[0].aval, "dtype", np.float32)) - - lhs_id = inputs[0] - if lhs_order != tuple(range(len(lhs_shape))): - lhs_id = f"{outputs[0]}__lhs_transpose" - _register_node( - graph, - IRNode( - ctx.node_id("dot_general_lhs_transpose"), - "permute", - [inputs[0]], - [lhs_id], - attrs={"permutation": lhs_order}, - meta=meta, - ), - out_avals=(_SyntheticAval(tuple(lhs_shape[dim] for dim in lhs_order), dtype),), - ) - lhs_2d = f"{outputs[0]}__lhs_reshape" - _register_node( - graph, - IRNode( - ctx.node_id("dot_general_lhs_reshape"), - "reshape", - [lhs_id], - [lhs_2d], - attrs={"shape": lhs_matrix_shape}, - meta=meta, - ), - out_avals=(_SyntheticAval(lhs_matrix_shape, dtype),), - ) - - rhs_id = inputs[1] - if rhs_order != tuple(range(len(rhs_shape))): - rhs_id = f"{outputs[0]}__rhs_transpose" - _register_node( - graph, - IRNode( - ctx.node_id("dot_general_rhs_transpose"), - "permute", - [inputs[1]], - [rhs_id], - attrs={"permutation": rhs_order}, - meta=meta, - ), - out_avals=(_SyntheticAval(tuple(rhs_shape[dim] for dim in rhs_order), dtype),), - ) - rhs_2d = f"{outputs[0]}__rhs_reshape" - _register_node( - graph, - IRNode( - ctx.node_id("dot_general_rhs_reshape"), - "reshape", - [rhs_id], - [rhs_2d], - attrs={"shape": rhs_matrix_shape}, - meta=meta, - ), - out_avals=(_SyntheticAval(rhs_matrix_shape, dtype),), - ) - - matmul_id = f"{outputs[0]}__matmul" - _register_node( - graph, - IRNode(ctx.node_id("dot_general_matmul"), "matmul", [lhs_2d, rhs_2d], [matmul_id], meta=meta), - out_avals=(_SyntheticAval(batch_shape + (lhs_matrix_shape[-2], rhs_matrix_shape[-1]), dtype),), - ) - return IRNode( - ctx.node_id("dot_general_output_reshape"), - "reshape", - [matmul_id], - outputs, - attrs={"shape": output_shape}, - meta=meta, - ) - return IRNode(node_id, "matmul", inputs, outputs, meta=meta) - if prim == "conv_general_dilated": - if len(inputs) != 2 or len(outputs) != 1: - raise NotImplementedError("JAX conv_general_dilated import supports input and kernel only") - if int(params.get("batch_group_count", 1)) != 1: - raise NotImplementedError("JAX conv_general_dilated batch_group_count is unsupported") - lhs_dilation = tuple(int(value) for value in params.get("lhs_dilation", ())) - if any(value != 1 for value in lhs_dilation): - raise NotImplementedError(f"JAX conv_general_dilated lhs_dilation is unsupported: {lhs_dilation}") - dimension_numbers = params.get("dimension_numbers") - if dimension_numbers is None: - raise NotImplementedError("JAX conv_general_dilated requires static dimension_numbers") - - lhs_spec = tuple(int(axis) for axis in dimension_numbers.lhs_spec) - rhs_spec = tuple(int(axis) for axis in dimension_numbers.rhs_spec) - out_spec = tuple(int(axis) for axis in dimension_numbers.out_spec) - lhs_shape = tuple(int(dim) for dim in getattr(eqn.invars[0].aval, "shape", ())) - rhs_shape = tuple(int(dim) for dim in getattr(eqn.invars[1].aval, "shape", ())) - out_shape = tuple(int(dim) for dim in getattr(eqn.outvars[0].aval, "shape", ())) - rank = len(lhs_shape) - if rank not in {3, 4}: - raise NotImplementedError(f"JAX conv_general_dilated supports 1D/2D convs, got rank {rank}") - if tuple(sorted(lhs_spec)) != tuple(range(rank)) or tuple(sorted(rhs_spec)) != tuple(range(rank)): - raise NotImplementedError("JAX conv_general_dilated requires complete dimension specs") - if tuple(sorted(out_spec)) != tuple(range(rank)): - raise NotImplementedError("JAX conv_general_dilated requires complete output dimension spec") - - stride = _conv_static_int(params.get("window_strides", (1,) * (rank - 2)), name="stride") - dilation = _conv_static_int(params.get("rhs_dilation", (1,) * (rank - 2)), name="rhs_dilation") - padding = _conv_padding_int(params.get("padding", ((0, 0),) * (rank - 2))) - groups = int(params.get("feature_group_count", 1)) - dtype = getattr(eqn.outvars[0].aval, "dtype", getattr(eqn.invars[0].aval, "dtype", np.float32)) - - x_id = inputs[0] - x_shape = tuple(lhs_shape[axis] for axis in lhs_spec) - if lhs_spec != tuple(range(rank)): - x_id = f"{outputs[0]}__conv_lhs" - _register_node( - graph, - IRNode( - ctx.node_id("conv_lhs_transpose"), - "permute", - [inputs[0]], - [x_id], - attrs={"permutation": lhs_spec}, - meta=meta, - ), - out_avals=(_SyntheticAval(x_shape, dtype),), - ) - - weight_id = inputs[1] - weight_shape = tuple(rhs_shape[axis] for axis in rhs_spec) - if rhs_spec != tuple(range(rank)): - weight_id = f"{outputs[0]}__conv_rhs" - _register_node( - graph, - IRNode( - ctx.node_id("conv_rhs_transpose"), - "permute", - [inputs[1]], - [weight_id], - attrs={"permutation": rhs_spec}, - meta=meta, - ), - out_avals=(_SyntheticAval(weight_shape, dtype),), - ) - - conv_id = outputs[0] - output_permutation = _conv_output_permutation(out_spec) - conv_shape = tuple(out_shape[axis] for axis in out_spec) - if output_permutation != tuple(range(rank)): - conv_id = f"{outputs[0]}__conv" - _register_node( - graph, - IRNode( - node_id, - "conv1d" if rank == 3 else "conv2d", - [x_id, weight_id], - [conv_id], - attrs={"stride": stride, "padding": padding, "dilation": dilation, "groups": groups}, - meta=meta, - ), - out_avals=(_SyntheticAval(conv_shape, dtype),), - ) - if conv_id == outputs[0]: - return None - return IRNode( - ctx.node_id("conv_output_transpose"), - "permute", - [conv_id], - outputs, - attrs={"permutation": output_permutation}, - meta=meta, - ) - if prim == "reshape": - if inputs[0] in ctx.gather_index_aliases: - ctx.gather_index_aliases[outputs[0]] = ctx.gather_index_aliases[inputs[0]] - new_shape = tuple(int(dim) for dim in params["new_sizes"]) - input_shape = tuple(int(dim) for dim in getattr(eqn.invars[0].aval, "shape", ())) - source_value = graph.values.get(inputs[0]) - graph_input_shape = tuple(int(dim) for dim in (source_value.shape or ())) if source_value is not None else input_shape - if new_shape == () and graph_input_shape == (1,): - _alias_output(ctx, eqn.outvars[0], inputs[0]) - return None - return IRNode( - node_id, - "reshape", - inputs, - outputs, - attrs={"shape": new_shape}, - meta=meta, - ) - if prim == "squeeze": - shape = tuple(int(dim) for dim in getattr(eqn.outvars[0].aval, "shape", ())) - return IRNode(node_id, "reshape", inputs, outputs, attrs={"shape": shape}, meta=meta) - if prim == "transpose": - permutation = params.get("permutation") - if permutation is None: - permutation = params.get("permutation_or_none") - if permutation is None: - raise NotImplementedError("JAX transpose missing permutation") - return IRNode( - node_id, - "permute", - inputs, - outputs, - attrs={"permutation": tuple(int(dim) for dim in permutation)}, - meta=meta, - ) - if prim == "convert_element_type": - array = _constant_array(graph, inputs[0]) - if array is not None: - _add_constant( - graph, - ctx, - eqn.outvars[0], - array.astype(np.dtype(params["new_dtype"])), - source_name="folded:convert_element_type", - meta=_derived_meta(inputs, op=prim), - ) - return None - return IRNode( - node_id, - "precision_cast", - inputs, - outputs, - attrs={"dtype": _dtype_to_ir(params["new_dtype"])}, - meta=meta, - ) - if prim == "broadcast_in_dim": - shape = tuple(int(dim) for dim in params["shape"]) - scalar = _constant_singleton_scalar(graph, inputs[0]) - if scalar is not None: - _add_constant( - graph, - ctx, - eqn.outvars[0], - np.full(shape, scalar, dtype=np.float32), - source_name="folded:broadcast_in_dim", - meta=_derived_meta(inputs, op=prim), - ) - return None - if inputs[0] in ctx.gather_index_aliases: - ctx.gather_index_aliases[outputs[0]] = ctx.gather_index_aliases[inputs[0]] - broadcast_dimensions = tuple(int(dim) for dim in params.get("broadcast_dimensions", ())) - input_shape = tuple(graph.values[inputs[0]].shape or ()) - if broadcast_dimensions and broadcast_dimensions != tuple(range(len(input_shape))): - reshape_shape = [1] * len(shape) - for input_axis, output_axis in enumerate(broadcast_dimensions): - reshape_shape[output_axis] = input_shape[input_axis] - reshaped_shape = tuple(reshape_shape) - if reshaped_shape == shape: - return IRNode(node_id, "view", inputs, outputs, attrs={"shape": reshaped_shape}, meta=meta) - reshaped_id = _generated_value( - graph, - ctx, - stem="broadcast_base", - shape=reshaped_shape, - dtype=graph.values[inputs[0]].dtype, - producer=node_id, - meta=_derived_meta(inputs, op=f"{prim}:reshape"), - ) - reshape_node = IRNode(node_id, "view", inputs, [reshaped_id], attrs={"shape": reshaped_shape}, meta=meta) - _register_generated_node( - graph, - reshape_node, - output_shapes=(reshaped_shape,), - output_dtypes=(graph.values[inputs[0]].dtype,), - ) - node_id = ctx.node_id(prim) - return IRNode(node_id, "expand", [reshaped_id], outputs, attrs={"shape": shape}, meta=meta) - return IRNode(node_id, "expand", inputs, outputs, attrs={"shape": shape}, meta=meta) - if prim == "concatenate": - return IRNode(node_id, "cat", inputs, outputs, attrs={"axis": int(params["dimension"])}, meta=meta) - if prim == "stack": - axis = int(params.get("axis", 0)) - input_shape = tuple(getattr(eqn.invars[0].aval, "shape", ())) - rank = len(input_shape) + 1 - if axis < 0: - axis += rank - if axis < 0 or axis > len(input_shape): - raise NotImplementedError(f"JAX stack axis out of range: {axis}") - dtype = getattr(eqn.outvars[0].aval, "dtype", getattr(eqn.invars[0].aval, "dtype", np.float32)) - expanded_inputs: list[str] = [] - for input_index, input_id in enumerate(inputs): - expanded_id = f"{outputs[0]}__stack_{input_index}" - expanded_shape = input_shape[:axis] + (1,) + input_shape[axis:] - _register_node( - graph, - IRNode( - ctx.node_id("stack_view"), - "view", - [input_id], - [expanded_id], - attrs={"shape": expanded_shape}, - meta=meta, - ), - out_avals=(_SyntheticAval(expanded_shape, dtype),), - ) - expanded_inputs.append(expanded_id) - return IRNode(node_id, "cat", expanded_inputs, outputs, attrs={"axis": axis}, meta=meta) - if prim == "tile": - reps = tuple(int(value) for value in params.get("reps", ())) - if any(rep <= 0 for rep in reps): - raise NotImplementedError(f"JAX tile import requires positive reps, got {reps}") - input_shape = tuple(int(dim) for dim in getattr(eqn.invars[0].aval, "shape", ())) - if len(reps) < len(input_shape): - reps = (1,) * (len(input_shape) - len(reps)) + reps - current_id = inputs[0] - current_shape = input_shape - dtype = getattr(eqn.invars[0].aval, "dtype", getattr(eqn.outvars[0].aval, "dtype", np.float32)) - if len(reps) > len(current_shape): - current_shape = (1,) * (len(reps) - len(current_shape)) + current_shape - viewed_id = f"{outputs[0]}__tile_view" - _register_node( - graph, - IRNode( - ctx.node_id("tile_view"), - "view", - [current_id], - [viewed_id], - attrs={"shape": current_shape}, - meta=meta, - ), - out_avals=(_SyntheticAval(current_shape, dtype),), - ) - current_id = viewed_id - for axis, rep in enumerate(reps): - if rep == 1: - continue - next_id = outputs[0] if axis == len(reps) - 1 and all(factor == 1 for factor in reps[axis + 1 :]) else f"{outputs[0]}__tile_axis_{axis}" - next_shape = current_shape[:axis] + (current_shape[axis] * rep,) + current_shape[axis + 1 :] - _register_node( - graph, - IRNode( - ctx.node_id("tile_cat"), - "cat", - [current_id] * rep, - [next_id], - attrs={"axis": axis}, - meta=meta, - ), - out_avals=(_SyntheticAval(next_shape, dtype),), - ) - current_id = next_id - current_shape = next_shape - if current_id != outputs[0]: - _alias_output(ctx, eqn.outvars[0], current_id) - return None - if prim == "split": - axis = int(params.get("axis", 0)) - start = 0 - for output_id, output_var, size in zip(outputs, eqn.outvars, params.get("sizes", ()), strict=True): - end = start + int(size) - _register_node( - graph, - IRNode( - ctx.node_id("split_slice"), - "slice", - inputs, - [output_id], - attrs={"axis": axis, "start": start, "end": end, "step": 1}, - meta=meta, - ), - out_avals=(output_var.aval,), - ) - start = end - return None - if prim == "slice": - starts = tuple(int(value) for value in params["start_indices"]) - limits = tuple(int(value) for value in params["limit_indices"]) - raw_strides = params.get("strides") - strides = (1,) * len(starts) if raw_strides is None else tuple(int(value) for value in raw_strides) - input_shape = getattr(eqn.invars[0].aval, "shape", ()) - changed_axes = [axis for axis, (start, limit, stride) in enumerate(zip(starts, limits, strides, strict=True)) if start != 0 or limit != input_shape[axis] or stride != 1] - if not changed_axes: - _alias_output(ctx, eqn.outvars[0], inputs[0]) - return None - if len(changed_axes) != 1: - raise NotImplementedError("JAX slice import supports one non-trivial axis") - axis = changed_axes[0] - return IRNode( - node_id, - "slice", - inputs, - outputs, - attrs={"axis": axis, "start": starts[axis], "end": limits[axis], "step": strides[axis]}, - meta=meta, - ) - if prim == "select_n": - if len(inputs) != 3: - raise NotImplementedError("JAX select_n import supports ternary select") - ctx.gather_index_aliases[outputs[0]] = inputs[1] - false_scalar = _constant_scalar(graph, inputs[1]) - true_scalar = _constant_scalar(graph, inputs[2]) - where_inputs = [inputs[0]] - attrs: dict[str, object] = {} - if true_scalar is None: - where_inputs.append(inputs[2]) - else: - attrs["true_is_scalar"] = True - attrs["true_value"] = _where_scalar_value(true_scalar) - if false_scalar is None: - where_inputs.append(inputs[1]) - else: - attrs["false_is_scalar"] = True - attrs["false_value"] = _where_scalar_value(false_scalar) - return IRNode(node_id, "where", where_inputs, outputs, attrs=attrs, meta=meta) - if prim == "gather": - if len(inputs) != 2: - raise NotImplementedError("JAX gather import supports embedding-style gather") - indices_id = ctx.gather_index_aliases.get(inputs[1], inputs[1]) - index_shape = tuple(getattr(eqn.invars[1].aval, "shape", ())) - if index_shape and index_shape[-1] == 1: - squeezed_id = f"{indices_id}__squeezed" - squeeze_shape = tuple(int(dim) for dim in index_shape[:-1]) - if squeezed_id not in graph.values: - squeeze_node = IRNode( - node_id, - "reshape", - [indices_id], - [squeezed_id], - attrs={"shape": squeeze_shape}, - meta=meta, - ) - _register_node(graph, squeeze_node, out_avals=(eqn.invars[1].aval,)) - graph.values[squeezed_id].shape = squeeze_shape - if indices_id in ctx.gather_index_aliases: - ctx.gather_index_aliases[squeezed_id] = ctx.gather_index_aliases[indices_id] - indices_id = squeezed_id - node_id = ctx.node_id("gather_embedding") - indices_id = ctx.gather_index_aliases.get(indices_id, indices_id) - return IRNode(node_id, "embedding", [inputs[0], indices_id], outputs, meta=meta) - if prim in {"max", "min"}: - lhs_literal = _literal_number(eqn.invars[0]) - rhs_literal = _literal_number(eqn.invars[1]) - identity = -math.inf if prim == "max" else math.inf - if lhs_literal == identity: - _alias_output(ctx, eqn.outvars[0], inputs[1]) - return None - if rhs_literal == identity: - _alias_output(ctx, eqn.outvars[0], inputs[0]) - return None - - lhs_scalar = _constant_scalar(graph, inputs[0]) - rhs_scalar = _constant_scalar(graph, inputs[1]) - condition_id = f"{outputs[0]}__{prim}_condition" - condition_shape = tuple(getattr(eqn.outvars[0].aval, "shape", ())) - condition_dtype = np.bool_ - if rhs_scalar is not None: - compare_op = "scalar_greater" if prim == "max" else "scalar_less" - _register_node( - graph, - IRNode( - ctx.node_id(f"{prim}_condition"), - compare_op, - [inputs[0]], - [condition_id], - attrs={"value": rhs_scalar}, - meta=meta, - ), - out_avals=(_SyntheticAval(condition_shape, condition_dtype),), - ) - attrs = {"false_is_scalar": True, "false_value": _where_scalar_value(rhs_scalar)} - return IRNode(ctx.node_id(f"{prim}_where"), "where", [condition_id, inputs[0]], outputs, attrs=attrs, meta=meta) - if lhs_scalar is not None: - compare_op = "scalar_less" if prim == "max" else "scalar_greater" - _register_node( - graph, - IRNode( - ctx.node_id(f"{prim}_condition"), - compare_op, - [inputs[1]], - [condition_id], - attrs={"value": lhs_scalar}, - meta=meta, - ), - out_avals=(_SyntheticAval(condition_shape, condition_dtype),), - ) - attrs = {"true_is_scalar": True, "true_value": _where_scalar_value(lhs_scalar)} - return IRNode(ctx.node_id(f"{prim}_where"), "where", [condition_id, inputs[1]], outputs, attrs=attrs, meta=meta) - - compare_op = "greater" if prim == "max" else "less" - _register_node( - graph, - IRNode(ctx.node_id(f"{prim}_condition"), compare_op, inputs, [condition_id], meta=meta), - out_avals=(_SyntheticAval(condition_shape, condition_dtype),), - ) - return IRNode(ctx.node_id(f"{prim}_where"), "where", [condition_id, inputs[0], inputs[1]], outputs, meta=meta) - if prim in {"reduce_sum", "reduce_max", "reduce_min"}: - reduce_ops = { - "reduce_sum": "sum", - "reduce_max": "max", - "reduce_min": "min", - } - return IRNode( - node_id, - reduce_ops[prim], - inputs, - outputs, - attrs={"axis": tuple(int(axis) for axis in params.get("axes", ()))}, - meta=meta, - ) - - raise NotImplementedError(f"unsupported JAX primitive: {prim}") - - -def capture_jax_function( - fn: Callable[..., Any], - example_args: Sequence[Any], - *, - constant_names: Sequence[str] | None = None, - weight_bindings: dict[str, dict[str, str]] | None = None, - weights_dir: str | None = None, - graph_meta: dict[str, object] | None = None, -) -> IRGraph: - try: - import jax - except Exception as exc: # pragma: no cover - depends on optional dependency - raise RuntimeError("capture_jax_function requires the optional jax package") from exc - - closed = jax.make_jaxpr(fn)(*example_args) - jaxpr = closed.jaxpr - constants = tuple(getattr(closed, "consts", ()) or ()) - names = tuple(constant_names or ()) - bindings = weight_bindings or {} - ctx = _JaxImportContext() - graph = IRGraph( - values={}, - nodes={}, - order=[], - inputs=[], - outputs=[], - constants={}, - meta={ - "frontend": "jax", - "adapter_family": "generic", - **({"weights_dir": weights_dir} if weights_dir else {}), - **dict(graph_meta or {}), - }, - ) - - for var in jaxpr.invars: - value_id = ctx.value_id(var) - _add_value(graph, value_id, var.aval) - graph.inputs.append(value_id) - - for index, (var, value) in enumerate(zip(jaxpr.constvars, constants, strict=True)): - value_id = ctx.value_id(var) - name = names[index] if index < len(names) else f"const_{index}" - tensor = _constant_to_torch(value) - graph.add_value( - IRValue( - id=value_id, - shape=tuple(int(dim) for dim in tensor.shape), - dtype=_dtype_to_ir(tensor.numpy().dtype), - meta={ - "source_name": name, - "jax_closed_constant": True, - **_binding_meta(name=name, weights_dir=weights_dir, explicit=bindings), - }, - ) - ) - graph.constants[value_id] = tensor - if "path" in graph.values[value_id].meta: - graph.meta.setdefault("weight_bindings", {})[value_id] = dict(graph.values[value_id].meta) - - for eqn in jaxpr.eqns: - _import_eqn(graph, ctx, eqn) - - graph.outputs = [ctx.value_id(var) for var in jaxpr.outvars] - apply_jax_semantic_rewrites(graph) - verify_ir(graph) - return graph - - -def capture_jax_function_with_params( - fn: Callable[..., Any], - params: Any, - example_args: Sequence[Any], - *, - weights_dir: str | None = None, - weight_bindings: dict[str, dict[str, str]] | None = None, - graph_meta: dict[str, object] | None = None, -) -> IRGraph: - try: - import jax - except Exception as exc: # pragma: no cover - depends on optional dependency - raise RuntimeError("capture_jax_function_with_params requires the optional jax package") from exc - - named_leaves = _flatten_named_leaves(jax.tree_util, params) - flat_leaves, treedef = jax.tree_util.tree_flatten(params) - if len(flat_leaves) != len(named_leaves): - raise ValueError("JAX param flattening produced mismatched leaf counts") - leaf_count = len(flat_leaves) - - def bound_fn(*flat_params_and_args: Any) -> Any: - flat_param_values = flat_params_and_args[:leaf_count] - runtime_args = flat_params_and_args[leaf_count:] - rebuilt_params = jax.tree_util.tree_unflatten(treedef, flat_param_values) - return fn(rebuilt_params, *runtime_args) - - graph = capture_jax_function( - bound_fn, - (*flat_leaves, *example_args), - weight_bindings=weight_bindings, - weights_dir=weights_dir, - graph_meta=graph_meta, - ) - explicit = weight_bindings or {} - _freeze_leading_param_inputs( - graph, - named_leaves, - weights_dir=weights_dir, - explicit=explicit, - ) - _propagate_weight_binding_meta(graph) - verify_ir(graph) - return graph - - -def capture_jax_graphs( - params: Any, - specs: Sequence[JaxGraphSpec], - *, - weights_dir: str | None = None, - weight_bindings: dict[str, dict[str, str]] | None = None, - graph_meta: dict[str, object] | None = None, -) -> CapturedJaxGraphBundle: - seen_names: set[str] = set() - for spec in specs: - if spec.name in seen_names: - raise ValueError(f"duplicate JAX graph spec name: {spec.name!r}") - seen_names.add(spec.name) - if spec.input_names is not None and len(spec.input_names) != len(spec.example_args): - raise ValueError( - f"graph {spec.name!r} has {len(spec.input_names)} input names for " - f"{len(spec.example_args)} example inputs" - ) - - captured: dict[str, CapturedJaxGraph] = {} - for spec in specs: - ir = capture_jax_function_with_params( - spec.fn, - params, - tuple(spec.example_args), - weights_dir=weights_dir, - weight_bindings=weight_bindings, - graph_meta={ - **dict(graph_meta or {}), - **dict(spec.graph_meta or {}), - "jax_graph_name": spec.name, - "jax_graph_role": spec.role, - **({"input_names": tuple(spec.input_names)} if spec.input_names is not None else {}), - **({"output_names": tuple(spec.output_names)} if spec.output_names is not None else {}), - }, - ) - raw_ir = copy.deepcopy(ir) - transpiled_graph = transpile_ir(ir) - captured[spec.name] = CapturedJaxGraph( - spec=spec, - raw_ir_graph=raw_ir, - ir_graph=ir, - graph=transpiled_graph, - ) - return CapturedJaxGraphBundle(graphs=captured, params=params, weights_dir=weights_dir) - - -__all__ = [ - "CapturedJaxGraph", - "CapturedJaxGraphBundle", - "JaxGraphSpec", - "capture_jax_function", - "capture_jax_function_with_params", - "capture_jax_graphs", -] diff --git a/python/cactus/transpile/capture_pytorch.py b/python/cactus/transpile/capture_pytorch.py deleted file mode 100644 index 22c63833e..000000000 --- a/python/cactus/transpile/capture_pytorch.py +++ /dev/null @@ -1,375 +0,0 @@ -from dataclasses import dataclass -import inspect -from contextlib import contextmanager -from typing import Any -from typing import Iterable - -import torch -from torch.export import export -from torch.fx.passes.shape_prop import ShapeProp - -from cactus.transpile.aten_ops import canonical_torch_op -from cactus.transpile.graph_ir import IRGraph - - -class CapturePhaseError(RuntimeError): - def __init__(self, phase: str, message: str, *, cause: Exception | None = None): - super().__init__(f"[capture:{phase}] {message}") - self.phase = phase - self.cause = cause - - -@dataclass -class TensorMetadata: - shape: tuple[Any, ...] - dtype: Any - stride: tuple[Any, ...] | None = None - device: Any | None = None - - -@dataclass -class CapturedModel: - exported_program: Any - graph_module: Any - graph: Any - state_dict: dict[str, Any] - ir_graph: IRGraph - source_module: Any - example_args: tuple[Any, ...] - example_kwargs: dict[str, Any] - strict: bool - transpile_metadata: dict[str, Any] - - def named_parameters(self): - return self.exported_program.named_parameters() - - def named_buffers(self): - return self.exported_program.named_buffers() - - def parameters_dict(self) -> dict[str, Any]: - return dict(self.named_parameters()) - - def buffers_dict(self) -> dict[str, Any]: - return dict(self.named_buffers()) - - def placeholders(self) -> list[Any]: - return [node for node in self.graph.nodes if node.op == "placeholder"] - - def outputs(self) -> list[Any]: - return [node for node in self.graph.nodes if node.op == "output"] - - -def _normalize_args(args: Any) -> tuple[Any, ...]: - if args is None: - return () - if isinstance(args, tuple): - return args - if isinstance(args, list): - return tuple(args) - return (args,) - - -def _normalize_kwargs(kwargs: dict[str, Any] | None) -> dict[str, Any]: - return {} if kwargs is None else dict(kwargs) - - -def _clone_example_value(value: Any) -> Any: - if isinstance(value, torch.Tensor): - return value.detach().clone() - if isinstance(value, tuple): - return tuple(_clone_example_value(v) for v in value) - if isinstance(value, list): - return [_clone_example_value(v) for v in value] - if isinstance(value, dict): - return {k: _clone_example_value(v) for k, v in value.items()} - return value - - -def _clone_examples(args: tuple[Any, ...], kwargs: dict[str, Any]) -> tuple[tuple[Any, ...], dict[str, Any]]: - return ( - tuple(_clone_example_value(arg) for arg in args), - {k: _clone_example_value(v) for k, v in kwargs.items()}, - ) - - -def _inject_export_safe_kwargs(model: torch.nn.Module, kwargs: dict[str, Any]) -> dict[str, Any]: - try: - signature = inspect.signature(model.forward) - except (TypeError, ValueError): - return kwargs - - updated = dict(kwargs) - parameters = signature.parameters - - # HF causal LM forwards often default to returning cache objects and ModelOutput - # containers, which torch.export cannot serialize cleanly. Keep the capture - # boundary tensor-only unless the caller explicitly asked otherwise. - if "use_cache" in parameters and "use_cache" not in updated: - updated["use_cache"] = False - if "return_dict" in parameters and "return_dict" not in updated: - updated["return_dict"] = False - return updated - - -def _call_transpile_metadata_provider(module: torch.nn.Module, args: tuple[Any, ...], kwargs: dict[str, Any]) -> dict[str, Any]: - for attr_name in ("get_transpile_metadata", "transpile_metadata"): - provider = getattr(module, attr_name, None) - if not callable(provider): - continue - try: - metadata = provider(args=args, kwargs=kwargs) - except TypeError: - metadata = provider() - if metadata is None: - return {} - if not isinstance(metadata, dict): - raise TypeError(f"{type(module).__name__}.{attr_name} must return a dict or None") - return metadata - return {} - - -def _collect_transpile_metadata(model: torch.nn.Module, args: tuple[Any, ...], kwargs: dict[str, Any]) -> dict[str, Any]: - graph_meta: dict[str, Any] = {} - provider_graph_meta: dict[str, dict[str, Any]] = {} - seen_modules: set[int] = set() - - for module_path, module in model.named_modules(): - module_id = id(module) - if module_id in seen_modules: - continue - seen_modules.add(module_id) - - metadata = _call_transpile_metadata_provider(module, args, kwargs) - if not metadata: - continue - - provider_key = module_path or type(module).__name__ - provider_meta = metadata.get("graph", {}) - if isinstance(provider_meta, dict) and provider_meta: - provider_graph_meta[provider_key] = dict(provider_meta) - if not module_path: - graph_meta.update(provider_meta) - - if provider_graph_meta: - graph_meta.setdefault("transpile_metadata_providers", provider_graph_meta) - - if not graph_meta: - return {} - return {"graph": graph_meta} - - -def _prepare_import_graph_module(ep: Any, args: tuple[Any, ...], kwargs: dict[str, Any]) -> Any: - graph_module = ep.graph_module - if kwargs: - return graph_module - - try: - graph_module = ep.module() - ShapeProp(graph_module).propagate(*args) - except Exception: - return ep.graph_module - return graph_module - - -@contextmanager -def _suppress_transformers_model_output_registration(): - """Avoid tracing a Transformers global-set mutation inside ModelOutput init. - - Recent Transformers versions lazily register ModelOutput dataclasses as - PyTrees from their __post_init__. That side effect is irrelevant for our - capture wrapper because the exported boundary returns tensors, but Dynamo - sees the intermediate set membership check and rejects it. - """ - - try: - import transformers.utils.generic as hf_generic - except Exception: - yield - return - - original = getattr(hf_generic, "_register_model_output_pytree_node", None) - if original is None: - yield - return - - def _noop_register_model_output_pytree_node(*_args: Any, **_kwargs: Any) -> None: - return None - - hf_generic._register_model_output_pytree_node = _noop_register_model_output_pytree_node - try: - yield - finally: - hf_generic._register_model_output_pytree_node = original - - -def capture_model(model, args, kwargs=None, *, strict=True) -> CapturedModel: - if not isinstance(model, torch.nn.Module): - raise TypeError("model must be a torch.nn.Module") - - normalized_args = _normalize_args(args) - normalized_kwargs = _inject_export_safe_kwargs(model, _normalize_kwargs(kwargs)) - - if not normalized_args and not normalized_kwargs: - raise ValueError("capture_model requires example args or kwargs") - - # nn.Module.eval() is an in-place mutation on the caller's model. Remember - # the prior training flag so capture failures don't leave the caller's - # module silently in eval mode. - was_training = bool(getattr(model, "training", False)) - model.eval() - try: - example_args, example_kwargs = _clone_examples(normalized_args, normalized_kwargs) - transpile_metadata = _collect_transpile_metadata(model, example_args, example_kwargs) - - try: - with _suppress_transformers_model_output_registration(): - ep = export(model, args=example_args, kwargs=example_kwargs, strict=strict) - except Exception as exc: - raise CapturePhaseError( - "export", - f"torch.export failed for model={type(model).__name__} strict={strict}: {exc}", - cause=exc, - ) from exc - from cactus.transpile.import_ir import import_captured_to_ir - - import_graph_module = _prepare_import_graph_module(ep, example_args, example_kwargs) - - raw_captured = CapturedModel( - exported_program=ep, - graph_module=import_graph_module, - graph=import_graph_module.graph, - state_dict=dict(ep.state_dict), - ir_graph=IRGraph(values={}, nodes={}, order=[], inputs=[], outputs=[]), - source_module=model, - example_args=example_args, - example_kwargs=example_kwargs, - strict=strict, - transpile_metadata=transpile_metadata, - ) - try: - ir_graph = import_captured_to_ir(raw_captured, strict=strict) - except Exception as exc: - raise CapturePhaseError( - "import", - f"failed to import exported graph to IR for model={type(model).__name__} strict={strict}: {exc}", - cause=exc, - ) from exc - except Exception: - if was_training: - try: - model.train() - except Exception: - pass - raise - - return CapturedModel( - exported_program=ep, - graph_module=ep.graph_module, - graph=ep.graph, - state_dict=dict(ep.state_dict), - ir_graph=ir_graph, - source_module=model, - example_args=example_args, - example_kwargs=example_kwargs, - strict=strict, - transpile_metadata=transpile_metadata, - ) - - -def capture_model_with_fallback(model, args, kwargs=None) -> CapturedModel: - try: - return capture_model(model, args, kwargs=kwargs, strict=True) - except Exception: - return capture_model(model, args, kwargs=kwargs, strict=False) - - -def resolve_attr(root: Any, target: str) -> Any: - obj = root - for atom in target.split("."): - obj = getattr(obj, atom) - return obj - - -def format_target(node: Any) -> str: - return canonical_torch_op(node.target) - - -def _try_materialize_int_tuple(values: Iterable[Any]) -> tuple[int, ...] | None: - materialized: list[int] = [] - try: - for value in values: - materialized.append(int(value)) - except Exception: - return None - return tuple(materialized) - - -def get_tensor_metadata(node: Any) -> TensorMetadata | None: - meta = getattr(node, "meta", {}) or {} - tensor_meta = meta.get("tensor_meta") - if tensor_meta is not None: - shape = tuple(getattr(tensor_meta, "shape", ())) - dtype = getattr(tensor_meta, "dtype", None) - stride_value = getattr(tensor_meta, "stride", None) - stride = _try_materialize_int_tuple(stride_value) if isinstance(stride_value, Iterable) else None - device = getattr(tensor_meta, "device", None) - if shape or dtype is not None or stride is not None or device is not None: - return TensorMetadata( - shape=shape, - dtype=dtype, - stride=stride, - device=device, - ) - - value = meta.get("val") - if value is None: - return None - - if isinstance(value, torch.Tensor): - stride = _try_materialize_int_tuple(value.stride()) if value.layout == torch.strided else None - return TensorMetadata( - shape=tuple(value.shape), - dtype=value.dtype, - stride=stride, - device=value.device, - ) - - shape = tuple(getattr(value, "shape", ())) - dtype = getattr(value, "dtype", None) - stride_value = getattr(value, "stride", None) - stride = tuple(stride_value) if isinstance(stride_value, Iterable) else None - device = getattr(value, "device", None) - if not shape and dtype is None and stride is None and device is None: - return None - return TensorMetadata(shape=shape, dtype=dtype, stride=stride, device=device) - - -def get_shape(node: Any) -> tuple[Any, ...] | None: - metadata = get_tensor_metadata(node) - return None if metadata is None else metadata.shape - - -def get_dtype(node: Any) -> Any | None: - metadata = get_tensor_metadata(node) - return None if metadata is None else metadata.dtype - - -def dump_graph(captured: CapturedModel, *, include_meta: bool = True) -> str: - lines: list[str] = [] - for index, node in enumerate(captured.graph.nodes): - lines.append(f"[{index}] {node.op} {node.name} target={node.target}") - lines.append(f" args={node.args}") - lines.append(f" kwargs={node.kwargs}") - if include_meta: - tensor_meta = get_tensor_metadata(node) - if tensor_meta is not None: - lines.append( - " tensor_meta=" - f"shape={tensor_meta.shape} dtype={tensor_meta.dtype} " - f"stride={tensor_meta.stride} device={tensor_meta.device}" - ) - else: - lines.append(f" meta={getattr(node, 'meta', {})}") - return "\n".join(lines) - - diff --git a/python/cactus/transpile/component_partition.py b/python/cactus/transpile/component_partition.py deleted file mode 100644 index a0a76dd32..000000000 --- a/python/cactus/transpile/component_partition.py +++ /dev/null @@ -1,335 +0,0 @@ -from __future__ import annotations - -import copy -from collections import Counter - -from cactus.transpile.graph_ir import IRGraph -from cactus.transpile.graph_ir import IRNode -from cactus.transpile.graph_ir import IRValue -from cactus.transpile.graph_ir import verify_ir -from cactus.transpile.model_profiles import profile_for_family - - -COMPONENT_AUDIO_ENCODER = "audio_encoder" -COMPONENT_VISION_ENCODER = "vision_encoder" -COMPONENT_LM_ENCODER = "lm_encoder" -COMPONENT_DECODER = "decoder" -COMPONENT_UNSPECIFIED = "unspecified" - -COMPONENT_ORDER = ( - COMPONENT_AUDIO_ENCODER, - COMPONENT_VISION_ENCODER, - COMPONENT_LM_ENCODER, - COMPONENT_DECODER, -) - - -def _normalized_name_candidates( - *, - value_id: str | None = None, - meta: dict[str, object] | None = None, -) -> tuple[str, ...]: - names: list[str] = [] - - def _add(value: object) -> None: - if not isinstance(value, str): - return - normalized = value.strip().lower() - if normalized and normalized not in names: - names.append(normalized) - - if value_id is not None: - _add(value_id) - if isinstance(meta, dict): - for key in ("logical_name", "input_name", "output_name", "component_role"): - _add(meta.get(key)) - return tuple(names) - - -def _contains_any(name: str, needles: tuple[str, ...]) -> bool: - return any(needle in name for needle in needles) - - -def classify_component_semantic_candidates( - candidates: tuple[str, ...], - *, - task: str, -) -> str | None: - if not candidates: - return None - - for name in candidates: - if _contains_any( - name, - ( - "pixel_values", - "image_features", - ), - ): - return COMPONENT_VISION_ENCODER - - for name in candidates: - if _contains_any( - name, - ( - "input_features", - "audio_features", - "speech_features", - "mel_features", - ), - ): - return COMPONENT_AUDIO_ENCODER - - for name in candidates: - if _contains_any( - name, - ( - "logits", - "next_token", - "ctc_head", - "ctc_logits", - "decoder_output", - ), - ): - return COMPONENT_DECODER - - for name in candidates: - if task == "multimodal_causal_lm_logits" and _contains_any( - name, - ( - "input_ids", - "inputs_embeds", - "text_embeds", - "position_ids", - "placeholder", - "merge", - "masked_scatter", - "token_type", - "attention_mask", - "pixel_position", - ), - ): - return COMPONENT_LM_ENCODER - - return None - - -def classify_value_component(value_id: str, value: IRValue, *, family: str, task: str) -> str | None: - return classify_component_semantic_candidates( - _normalized_name_candidates( - value_id=value_id, - meta=value.meta, - ), - task=task, - ) - - -def classify_node_component(node: IRNode, *, family: str, task: str) -> str | None: - if task == "multimodal_causal_lm_logits" and node.op in {"embedding", "masked_scatter", "advanced_index"}: - return COMPONENT_LM_ENCODER - return classify_component_semantic_candidates( - _normalized_name_candidates(meta=node.meta), - task=task, - ) - - -def _task_default_component(*, family: str, task: str) -> str: - if task == "multimodal_causal_lm_logits": - return COMPONENT_DECODER - if task == "causal_lm_logits": - return COMPONENT_DECODER - if task in {"ctc_logits", "encoder_hidden_states"}: - return COMPONENT_AUDIO_ENCODER - profile = profile_for_family(family) - if profile is not None and profile.default_task == "tdt_transcription": - return COMPONENT_AUDIO_ENCODER - return COMPONENT_UNSPECIFIED - - -def _rebuild_users(graph: IRGraph) -> None: - for value in graph.values.values(): - value.users = [] - for value_id in graph.inputs: - graph.values[value_id].producer = None - for value_id in graph.constants: - graph.values[value_id].producer = None - for node_id in graph.order: - node = graph.nodes[node_id] - for output_id in node.outputs: - graph.values[output_id].producer = node_id - for input_id in node.inputs: - graph.values[input_id].users.append(node_id) - - -def annotate_ir_components(graph: IRGraph) -> None: - family = str(graph.meta.get("adapter_family", "") or "").lower() - task = str(graph.meta.get("task", "") or "").lower() - graph_component = str(graph.meta.get("component", "") or "").strip() - if graph_component in COMPONENT_ORDER: - for value in graph.values.values(): - value.meta["component"] = graph_component - for node_id in graph.order: - node = graph.nodes[node_id] - node.meta["component"] = graph_component - for output_id in node.outputs: - graph.values[output_id].meta["component"] = graph_component - graph.meta["component_counts"] = {graph_component: len(graph.order)} - return - - default_component = _task_default_component(family=family, task=task) - - for value_id, value in graph.values.items(): - component = classify_value_component(value_id, value, family=family, task=task) - if component is not None: - value.meta["component"] = component - - for node_id in graph.order: - node = graph.nodes[node_id] - component = classify_node_component(node, family=family, task=task) - if component is not None: - node.meta["component"] = component - for output_id in node.outputs: - graph.values[output_id].meta["component"] = component - - for _ in range(8): - changed = False - for node_id in graph.order: - node = graph.nodes[node_id] - component = node.meta.get("component") - if component not in COMPONENT_ORDER: - input_components = { - graph.values[input_id].meta.get("component") - for input_id in node.inputs - if graph.values[input_id].meta.get("component") in COMPONENT_ORDER - } - if len(input_components) == 1: - component = next(iter(input_components)) - node.meta["component"] = component - changed = True - component = node.meta.get("component") - if component in COMPONENT_ORDER: - for output_id in node.outputs: - if graph.values[output_id].meta.get("component") != component: - graph.values[output_id].meta["component"] = component - changed = True - - for node_id in reversed(graph.order): - node = graph.nodes[node_id] - component = node.meta.get("component") - if component in COMPONENT_ORDER: - continue - output_components: set[str] = set() - for output_id in node.outputs: - value = graph.values[output_id] - for user_id in value.users: - user_component = graph.nodes[user_id].meta.get("component") - if user_component in COMPONENT_ORDER: - output_components.add(user_component) - if len(output_components) == 1: - component = next(iter(output_components)) - node.meta["component"] = component - for output_id in node.outputs: - graph.values[output_id].meta["component"] = component - changed = True - - if not changed: - break - - for value in graph.values.values(): - component = value.meta.get("component") - if component not in COMPONENT_ORDER and default_component != COMPONENT_UNSPECIFIED: - value.meta["component"] = default_component - - for node_id in graph.order: - node = graph.nodes[node_id] - component = node.meta.get("component") - if component not in COMPONENT_ORDER and default_component != COMPONENT_UNSPECIFIED: - node.meta["component"] = default_component - for output_id in node.outputs: - graph.values[output_id].meta["component"] = default_component - - counts = Counter( - str(graph.nodes[node_id].meta.get("component", COMPONENT_UNSPECIFIED)) - for node_id in graph.order - ) - graph.meta["component_counts"] = dict(sorted(counts.items())) - - -def summarize_ir_components(graph: IRGraph) -> dict[str, int]: - annotate_ir_components(graph) - counts = graph.meta.get("component_counts", {}) - if isinstance(counts, dict): - return {str(key): int(value) for key, value in counts.items()} - return {} - - -def _clone_value_for_subgraph(value: IRValue, *, as_input: bool) -> IRValue: - cloned = copy.deepcopy(value) - if as_input: - cloned.producer = None - cloned.users = [] - return cloned - - -def extract_component_subgraphs(graph: IRGraph) -> dict[str, IRGraph]: - annotate_ir_components(graph) - component_graphs: dict[str, IRGraph] = {} - - for component in COMPONENT_ORDER: - node_ids = [ - node_id - for node_id in graph.order - if graph.nodes[node_id].meta.get("component") == component - ] - if not node_ids: - continue - - node_id_set = set(node_ids) - inputs: list[str] = [] - outputs: list[str] = [] - constants: dict[str, object] = {} - values: dict[str, IRValue] = {} - nodes: dict[str, IRNode] = {} - - for node_id in node_ids: - node = copy.deepcopy(graph.nodes[node_id]) - nodes[node_id] = node - for input_id in node.inputs: - if input_id in graph.constants: - constants[input_id] = graph.constants[input_id] - if input_id not in values: - values[input_id] = _clone_value_for_subgraph(graph.values[input_id], as_input=True) - continue - producer = graph.values[input_id].producer - if producer not in node_id_set and input_id not in inputs: - inputs.append(input_id) - values[input_id] = _clone_value_for_subgraph(graph.values[input_id], as_input=True) - for output_id in node.outputs: - if output_id not in values: - values[output_id] = _clone_value_for_subgraph(graph.values[output_id], as_input=False) - - for node_id in node_ids: - node = graph.nodes[node_id] - for output_id in node.outputs: - value = graph.values[output_id] - if output_id in graph.outputs or any(user_id not in node_id_set for user_id in value.users): - if output_id not in outputs: - outputs.append(output_id) - - component_graph = IRGraph( - values=values, - nodes=nodes, - order=list(node_ids), - inputs=inputs, - outputs=outputs, - constants=constants, - meta={ - **copy.deepcopy(graph.meta), - "component": component, - }, - ) - _rebuild_users(component_graph) - verify_ir(component_graph) - component_graphs[component] = component_graph - - return component_graphs diff --git a/python/cactus/transpile/component_pipeline.py b/python/cactus/transpile/component_pipeline.py deleted file mode 100644 index 5cc9af49b..000000000 --- a/python/cactus/transpile/component_pipeline.py +++ /dev/null @@ -1,213 +0,0 @@ -from __future__ import annotations - -import copy -from contextlib import AbstractContextManager -from dataclasses import dataclass -from dataclasses import field -import re -from typing import Any, Callable - -import numpy as np -import torch - -from cactus.transpile.capture_pytorch import capture_model -from cactus.transpile.canonicalize.cleanup import canonicalize_exported_graph -from cactus.transpile.graph_ir import IRGraph -from cactus.transpile.lower import TranspiledGraph -from cactus.transpile.lower import transpile_preoptimized_ir -from cactus.transpile.optimize_graph import FusionConfig -from cactus.transpile.optimize_graph import optimize_graph -from cactus.transpile.optimize_graph import precompute_rope_tables - - -@dataclass -class ComponentModuleSpec: - component: str - module: torch.nn.Module - example_inputs: tuple[torch.Tensor, ...] - input_keys: tuple[str, ...] - output_keys: tuple[str, ...] - graph_meta: dict[str, object] = field(default_factory=dict) - metadata: dict[str, object] = field(default_factory=dict) - dynamic_batch_axis: int | None = None - npu_module: torch.nn.Module | None = None - npu_runtime_input_count: int = 1 - npu_reparam: Callable[[torch.nn.Module], AbstractContextManager] | None = None - npu_example_inputs: tuple[torch.Tensor, ...] | None = None - - -@dataclass -class CapturedComponent: - component: str - input_keys: tuple[str, ...] - output_keys: tuple[str, ...] - raw_ir_graph: IRGraph - optimized_ir_graph: IRGraph - transpiled_graph: TranspiledGraph - metadata: dict[str, object] = field(default_factory=dict) - - -_GRAPH_ARG_INPUT_RE = re.compile(r"^v_args_(\d+)$") - - -def _align_input_keys_to_graph_inputs( - declared_input_keys: tuple[str, ...], - graph_inputs: tuple[str, ...] | list[str], -) -> tuple[str, ...]: - aligned: list[str] = [] - for value_id in graph_inputs: - match = _GRAPH_ARG_INPUT_RE.match(str(value_id)) - if match is None: - return declared_input_keys - input_index = int(match.group(1)) - if input_index < 0 or input_index >= len(declared_input_keys): - return declared_input_keys - aligned.append(declared_input_keys[input_index]) - return tuple(aligned) - - -def _call_optional_capture_hook(spec: ComponentModuleSpec, attr_name: str) -> None: - hook = getattr(spec.module, attr_name, None) - if not callable(hook): - return - try: - hook( - component=spec.component, - example_inputs=spec.example_inputs, - input_keys=spec.input_keys, - output_keys=spec.output_keys, - ) - except TypeError: - hook() - - -def capture_component_spec( - spec: ComponentModuleSpec, - *, - fusion_config: FusionConfig | None = None, -) -> CapturedComponent: - wrapped = _ComponentCaptureWrapper(spec.module, graph_meta=spec.graph_meta).eval() - capture_error: Exception | None = None - _call_optional_capture_hook(spec, "prepare_for_capture") - try: - captured = capture_model(wrapped, spec.example_inputs) - except Exception as exc: - capture_error = exc - raise - finally: - try: - _call_optional_capture_hook(spec, "restore_after_capture") - except Exception: - if capture_error is None: - raise - raw_ir_graph = copy.deepcopy(captured.ir_graph) - optimized_ir_graph = copy.deepcopy(captured.ir_graph) - canonicalize_exported_graph(optimized_ir_graph) - optimize_graph(optimized_ir_graph, config=fusion_config, precompute_rope=False) - # Bake rope tables only into the lowered graph.cactus; the saved optimized IR stays un-baked. - lowered_ir = copy.deepcopy(optimized_ir_graph) - if precompute_rope_tables(lowered_ir): - canonicalize_exported_graph(lowered_ir) - if spec.dynamic_batch_axis is not None: - axis = int(spec.dynamic_batch_axis) - for vid in lowered_ir.inputs: - v = lowered_ir.values.get(vid) - if v is None or v.shape is None or axis >= len(v.shape): - continue - mask = [False] * len(v.shape) - mask[axis] = True - v.meta["dynamic_dims"] = tuple(mask) - transpiled_graph = transpile_preoptimized_ir(lowered_ir) - if len(spec.output_keys) != len(transpiled_graph.outputs): - raise ValueError( - f"component {spec.component} declared {len(spec.output_keys)} output keys but lowered " - f"{len(transpiled_graph.outputs)} graph outputs" - ) - aligned_input_keys = _align_input_keys_to_graph_inputs(spec.input_keys, optimized_ir_graph.inputs) - return CapturedComponent( - component=spec.component, - input_keys=aligned_input_keys, - output_keys=spec.output_keys, - raw_ir_graph=raw_ir_graph, - optimized_ir_graph=optimized_ir_graph, - transpiled_graph=transpiled_graph, - metadata=dict(spec.metadata), - ) - - -def execute_component_pipeline( - components: list[CapturedComponent], - *, - initial_store: dict[str, Any], -) -> tuple[dict[str, np.ndarray], dict[str, list[np.ndarray]]]: - store: dict[str, np.ndarray] = {} - for key, value in initial_store.items(): - store[key] = _to_numpy(value) - - outputs_by_component: dict[str, list[np.ndarray]] = {} - for component in components: - runtime_inputs = [] - for input_key in component.input_keys: - if input_key not in store: - raise KeyError( - f"component {component.component} is missing pipeline input {input_key!r}" - ) - runtime_inputs.append(store[input_key]) - component.transpiled_graph.set_inputs(runtime_inputs) - raw_outputs = component.transpiled_graph.execute() - numpy_outputs = [output.numpy().copy() for output in raw_outputs] - if len(numpy_outputs) != len(component.output_keys): - raise ValueError( - f"component {component.component} produced {len(numpy_outputs)} outputs, " - f"expected {len(component.output_keys)}" - ) - for output_key, value in zip(component.output_keys, numpy_outputs, strict=True): - store[output_key] = value - outputs_by_component[component.component] = numpy_outputs - - return store, outputs_by_component - - -def initial_store_from_named_tensors( - keys: tuple[str, ...] | list[str], - tensors: tuple[torch.Tensor, ...] | list[torch.Tensor], -) -> dict[str, np.ndarray]: - if len(keys) != len(tensors): - raise ValueError(f"expected {len(keys)} tensors, got {len(tensors)}") - return { - str(key): _to_numpy(tensor) - for key, tensor in zip(keys, tensors, strict=True) - } - - -def _to_numpy(value: Any) -> np.ndarray: - if isinstance(value, np.ndarray): - return np.ascontiguousarray(value) - if isinstance(value, torch.Tensor): - return np.ascontiguousarray(value.detach().cpu().numpy()) - raise TypeError(f"unsupported pipeline value type: {type(value).__name__}") - - -class _ComponentCaptureWrapper(torch.nn.Module): - def __init__(self, module: torch.nn.Module, *, graph_meta: dict[str, object]): - super().__init__() - self.module = module - self.graph_meta = dict(graph_meta) - - def forward(self, *args: torch.Tensor) -> Any: - return self.module(*args) - - def get_transpile_metadata(self) -> dict[str, object]: - metadata: dict[str, object] = {} - provider = getattr(self.module, "get_transpile_metadata", None) - if callable(provider): - provided = provider() - if isinstance(provided, dict): - metadata.update(provided) - graph_meta = {} - existing = metadata.get("graph", {}) - if isinstance(existing, dict): - graph_meta.update(existing) - graph_meta.update(self.graph_meta) - metadata["graph"] = graph_meta - return metadata diff --git a/python/cactus/transpile/component_plan.py b/python/cactus/transpile/component_plan.py deleted file mode 100644 index c9e63cabd..000000000 --- a/python/cactus/transpile/component_plan.py +++ /dev/null @@ -1,239 +0,0 @@ -from __future__ import annotations - -from dataclasses import dataclass -import json -from pathlib import Path -from typing import Mapping - -from cactus.transpile.model_profiles import ModelProfile -from cactus.transpile.model_profiles import profile_for_model_id -from cactus.transpile.model_profiles import profile_for_model_type - - -@dataclass(frozen=True) -class ComponentPlan: - task: str - components: tuple[str, ...] = () - default_max_new_tokens: int | None = None - needs_image: bool = False - needs_audio: bool = False - force_component_pipeline: bool = False - - -def _plan_from_profile( - profile: ModelProfile, - config: Mapping[str, object] | None = None, -) -> ComponentPlan | None: - if profile.default_task is None or not profile.default_components: - return None - if config is not None and profile.needs_image: - has_vision = _has_dict_config(config, "vision_config", "visual_config", "image_config") - if not has_vision: - architectures = [str(a).lower() for a in (config.get("architectures") or ())] - for arch_marker, components in profile.text_component_plans: - if any(arch_marker in a for a in architectures): - return ComponentPlan( - task="causal_lm_logits", - components=components, - force_component_pipeline=True, - ) - return ComponentPlan( - task="causal_lm_logits", - components=("decoder", "decoder_step"), - force_component_pipeline=True, - ) - return ComponentPlan( - task=profile.default_task, - components=profile.default_components, - default_max_new_tokens=profile.default_max_new_tokens, - needs_image=profile.needs_image, - needs_audio=profile.needs_audio, - force_component_pipeline=profile.force_component_pipeline, - ) - - -def _load_json(path: Path) -> dict[str, object]: - try: - value = json.loads(path.read_text(encoding="utf-8")) - except Exception: - return {} - return value if isinstance(value, dict) else {} - - -def _load_config_txt(path: Path) -> dict[str, str]: - config: dict[str, str] = {} - try: - lines = path.read_text(encoding="utf-8").splitlines() - except Exception: - return config - for line in lines: - if "=" not in line: - continue - key, value = line.split("=", 1) - config[key.strip()] = value.strip() - return config - - -def _architecture_names(config: Mapping[str, object], config_txt: Mapping[str, str] | None = None) -> tuple[str, ...]: - raw_architectures = config.get("architectures", []) - names = tuple(str(value) for value in raw_architectures if isinstance(value, str)) - if names or config_txt is None: - return names - raw_txt = config_txt.get("architectures") - if not raw_txt: - return () - return tuple(value.strip() for value in raw_txt.split(",") if value.strip()) - - -def _has_dict_config(config: Mapping[str, object], *keys: str) -> bool: - return any(isinstance(config.get(key), dict) for key in keys) - - -def _has_positive_txt_int(config_txt: Mapping[str, str] | None, *keys: str) -> bool: - if config_txt is None: - return False - for key in keys: - try: - if int(config_txt.get(key, "0") or 0) > 0: - return True - except ValueError: - continue - return False - - -def _model_type(config: Mapping[str, object], config_txt: Mapping[str, str] | None = None) -> str: - value = config.get("model_type") - if value is None and config_txt is not None: - value = config_txt.get("model_type") - return str(value or "").strip().lower() - - -def _lowered_architectures(config: Mapping[str, object], config_txt: Mapping[str, str] | None = None) -> tuple[str, ...]: - return tuple(value.lower() for value in _architecture_names(config, config_txt)) - - -def _is_tdt_config(config: Mapping[str, object], model_type: str, lowered_id: str) -> bool: - decoding_cfg = config.get("decoding") - if isinstance(decoding_cfg, Mapping) and str(decoding_cfg.get("model_type", "") or "").lower() == "tdt": - return True - loss_cfg = config.get("loss") - if isinstance(loss_cfg, Mapping) and str(loss_cfg.get("loss_name", "") or "").lower() == "tdt": - return True - return model_type == "parakeet_tdt" or "parakeet-tdt" in lowered_id - - -def _looks_like_vision_language_model( - *, - model_type: str, - architectures: tuple[str, ...], - lowered_id: str, - has_vision: bool, -) -> bool: - if not has_vision: - return False - if any(token in model_type for token in ("vl", "vision", "image", "multimodal")): - return True - if any( - token in architecture - for architecture in architectures - for token in ("vision", "image", "vl", "imagetexttotext", "conditionalgeneration") - ): - return True - return any(token in lowered_id for token in ("-vl", "_vl", "vision", "image")) - - -def infer_component_plan_from_config( - config: Mapping[str, object], - *, - model_id: str = "", - config_txt: Mapping[str, str] | None = None, -) -> ComponentPlan | None: - model_type = _model_type(config, config_txt) - lowered_id = model_id.lower() - architectures = _lowered_architectures(config, config_txt) - profile = profile_for_model_type(model_type) or profile_for_model_id(model_id) - profile_plan = _plan_from_profile(profile, config) if profile is not None else None - if profile_plan is not None: - return profile_plan - - if "nomic" in model_type or "nomic-embed" in lowered_id or any("nomicbert" in value for value in architectures): - return ComponentPlan( - task="text_embedding", - components=("text_embedding",), - force_component_pipeline=True, - ) - - if _is_tdt_config(config, model_type, lowered_id): - return ComponentPlan( - task="tdt_transcription", - components=("audio_encoder", "decoder"), - needs_audio=True, - force_component_pipeline=True, - ) - - if model_type == "whisper" or "whisper" in lowered_id or any("whisper" in value for value in architectures): - return ComponentPlan( - task="seq2seq_transcription", - components=("audio_encoder", "decoder"), - needs_audio=True, - force_component_pipeline=True, - ) - - if any("ctc" in value for value in architectures) or "ctc" in lowered_id: - return ComponentPlan( - task="ctc_logits", - components=("audio_encoder",), - needs_audio=True, - ) - - has_vision = ( - _has_dict_config(config, "vision_config", "visual_config", "image_config") - or _has_positive_txt_int(config_txt, "vision_num_layers") - ) - has_audio = ( - _has_dict_config(config, "audio_config", "speech_config", "acoustic_config") - or ( - _has_dict_config(config, "encoder_config") - and ( - "audio" in model_type - or "speech" in model_type - or "audio_token_index" in config - or any(token in lowered_id for token in ("speech", "audio")) - or any(token in value for value in architectures for token in ("speech", "audio")) - ) - ) - or _has_positive_txt_int(config_txt, "audio_num_layers") - ) - if has_vision or has_audio: - if has_audio or _looks_like_vision_language_model( - model_type=model_type, - architectures=architectures, - lowered_id=lowered_id, - has_vision=has_vision, - ): - components: list[str] = [] - if has_vision: - components.append("vision_encoder") - if has_audio: - components.append("audio_encoder") - components.extend(("lm_encoder", "decoder")) - return ComponentPlan( - task="multimodal_causal_lm_logits", - components=tuple(components), - needs_image=has_vision, - needs_audio=has_audio, - force_component_pipeline=True, - ) - - if any("causallm" in value for value in architectures): - return ComponentPlan(task="causal_lm_logits", components=("decoder",)) - if any(token in lowered_id or token in model_type for token in ("qwen", "gemma", "llama", "mistral", "lfm")): - return ComponentPlan(task="causal_lm_logits", components=("decoder",)) - return None - - -def infer_component_plan_from_output(output_dir: str | Path, *, model_id: str = "") -> ComponentPlan | None: - root = Path(output_dir) - hf_config = _load_json(root / "hf_config.json") - config_txt = _load_config_txt(root / "config.txt") - return infer_component_plan_from_config(hf_config, model_id=model_id, config_txt=config_txt) diff --git a/python/cactus/transpile/fusion/__init__.py b/python/cactus/transpile/fusion/__init__.py deleted file mode 100644 index 463478d1c..000000000 --- a/python/cactus/transpile/fusion/__init__.py +++ /dev/null @@ -1,47 +0,0 @@ -from cactus.transpile.fusion.conv import ConvModuleMatch -from cactus.transpile.fusion.conv import match_conv_module -from cactus.transpile.fusion.attention import AttentionBlockMatch -from cactus.transpile.fusion.attention import AttentionMatch -from cactus.transpile.fusion.attention import SelfAttentionBlockMatch -from cactus.transpile.fusion.attention import match_attention -from cactus.transpile.fusion.attention import match_attention_block -from cactus.transpile.fusion.attention import match_self_attention_block -from cactus.transpile.fusion.deltanet import GatedDeltaNetMatch -from cactus.transpile.fusion.deltanet import match_gated_deltanet -from cactus.transpile.fusion.linear import LinearMatch -from cactus.transpile.fusion.linear import match_linear -from cactus.transpile.fusion.lstm import LSTMCellMatch -from cactus.transpile.fusion.lstm import match_lstm_cell -from cactus.transpile.fusion.mlp import GatedMLPMatch -from cactus.transpile.fusion.mlp import match_gated_mlp -from cactus.transpile.fusion.rel_pos_bias import RelPosBiasMatch -from cactus.transpile.fusion.rel_pos_bias import match_rel_pos_bias -from cactus.transpile.fusion.rms_norm import RMSNormMatch -from cactus.transpile.fusion.rms_norm import match_rms_norm -from cactus.transpile.fusion.rope import RoPEMatch -from cactus.transpile.fusion.rope import match_rope - -__all__ = [ - "ConvModuleMatch", - "AttentionBlockMatch", - "AttentionMatch", - "SelfAttentionBlockMatch", - "GatedDeltaNetMatch", - "GatedMLPMatch", - "LinearMatch", - "LSTMCellMatch", - "RelPosBiasMatch", - "RMSNormMatch", - "RoPEMatch", - "match_attention", - "match_attention_block", - "match_self_attention_block", - "match_conv_module", - "match_gated_deltanet", - "match_gated_mlp", - "match_linear", - "match_lstm_cell", - "match_rel_pos_bias", - "match_rms_norm", - "match_rope", -] diff --git a/python/cactus/transpile/fusion/attention.py b/python/cactus/transpile/fusion/attention.py deleted file mode 100644 index d6a8e14e0..000000000 --- a/python/cactus/transpile/fusion/attention.py +++ /dev/null @@ -1,550 +0,0 @@ -from __future__ import annotations - -from dataclasses import dataclass - -from cactus.transpile.graph_ir import IRGraph -from cactus.transpile.graph_ir import IRNode -from cactus.transpile.fusion.common import collect_node_ids -from cactus.transpile.fusion.common import producer -from cactus.transpile.fusion.common import strip_passthrough -from cactus.transpile.fusion.linear import match_linear -from cactus.transpile.fusion.rms_norm import match_rms_norm -from cactus.transpile.fusion.rope import match_rope - - -@dataclass(frozen=True) -class AttentionMatch: - query_value_id: str - key_value_id: str - value_value_id: str - source_input_value_ids: tuple[str, str, str] - weight_value_ids: tuple[str | None, str | None, str | None] - has_rope: bool - has_qk_norm: bool - has_gqa_repeat: bool - is_causal: bool - scale: float - window_size: int - node_ids: tuple[str, ...] - - -@dataclass(frozen=True) -class AttentionBlockMatch: - attention_node_id: str - query_value_id: str - key_value_id: str - value_value_id: str - mask_value_id: str | None - gate_value_id: str | None - output_projection_weight_value_id: str - output_projection_bias_value_id: str | None - attention_output_shape: tuple[int, ...] - is_causal: bool - additive_mask: bool - scale: float - window_size: int - node_ids: tuple[str, ...] - - -@dataclass(frozen=True) -class SelfAttentionBlockMatch: - attention_node_id: str - hidden_value_id: str - query_weight_value_id: str - query_projection_bias_value_id: str | None - query_add_value_id: str | None - rel_query_add_value_id: str | None - key_weight_value_id: str - key_projection_bias_value_id: str | None - value_weight_value_id: str - value_projection_bias_value_id: str | None - mask_value_id: str | None - relative_key_input_value_id: str | None - relative_key_weight_value_id: str | None - relative_key_projection_bias_value_id: str | None - gate_value_id: str | None - output_projection_weight_value_id: str - output_projection_bias_value_id: str | None - query_shape: tuple[int, ...] - key_shape: tuple[int, ...] - value_shape: tuple[int, ...] - relative_key_shape: tuple[int, ...] | None - attention_output_shape: tuple[int, ...] - is_causal: bool - additive_mask: bool - scale: float - rel_pos_scale: float | None - window_size: int - node_ids: tuple[str, ...] - - -@dataclass(frozen=True) -class _ProjectedAttentionInputMatch: - hidden_value_id: str - weight_value_id: str - projection_bias_value_id: str | None - projected_shape: tuple[int, ...] - add_value_id: str | None - node_ids: tuple[str, ...] - - -@dataclass(frozen=True) -class _RelativeKeyProjectionMatch: - input_value_id: str - weight_value_id: str - projection_bias_value_id: str | None - projected_shape: tuple[int, ...] - node_ids: tuple[str, ...] - - -def match_attention(graph: IRGraph, node: IRNode) -> AttentionMatch | None: - if node.op not in {"attention", "scaled_dot_product_attention"} or len(node.inputs) < 3: - return None - - q_info = _extract_attention_input(graph, node.inputs[0], role="q") - k_info = _extract_attention_input(graph, node.inputs[1], role="k") - v_info = _extract_attention_input(graph, node.inputs[2], role="v") - if q_info is None or k_info is None or v_info is None: - return None - - has_rope = bool(q_info["has_rope"] and k_info["has_rope"]) - has_qk_norm = bool(q_info["has_rms_norm"] and k_info["has_rms_norm"]) - has_gqa = bool(k_info["has_gqa_repeat"] or v_info["has_gqa_repeat"] or node.attrs.get("enable_gqa", False)) - node_ids = {node.id, *q_info["node_ids"], *k_info["node_ids"], *v_info["node_ids"]} - - return AttentionMatch( - query_value_id=node.inputs[0], - key_value_id=node.inputs[1], - value_value_id=node.inputs[2], - source_input_value_ids=(q_info["source_input"], k_info["source_input"], v_info["source_input"]), - weight_value_ids=(q_info.get("weight_value_id"), k_info.get("weight_value_id"), v_info.get("weight_value_id")), - has_rope=has_rope, - has_qk_norm=has_qk_norm, - has_gqa_repeat=has_gqa, - is_causal=bool(node.attrs.get("is_causal", False)), - scale=float(node.attrs.get("scale", 0.0)), - window_size=int(node.attrs.get("window_size", 0)), - node_ids=tuple(sorted(node_ids)), - ) - - -def match_self_attention_block(graph: IRGraph, node: IRNode) -> SelfAttentionBlockMatch | None: - projection = match_linear(graph, node) - if projection is None: - return None - - output_path = _extract_attention_output_path(graph, projection.input_value_id) - if output_path is None: - return None - attention_node = graph.nodes.get(output_path["attention_node_id"]) - if attention_node is None or attention_node.op not in {"attention", "scaled_dot_product_attention"} or len(attention_node.inputs) < 3: - return None - - query_match = _extract_projected_attention_input(graph, attention_node.inputs[0], role="q") - key_match = _extract_projected_attention_input(graph, attention_node.inputs[1], role="k") - value_match = _extract_projected_attention_input(graph, attention_node.inputs[2], role="v") - if query_match is None or key_match is None or value_match is None: - return None - - if not ( - query_match.hidden_value_id == key_match.hidden_value_id == value_match.hidden_value_id - ): - return None - - mask_value_id = attention_node.inputs[3] if len(attention_node.inputs) > 3 else None - rel_query_add_value_id: str | None = None - relative_key_input_value_id: str | None = None - relative_key_weight_value_id: str | None = None - relative_key_projection_bias_value_id: str | None = None - relative_key_shape: tuple[int, ...] | None = None - rel_pos_scale: float | None = None - rel_pos_node_ids: tuple[str, ...] = () - - if mask_value_id is not None: - rel_pos_node = producer(graph, strip_passthrough(graph, mask_value_id)) - if rel_pos_node is not None and rel_pos_node.op == "rel_pos_bias" and len(rel_pos_node.inputs) == 2: - rel_query_match = _extract_projected_attention_input(graph, rel_pos_node.inputs[0], role="q") - relative_key_match = _extract_relative_key_projection(graph, rel_pos_node.inputs[1]) - if rel_query_match is None or relative_key_match is None: - return None - if not ( - rel_query_match.hidden_value_id == query_match.hidden_value_id - and rel_query_match.weight_value_id == query_match.weight_value_id - and rel_query_match.projection_bias_value_id == query_match.projection_bias_value_id - and rel_query_match.projected_shape == query_match.projected_shape - ): - return None - rel_query_add_value_id = rel_query_match.add_value_id - relative_key_input_value_id = relative_key_match.input_value_id - relative_key_weight_value_id = relative_key_match.weight_value_id - relative_key_projection_bias_value_id = relative_key_match.projection_bias_value_id - relative_key_shape = relative_key_match.projected_shape - rel_pos_scale = float(rel_pos_node.attrs.get("scale", 1.0)) - rel_pos_node_ids = collect_node_ids( - rel_pos_node.id, - rel_pos_node.meta.get("rel_pos_bias_nodes", ()), - rel_query_match.node_ids, - relative_key_match.node_ids, - ) - mask_value_id = None - - return SelfAttentionBlockMatch( - attention_node_id=attention_node.id, - hidden_value_id=query_match.hidden_value_id, - query_weight_value_id=query_match.weight_value_id, - query_projection_bias_value_id=query_match.projection_bias_value_id, - query_add_value_id=query_match.add_value_id, - rel_query_add_value_id=rel_query_add_value_id, - key_weight_value_id=key_match.weight_value_id, - key_projection_bias_value_id=key_match.projection_bias_value_id, - value_weight_value_id=value_match.weight_value_id, - value_projection_bias_value_id=value_match.projection_bias_value_id, - mask_value_id=mask_value_id, - relative_key_input_value_id=relative_key_input_value_id, - relative_key_weight_value_id=relative_key_weight_value_id, - relative_key_projection_bias_value_id=relative_key_projection_bias_value_id, - gate_value_id=output_path["gate_value_id"], - output_projection_weight_value_id=projection.weight_value_id, - output_projection_bias_value_id=projection.bias_value_id, - query_shape=query_match.projected_shape, - key_shape=key_match.projected_shape, - value_shape=value_match.projected_shape, - relative_key_shape=relative_key_shape, - attention_output_shape=output_path["attention_output_shape"], - is_causal=bool(attention_node.attrs.get("is_causal", True)), - additive_mask=bool(attention_node.attrs.get("additive_mask", False)), - scale=float(attention_node.attrs.get("scale", 0.0)), - rel_pos_scale=rel_pos_scale, - window_size=int(attention_node.attrs.get("window_size", 0)), - node_ids=collect_node_ids( - attention_node.id, - projection.node_ids, - output_path["node_ids"], - query_match.node_ids, - key_match.node_ids, - value_match.node_ids, - rel_pos_node_ids, - ), - ) - - -def match_attention_block(graph: IRGraph, node: IRNode) -> AttentionBlockMatch | None: - projection = match_linear(graph, node) - if projection is None: - return None - - output_path = _extract_attention_output_path(graph, projection.input_value_id) - if output_path is None: - return None - attention_node = graph.nodes.get(output_path["attention_node_id"]) - if attention_node is None or attention_node.op not in {"attention", "scaled_dot_product_attention"} or len(attention_node.inputs) < 3: - return None - - return AttentionBlockMatch( - attention_node_id=attention_node.id, - query_value_id=attention_node.inputs[0], - key_value_id=attention_node.inputs[1], - value_value_id=attention_node.inputs[2], - mask_value_id=attention_node.inputs[3] if len(attention_node.inputs) > 3 else None, - gate_value_id=output_path["gate_value_id"], - output_projection_weight_value_id=projection.weight_value_id, - output_projection_bias_value_id=projection.bias_value_id, - attention_output_shape=output_path["attention_output_shape"], - is_causal=bool(attention_node.attrs.get("is_causal", True)), - additive_mask=bool(attention_node.attrs.get("additive_mask", False)), - scale=float(attention_node.attrs.get("scale", 0.0)), - window_size=int(attention_node.attrs.get("window_size", 0)), - node_ids=tuple(sorted({attention_node.id, *projection.node_ids, *output_path["node_ids"]})), - ) - - -def _extract_attention_output_path(graph: IRGraph, value_id: str) -> dict[str, object] | None: - current = value_id - gate_value_id: str | None = None - - while True: - node = producer(graph, current) - if node is None: - return None - if node.op in {"precision_cast", "contiguous"} and len(node.inputs) == 1: - current = node.inputs[0] - continue - if node.op == "type_as" and len(node.inputs) >= 1: - current = node.inputs[0] - continue - break - - node = producer(graph, current) - if node is not None and node.op == "multiply" and len(node.inputs) == 2: - for attn_candidate, gate_candidate in ((node.inputs[0], node.inputs[1]), (node.inputs[1], node.inputs[0])): - attn_match = _extract_attention_output_path(graph, attn_candidate) - if attn_match is None: - continue - attn_match["gate_value_id"] = gate_candidate - attn_match["node_ids"] = tuple(sorted({node.id, *attn_match["node_ids"]})) - return attn_match - - reshape_node = producer(graph, current) - if reshape_node is None or reshape_node.op not in {"reshape", "view"} or len(reshape_node.inputs) != 1: - return None - - attention_output_shape = tuple(int(v) for v in reshape_node.attrs.get("shape", ())) - current = reshape_node.inputs[0] - transpose_node = producer(graph, current) - if transpose_node is None: - return None - - if transpose_node.op == "transpose": - if int(transpose_node.attrs.get("dim0", -1)) != 1 or int(transpose_node.attrs.get("dim1", -1)) != 2: - return None - current = transpose_node.inputs[0] - elif transpose_node.op == "permute": - permutation = tuple(int(v) for v in transpose_node.attrs.get("permutation", ())) - if permutation != (0, 2, 1, 3): - return None - current = transpose_node.inputs[0] - else: - return None - - attention_node = producer(graph, current) - if attention_node is None or attention_node.op not in {"attention", "scaled_dot_product_attention"}: - return None - - return { - "attention_node_id": attention_node.id, - "gate_value_id": gate_value_id, - "attention_output_shape": attention_output_shape, - "node_ids": tuple(sorted({reshape_node.id, transpose_node.id})), - } - - -def _extract_projected_attention_input(graph: IRGraph, value_id: str, *, role: str) -> _ProjectedAttentionInputMatch | None: - current = strip_passthrough(graph, value_id) - add_node = producer(graph, current) - if add_node is not None and add_node.op == "add" and len(add_node.inputs) == 2: - for projected_candidate, add_candidate in ((add_node.inputs[0], add_node.inputs[1]), (add_node.inputs[1], add_node.inputs[0])): - projected = _extract_transposed_projection(graph, projected_candidate) - if projected is None: - continue - if not _looks_like_attention_addend(graph, add_candidate, projected.projected_shape): - continue - return _ProjectedAttentionInputMatch( - hidden_value_id=projected.hidden_value_id, - weight_value_id=projected.weight_value_id, - projection_bias_value_id=projected.projection_bias_value_id, - projected_shape=projected.projected_shape, - add_value_id=strip_passthrough(graph, add_candidate), - node_ids=collect_node_ids(add_node.id, projected.node_ids), - ) - - if role == "q": - return None - - return _extract_transposed_projection(graph, current) - - -def _extract_transposed_projection(graph: IRGraph, value_id: str) -> _ProjectedAttentionInputMatch | None: - current = strip_passthrough(graph, value_id) - layout_node = producer(graph, current) - if layout_node is None or len(layout_node.inputs) != 1: - return None - - if layout_node.op == "permute": - permutation = tuple(int(v) for v in layout_node.attrs.get("permutation", ())) - if permutation != (0, 2, 1, 3): - return None - elif layout_node.op == "transpose": - if int(layout_node.attrs.get("dim0", -1)) != 1 or int(layout_node.attrs.get("dim1", -1)) != 2: - return None - else: - return None - - projected_value_id = strip_passthrough(graph, layout_node.inputs[0]) - projected_value = graph.values.get(projected_value_id) - if projected_value is None or projected_value.shape is None or len(projected_value.shape) != 4: - return None - - reshape_node = producer(graph, projected_value_id) - if reshape_node is None or reshape_node.op not in {"reshape", "view"} or len(reshape_node.inputs) != 1: - return None - - linear_node = producer(graph, reshape_node.inputs[0]) - if linear_node is None: - return None - linear = match_linear(graph, linear_node) - if linear is None: - return None - - return _ProjectedAttentionInputMatch( - hidden_value_id=linear.input_value_id, - weight_value_id=linear.weight_value_id, - projection_bias_value_id=linear.bias_value_id, - projected_shape=tuple(int(v) for v in projected_value.shape), - add_value_id=None, - node_ids=collect_node_ids(layout_node.id, reshape_node.id, linear.node_ids), - ) - - -def _extract_relative_key_projection(graph: IRGraph, value_id: str) -> _RelativeKeyProjectionMatch | None: - current = strip_passthrough(graph, value_id) - projected_value = graph.values.get(current) - if projected_value is None or projected_value.shape is None or len(projected_value.shape) != 4: - return None - - reshape_node = producer(graph, current) - if reshape_node is None or reshape_node.op not in {"reshape", "view"} or len(reshape_node.inputs) != 1: - return None - - linear_node = producer(graph, reshape_node.inputs[0]) - if linear_node is None: - return None - linear = match_linear(graph, linear_node) - if linear is None: - return None - - return _RelativeKeyProjectionMatch( - input_value_id=linear.input_value_id, - weight_value_id=linear.weight_value_id, - projection_bias_value_id=linear.bias_value_id, - projected_shape=tuple(int(v) for v in projected_value.shape), - node_ids=collect_node_ids(reshape_node.id, linear.node_ids), - ) - - -def _looks_like_attention_addend(graph: IRGraph, value_id: str, target_shape: tuple[int, ...]) -> bool: - value = graph.values.get(strip_passthrough(graph, value_id)) - if value is None or value.shape is None: - return False - - shape = tuple(int(v) for v in value.shape) - if len(shape) != 4 or len(target_shape) != 4: - return False - - batch, seq_len, heads, head_dim = target_shape - if shape[0] not in (1, batch) or shape[3] != head_dim: - return False - - # Allow either BHSD-style tensors like [B, H, 1, D] or BSHD-style tensors like [B, 1, H, D]. - if shape[1] == heads and shape[2] in (1, seq_len): - return True - if shape[1] in (1, seq_len) and shape[2] == heads: - return True - return False - - -def _extract_attention_input(graph: IRGraph, value_id: str, *, role: str) -> dict[str, object] | None: - node_ids: set[str] = set() - current = strip_passthrough(graph, value_id) - has_gqa_repeat = False - - while True: - node = producer(graph, current) - if node is None: - return { - "source_input": current, - "weight_value_id": None, - "has_rope": False, - "has_rms_norm": False, - "has_gqa_repeat": has_gqa_repeat, - "node_ids": tuple(sorted(node_ids)), - } - node_ids.add(node.id) - - if node.op in {"reshape", "view", "transpose", "permute"}: - current = node.inputs[0] - continue - - rope = match_rope(graph, current) - if rope is not None: - rope_info = _extract_attention_input(graph, rope.input_value_id, role=role) - if rope_info is None: - return None - rope_info["has_rope"] = True - rope_info["node_ids"] = tuple(sorted(set(rope_info["node_ids"]) | node_ids | set(rope.node_ids))) - rope_info["has_gqa_repeat"] = bool(rope_info["has_gqa_repeat"] or has_gqa_repeat) - return rope_info - - rms = match_rms_norm(graph, node) - if rms is not None: - if role == "v": - return None - rms_info = _extract_attention_input(graph, rms.input_value_id, role=role) - if rms_info is None: - return None - rms_info["has_rms_norm"] = True - rms_info["node_ids"] = tuple(sorted(set(rms_info["node_ids"]) | node_ids | set(rms.node_ids))) - rms_info["has_gqa_repeat"] = bool(rms_info["has_gqa_repeat"] or has_gqa_repeat) - return rms_info - - if _looks_like_gqa_repeat(graph, current): - has_gqa_repeat = True - current = _unwrap_gqa_repeat(graph, current) - if current is None: - return None - continue - - linear_node = producer(graph, current) - if linear_node is None: - return { - "source_input": strip_passthrough(graph, current), - "weight_value_id": None, - "has_rope": False, - "has_rms_norm": False, - "has_gqa_repeat": has_gqa_repeat, - "node_ids": tuple(sorted(node_ids)), - } - linear = match_linear(graph, linear_node) - if linear is None: - return { - "source_input": strip_passthrough(graph, current), - "weight_value_id": None, - "has_rope": False, - "has_rms_norm": False, - "has_gqa_repeat": has_gqa_repeat, - "node_ids": tuple(sorted(node_ids)), - } - return { - "source_input": strip_passthrough(graph, linear.input_value_id), - "weight_value_id": linear.weight_value_id, - "has_rope": False, - "has_rms_norm": False, - "has_gqa_repeat": has_gqa_repeat, - "node_ids": tuple(sorted(set(node_ids) | set(linear.node_ids))), - } - - -def _looks_like_gqa_repeat(graph: IRGraph, value_id: str) -> bool: - return _unwrap_gqa_repeat(graph, value_id) is not None - - -def _unwrap_gqa_repeat(graph: IRGraph, value_id: str) -> str | None: - current = strip_passthrough(graph, value_id) - reshape_node = producer(graph, current) - if reshape_node is None or reshape_node.op not in {"reshape", "view"}: - return None - current = reshape_node.inputs[0] - expand_node = producer(graph, current) - if expand_node is None or expand_node.op != "expand": - return None - current = expand_node.inputs[0] - while True: - node = producer(graph, current) - if node is None: - return None - if node.op == "slice": - current = node.inputs[0] - continue - if node.op != "unsqueeze": - return None - if int(node.attrs.get("dim", -999)) != 2: - return None - break - base_value_id = strip_passthrough(graph, node.inputs[0]) - base_shape = graph.values.get(base_value_id).shape if graph.values.get(base_value_id) is not None else None - target_shape = tuple(int(v) for v in expand_node.attrs.get("shape", ())) - if base_shape is None or len(base_shape) != 4 or len(target_shape) != 5: - return None - expected = (int(base_shape[0]), int(base_shape[1]), int(target_shape[2]), int(base_shape[2]), int(base_shape[3])) - if expected != target_shape: - return None - return base_value_id diff --git a/python/cactus/transpile/fusion/common.py b/python/cactus/transpile/fusion/common.py deleted file mode 100644 index d463cb804..000000000 --- a/python/cactus/transpile/fusion/common.py +++ /dev/null @@ -1,73 +0,0 @@ -from __future__ import annotations - -from dataclasses import dataclass -from typing import Any - -from cactus.transpile.graph_ir import IRGraph -from cactus.transpile.graph_ir import IRNode - - -@dataclass(frozen=True) -class FusionMatch: - kind: str - anchor_node_id: str - node_ids: tuple[str, ...] - value_ids: tuple[str, ...] - attrs: dict[str, object] - - -def producer(graph: IRGraph, value_id: str) -> IRNode | None: - value = graph.values.get(value_id) - if value is None or value.producer is None: - return None - return graph.nodes.get(value.producer) - - -def strip_passthrough(graph: IRGraph, value_id: str) -> str: - current = value_id - while True: - node = producer(graph, current) - if node is None: - return current - if node.op in {"precision_cast", "contiguous"} and len(node.inputs) == 1: - current = node.inputs[0] - continue - if node.op == "type_as" and len(node.inputs) >= 1: - current = node.inputs[0] - continue - return current - - -def strip_layout_passthrough(graph: IRGraph, value_id: str) -> str: - current = value_id - while True: - node = producer(graph, current) - if node is None: - return current - if node.op in {"precision_cast", "contiguous", "reshape", "view", "unsqueeze", "expand"} and len(node.inputs) == 1: - current = node.inputs[0] - continue - if node.op == "type_as" and len(node.inputs) >= 1: - current = node.inputs[0] - continue - return current - - -def collect_node_ids(*parts: Any) -> tuple[str, ...]: - node_ids: set[str] = set() - for part in parts: - if part is None: - continue - if isinstance(part, str): - node_ids.add(part) - continue - if isinstance(part, IRNode): - node_ids.add(part.id) - continue - if isinstance(part, (list, tuple, set)): - for item in part: - if isinstance(item, str): - node_ids.add(item) - elif isinstance(item, IRNode): - node_ids.add(item.id) - return tuple(sorted(node_ids)) diff --git a/python/cactus/transpile/fusion/conv.py b/python/cactus/transpile/fusion/conv.py deleted file mode 100644 index 265501f9a..000000000 --- a/python/cactus/transpile/fusion/conv.py +++ /dev/null @@ -1,123 +0,0 @@ -from __future__ import annotations - -from dataclasses import dataclass - -from cactus.transpile.fusion.common import collect_node_ids -from cactus.transpile.fusion.common import producer -from cactus.transpile.fusion.common import strip_passthrough -from cactus.transpile.graph_ir import IRGraph -from cactus.transpile.graph_ir import IRNode - - -@dataclass(frozen=True) -class ConvModuleMatch: - input_value_id: str - pointwise1_weight_value_id: str - pointwise1_bias_value_id: str | None - depthwise_weight_value_id: str - depthwise_bias_value_id: str | None - batch_norm_weight_value_id: str - batch_norm_bias_value_id: str - batch_norm_running_mean_value_id: str - batch_norm_running_var_value_id: str - pointwise2_weight_value_id: str - pointwise2_bias_value_id: str | None - eps: float - depthwise_kernel_size: int - depthwise_padding: int - node_ids: tuple[str, ...] - - -def match_conv_module(graph: IRGraph, node: IRNode) -> ConvModuleMatch | None: - if node.op != "permute" or len(node.inputs) != 1 or len(node.outputs) != 1: - return None - if tuple(node.attrs.get("permutation", ())) != (0, 2, 1): - return None - - pointwise2 = producer(graph, strip_passthrough(graph, node.inputs[0])) - if not _is_pointwise_conv1d(pointwise2): - return None - - silu = producer(graph, strip_passthrough(graph, pointwise2.inputs[0])) - if silu is None or silu.op != "silu" or len(silu.inputs) != 1: - return None - - batch_norm = producer(graph, strip_passthrough(graph, silu.inputs[0])) - if batch_norm is None or batch_norm.op != "batch_norm" or len(batch_norm.inputs) != 5: - return None - - depthwise = producer(graph, strip_passthrough(graph, batch_norm.inputs[0])) - if not _is_same_depthwise_conv1d(depthwise): - return None - - glu = producer(graph, strip_passthrough(graph, depthwise.inputs[0])) - if glu is None or glu.op != "glu" or len(glu.inputs) != 1: - return None - if int(glu.attrs.get("axis", -1)) not in {1, -2}: - return None - - pointwise1 = producer(graph, strip_passthrough(graph, glu.inputs[0])) - if not _is_pointwise_conv1d(pointwise1): - return None - - input_permute = producer(graph, strip_passthrough(graph, pointwise1.inputs[0])) - if input_permute is None or input_permute.op != "permute" or len(input_permute.inputs) != 1: - return None - if tuple(input_permute.attrs.get("permutation", ())) != (0, 2, 1): - return None - - depthwise_weight_shape = graph.values.get(depthwise.inputs[1]).shape if depthwise.inputs[1] in graph.values else None - if depthwise_weight_shape is None or len(depthwise_weight_shape) != 3: - return None - depthwise_kernel_size = int(depthwise_weight_shape[2]) - - return ConvModuleMatch( - input_value_id=strip_passthrough(graph, input_permute.inputs[0]), - pointwise1_weight_value_id=pointwise1.inputs[1], - pointwise1_bias_value_id=pointwise1.inputs[2] if len(pointwise1.inputs) > 2 else None, - depthwise_weight_value_id=depthwise.inputs[1], - depthwise_bias_value_id=depthwise.inputs[2] if len(depthwise.inputs) > 2 else None, - batch_norm_weight_value_id=batch_norm.inputs[1], - batch_norm_bias_value_id=batch_norm.inputs[2], - batch_norm_running_mean_value_id=batch_norm.inputs[3], - batch_norm_running_var_value_id=batch_norm.inputs[4], - pointwise2_weight_value_id=pointwise2.inputs[1], - pointwise2_bias_value_id=pointwise2.inputs[2] if len(pointwise2.inputs) > 2 else None, - eps=float(batch_norm.attrs.get("eps", 1e-5)), - depthwise_kernel_size=depthwise_kernel_size, - depthwise_padding=int(depthwise.attrs.get("padding", 0)), - node_ids=collect_node_ids( - node, - pointwise2, - silu, - batch_norm, - depthwise, - glu, - pointwise1, - input_permute, - ), - ) - - -def _is_pointwise_conv1d(node: IRNode | None) -> bool: - return ( - node is not None - and node.op == "conv1d" - and len(node.inputs) >= 2 - and int(node.attrs.get("stride", 1)) == 1 - and int(node.attrs.get("padding", 0)) == 0 - and int(node.attrs.get("dilation", 1)) == 1 - and int(node.attrs.get("groups", 1)) == 1 - ) - - -def _is_same_depthwise_conv1d(node: IRNode | None) -> bool: - if node is None or node.op != "conv1d" or len(node.inputs) < 2: - return False - groups = int(node.attrs.get("groups", 1)) - padding = int(node.attrs.get("padding", 0)) - stride = int(node.attrs.get("stride", 1)) - dilation = int(node.attrs.get("dilation", 1)) - if stride != 1 or dilation != 1: - return False - return groups > 1 and padding >= 0 diff --git a/python/cactus/transpile/fusion/deltanet.py b/python/cactus/transpile/fusion/deltanet.py deleted file mode 100644 index 94f4e9299..000000000 --- a/python/cactus/transpile/fusion/deltanet.py +++ /dev/null @@ -1,357 +0,0 @@ -from __future__ import annotations - -from dataclasses import dataclass - -from cactus.transpile.fusion.common import producer -from cactus.transpile.fusion.common import strip_passthrough -from cactus.transpile.fusion.linear import match_linear -from cactus.transpile.fusion.rms_norm import match_rms_norm -from cactus.transpile.graph_ir import IRGraph -from cactus.transpile.graph_ir import IRNode - - -@dataclass(frozen=True) -class GatedDeltaNetMatch: - output_value_id: str - normalized_input_value_id: str - qkv_weight_value_id: str - a_weight_value_id: str - b_weight_value_id: str - norm_weight_value_id: str - z_weight_value_id: str | None - dt_bias_value_id: str | None - a_log_value_id: str | None - conv_weight_value_id: str | None - num_k_heads: int - num_v_heads: int - key_dim: int - value_dim: int - eps: float - chunk_size: int - mode: str - node_ids: tuple[str, ...] - - -def match_gated_deltanet(graph: IRGraph, node: IRNode) -> GatedDeltaNetMatch | None: - if node.op not in {"view", "reshape"} or len(node.outputs) != 1: - return None - output_value = graph.values.get(node.outputs[0]) - if output_value is None or output_value.shape is None or len(output_value.shape) != 3: - return None - - body_value_id = strip_passthrough(graph, node.inputs[0]) - body_node = producer(graph, body_value_id) - if body_node is None: - return None - - z_weight_value_id: str | None = None - rms_match = None - normalized_input_value_id: str | None = None - node_ids: set[str] = {node.id} - - if body_node.op == "multiply" and len(body_node.inputs) == 2: - for rms_value_id, z_value_id in ((body_node.inputs[0], body_node.inputs[1]), (body_node.inputs[1], body_node.inputs[0])): - rms_match = _match_rms_from_value(graph, rms_value_id) - if rms_match is None: - continue - z_weight_value_id, normalized_input_value_id, z_node_ids = _extract_z_branch(graph, z_value_id) - if normalized_input_value_id is None: - continue - node_ids.update(z_node_ids) - node_ids.add(body_node.id) - break - if rms_match is None or normalized_input_value_id is None: - return None - else: - rms_match = _match_rms_from_value(graph, body_value_id) - if rms_match is None: - return None - - y2d_value_id = rms_match.input_value_id - if rms_match.weight_value_id is None: - return None - norm_weight_value_id = rms_match.weight_value_id - eps = float(rms_match.eps) - node_ids.update(rms_match.node_ids) - - if normalized_input_value_id is None: - normalized_input_value_id = _infer_normalized_input_from_y_branch(graph, y2d_value_id) - if normalized_input_value_id is None: - return None - - qkv_info = _find_qkv_linear(graph, normalized_input_value_id) - if qkv_info is None: - return None - qkv_weight_value_id, conv_weight_value_id, num_k_heads, key_dim, num_v_heads, value_dim, qkv_node_ids = qkv_info - node_ids.update(qkv_node_ids) - - a_info = _find_a_linear(graph, normalized_input_value_id) - b_info = _find_b_linear(graph, normalized_input_value_id) - if a_info is None or b_info is None: - return None - a_weight_value_id, dt_bias_value_id, a_log_value_id, a_node_ids = a_info - b_weight_value_id, b_node_ids = b_info - node_ids.update(a_node_ids) - node_ids.update(b_node_ids) - - mode, chunk_size = _classify_mode(graph, y2d_value_id, match_node_ids=node_ids) - - return GatedDeltaNetMatch( - output_value_id=node.outputs[0], - normalized_input_value_id=normalized_input_value_id, - qkv_weight_value_id=qkv_weight_value_id, - a_weight_value_id=a_weight_value_id, - b_weight_value_id=b_weight_value_id, - norm_weight_value_id=norm_weight_value_id, - z_weight_value_id=z_weight_value_id, - dt_bias_value_id=dt_bias_value_id, - a_log_value_id=a_log_value_id, - conv_weight_value_id=conv_weight_value_id, - num_k_heads=num_k_heads, - num_v_heads=num_v_heads, - key_dim=key_dim, - value_dim=value_dim, - eps=eps, - chunk_size=chunk_size, - mode=mode, - node_ids=tuple(sorted(node_ids)), - ) - - -def _match_rms_from_value(graph: IRGraph, value_id: str): - node = producer(graph, strip_passthrough(graph, value_id)) - if node is None: - return None - return match_rms_norm(graph, node) - - -def _extract_z_branch(graph: IRGraph, value_id: str) -> tuple[str | None, str | None, tuple[str, ...]]: - current = strip_passthrough(graph, value_id) - silu_node = producer(graph, current) - if silu_node is None or silu_node.op != "silu" or len(silu_node.inputs) != 1: - return None, None, () - linear_value_id = strip_passthrough(graph, silu_node.inputs[0]) - linear_node = producer(graph, linear_value_id) - while linear_node is not None and linear_node.op in {"view", "reshape"} and len(linear_node.inputs) == 1: - linear_value_id = strip_passthrough(graph, linear_node.inputs[0]) - linear_node = producer(graph, linear_value_id) - if linear_node is None: - return None, None, () - linear_match = match_linear(graph, linear_node) - if linear_match is None: - return None, None, () - return ( - linear_match.weight_value_id, - linear_match.input_value_id, - tuple(sorted({silu_node.id, *linear_match.node_ids})), - ) - - -def _infer_normalized_input_from_y_branch(graph: IRGraph, value_id: str) -> str | None: - current = strip_passthrough(graph, value_id) - node = producer(graph, current) - visited: set[str] = set() - while node is not None and node.id not in visited: - visited.add(node.id) - for input_value_id in node.inputs: - candidate = producer(graph, strip_passthrough(graph, input_value_id)) - if candidate is None: - continue - linear_match = match_linear(graph, candidate) - if linear_match is not None: - return linear_match.input_value_id - if len(node.inputs) != 1: - return None - node = producer(graph, strip_passthrough(graph, node.inputs[0])) - return None - - -def _find_qkv_linear( - graph: IRGraph, - normalized_input_value_id: str, -) -> tuple[str, str | None, int, int, int, int, tuple[str, ...]] | None: - for node_id in graph.order: - node = graph.nodes[node_id] - linear_match = match_linear(graph, node) - if linear_match is None or linear_match.input_value_id != normalized_input_value_id: - continue - split_node = _find_descendant_op(graph, linear_match.output_value_id, {"split_with_sizes"}, max_depth=8) - if split_node is None: - continue - split_sizes = tuple(int(v) for v in split_node.attrs.get("sizes", ())) - if len(split_sizes) != 3: - continue - - conv_node = _find_descendant_op(graph, linear_match.output_value_id, {"conv1d"}, max_depth=6) - conv_weight_value_id = conv_node.inputs[1] if conv_node is not None and len(conv_node.inputs) >= 2 else None - - q_shape = None - v_shape = None - node_ids = {node.id, split_node.id} - if conv_node is not None: - node_ids.add(conv_node.id) - - split_value = graph.values.get(split_node.outputs[0]) - if split_value is None: - continue - for getitem_node_id in split_value.users: - getitem_node = graph.nodes.get(getitem_node_id) - if getitem_node is None or getitem_node.op != "getitem" or not getitem_node.outputs: - continue - shape = _find_rank4_shape(graph, getitem_node.outputs[0], max_depth=5) - if shape is None: - continue - branch_has_sum = _find_descendant_op(graph, getitem_node.outputs[0], {"sum"}, max_depth=8) is not None - node_ids.add(getitem_node.id) - if branch_has_sum and q_shape is None: - q_shape = shape - elif not branch_has_sum and v_shape is None: - v_shape = shape - - if q_shape is None or v_shape is None: - continue - - num_k_heads = int(q_shape[2]) - key_dim = int(q_shape[3]) - num_v_heads = int(v_shape[2]) - value_dim = int(v_shape[3]) - if num_k_heads <= 0 or key_dim <= 0 or num_v_heads <= 0 or value_dim <= 0: - continue - - return ( - linear_match.weight_value_id, - conv_weight_value_id, - num_k_heads, - key_dim, - num_v_heads, - value_dim, - tuple(sorted(node_ids | set(linear_match.node_ids))), - ) - return None - - -def _find_a_linear( - graph: IRGraph, - normalized_input_value_id: str, -) -> tuple[str, str | None, str | None, tuple[str, ...]] | None: - for node_id in graph.order: - node = graph.nodes[node_id] - linear_match = match_linear(graph, node) - if linear_match is None or linear_match.input_value_id != normalized_input_value_id: - continue - softplus_node = _find_descendant_op(graph, linear_match.output_value_id, {"softplus"}, max_depth=6) - if softplus_node is None: - continue - dt_bias_value_id = None - softplus_input_node = producer(graph, strip_passthrough(graph, softplus_node.inputs[0])) - if softplus_input_node is not None and softplus_input_node.op == "add": - for input_value_id in softplus_input_node.inputs: - value = graph.values.get(input_value_id) - if value is not None and value.producer is None and input_value_id in graph.constants: - dt_bias_value_id = input_value_id - break - a_log_value_id = None - multiply_node = _find_descendant_op(graph, softplus_node.outputs[0], {"multiply"}, max_depth=6) - if multiply_node is not None: - for input_value_id in multiply_node.inputs: - a_log_value_id = _extract_exp_constant_input(graph, input_value_id) - if a_log_value_id is not None: - break - node_ids = {softplus_node.id, *linear_match.node_ids} - if softplus_input_node is not None: - node_ids.add(softplus_input_node.id) - if multiply_node is not None: - node_ids.add(multiply_node.id) - return linear_match.weight_value_id, dt_bias_value_id, a_log_value_id, tuple(sorted(node_ids)) - return None - - -def _find_b_linear(graph: IRGraph, normalized_input_value_id: str) -> tuple[str, tuple[str, ...]] | None: - for node_id in graph.order: - node = graph.nodes[node_id] - linear_match = match_linear(graph, node) - if linear_match is None or linear_match.input_value_id != normalized_input_value_id: - continue - sigmoid_node = _find_descendant_op(graph, linear_match.output_value_id, {"sigmoid"}, max_depth=4) - if sigmoid_node is None: - continue - return linear_match.weight_value_id, tuple(sorted({sigmoid_node.id, *linear_match.node_ids})) - return None - - -def _extract_exp_constant_input(graph: IRGraph, value_id: str) -> str | None: - node = producer(graph, strip_passthrough(graph, value_id)) - if node is None: - value = graph.values.get(value_id) - if value is not None and value.producer is None and value_id in graph.constants: - return value_id - return None - if node.op == "scalar_multiply": - return _extract_exp_constant_input(graph, node.inputs[0]) - if node.op == "scalar_exp": - source = strip_passthrough(graph, node.inputs[0]) - value = graph.values.get(source) - if value is not None and value.producer is None and source in graph.constants: - return source - return None - - -def _find_rank4_shape(graph: IRGraph, value_id: str, *, max_depth: int) -> tuple[int, int, int, int] | None: - frontier = [(value_id, 0)] - seen: set[str] = set() - while frontier: - current, depth = frontier.pop(0) - if current in seen or depth > max_depth: - continue - seen.add(current) - value = graph.values.get(current) - if value is not None and value.shape is not None and len(value.shape) == 4: - shape = tuple(int(dim) for dim in value.shape) - return shape # type: ignore[return-value] - for user_id in graph.values.get(current, None).users if graph.values.get(current, None) is not None else (): - user = graph.nodes.get(user_id) - if user is None or not user.outputs: - continue - frontier.extend((output_id, depth + 1) for output_id in user.outputs) - return None - - -def _find_descendant_op(graph: IRGraph, value_id: str, targets: set[str], *, max_depth: int) -> IRNode | None: - frontier = [(value_id, 0)] - seen: set[str] = set() - while frontier: - current, depth = frontier.pop(0) - if current in seen or depth > max_depth: - continue - seen.add(current) - value = graph.values.get(current) - if value is None: - continue - for user_id in value.users: - user = graph.nodes.get(user_id) - if user is None or not user.outputs: - continue - if user.op in targets: - return user - frontier.extend((output_id, depth + 1) for output_id in user.outputs) - return None - - -def _classify_mode(graph: IRGraph, y2d_value_id: str, *, match_node_ids: set[str]) -> tuple[str, int]: - del y2d_value_id - chunk_size = _infer_prefill_chunk_size(graph, match_node_ids) - # Current exported graphs are uncached/full-sequence captures, so fuse the - # linear-attention block to the prefill kernel by default. Decode should be - # emitted only once the capture path models cache/state explicitly. - return "prefill", chunk_size - - -def _infer_prefill_chunk_size(graph: IRGraph, match_node_ids: set[str]) -> int: - for node_id in match_node_ids: - node = graph.nodes.get(node_id) - if node is None or node.op != "aten.triu.default" or not node.inputs: - continue - ones_value = graph.values.get(node.inputs[0]) - if ones_value is not None and ones_value.shape and len(ones_value.shape) >= 1: - return int(ones_value.shape[-1]) - return 64 diff --git a/python/cactus/transpile/fusion/linear.py b/python/cactus/transpile/fusion/linear.py deleted file mode 100644 index 4841019f3..000000000 --- a/python/cactus/transpile/fusion/linear.py +++ /dev/null @@ -1,58 +0,0 @@ -from __future__ import annotations - -from dataclasses import dataclass - -from cactus.transpile.graph_ir import IRGraph -from cactus.transpile.graph_ir import IRNode -from cactus.transpile.fusion.common import producer -from cactus.transpile.fusion.common import strip_passthrough - - -@dataclass(frozen=True) -class LinearMatch: - input_value_id: str - weight_value_id: str - bias_value_id: str | None - output_value_id: str - node_ids: tuple[str, ...] - kind: str - - -def match_linear(graph: IRGraph, node: IRNode) -> LinearMatch | None: - if node.op == "linear" and len(node.inputs) >= 2 and len(node.outputs) == 1: - return LinearMatch( - input_value_id=strip_passthrough(graph, node.inputs[0]), - weight_value_id=node.inputs[1], - bias_value_id=node.inputs[2] if bool(node.attrs.get("has_bias", False)) and len(node.inputs) > 2 else None, - output_value_id=node.outputs[0], - node_ids=(node.id,), - kind="linear", - ) - - if node.op == "addmm" and len(node.inputs) == 3 and len(node.outputs) == 1: - return LinearMatch( - input_value_id=strip_passthrough(graph, node.inputs[1]), - weight_value_id=node.inputs[2], - bias_value_id=node.inputs[0], - output_value_id=node.outputs[0], - node_ids=(node.id,), - kind="addmm", - ) - - if node.op != "add" or len(node.inputs) != 2 or len(node.outputs) != 1: - return None - - for matmul_value_id, bias_value_id in ((node.inputs[0], node.inputs[1]), (node.inputs[1], node.inputs[0])): - matmul_node = producer(graph, strip_passthrough(graph, matmul_value_id)) - if matmul_node is None or matmul_node.op != "matmul" or len(matmul_node.inputs) != 2: - continue - return LinearMatch( - input_value_id=strip_passthrough(graph, matmul_node.inputs[0]), - weight_value_id=matmul_node.inputs[1], - bias_value_id=bias_value_id, - output_value_id=node.outputs[0], - node_ids=tuple(sorted({node.id, matmul_node.id})), - kind="matmul_add", - ) - - return None diff --git a/python/cactus/transpile/fusion/lstm.py b/python/cactus/transpile/fusion/lstm.py deleted file mode 100644 index 64831240e..000000000 --- a/python/cactus/transpile/fusion/lstm.py +++ /dev/null @@ -1,216 +0,0 @@ -from __future__ import annotations - -from dataclasses import dataclass - -from cactus.transpile.fusion.common import collect_node_ids -from cactus.transpile.fusion.common import producer -from cactus.transpile.fusion.common import strip_passthrough -from cactus.transpile.graph_ir import IRGraph -from cactus.transpile.graph_ir import IRNode - - -@dataclass(frozen=True) -class LSTMCellMatch: - anchor_node_id: str - node_ids: tuple[str, ...] - x_value_id: str - h_prev_value_id: str - c_prev_value_id: str - weight_ih_value_id: str - weight_hh_value_id: str - bias_ih_value_id: str - bias_hh_value_id: str | None - h_output_value_id: str - c_output_value_id: str - - -def match_lstm_cell(graph: IRGraph, node: IRNode) -> LSTMCellMatch | None: - if node.op != "multiply" or len(node.inputs) != 2 or len(node.outputs) != 1: - return None - - for output_gate_value_id, c_tanh_value_id in ( - (node.inputs[0], node.inputs[1]), - (node.inputs[1], node.inputs[0]), - ): - output_gate = _extract_gate_activation(graph, output_gate_value_id, expected_activation="sigmoid") - if output_gate is None or output_gate["index"] != 3: - continue - - c_tanh_node = producer(graph, strip_passthrough(graph, c_tanh_value_id)) - if c_tanh_node is None or c_tanh_node.op != "tanh" or len(c_tanh_node.inputs) != 1: - continue - - c_next_value_id = strip_passthrough(graph, c_tanh_node.inputs[0]) - c_next_node = producer(graph, c_next_value_id) - if c_next_node is None or c_next_node.op != "add" or len(c_next_node.inputs) != 2 or len(c_next_node.outputs) != 1: - continue - - for forget_term_value_id, input_term_value_id in ( - (c_next_node.inputs[0], c_next_node.inputs[1]), - (c_next_node.inputs[1], c_next_node.inputs[0]), - ): - forget_term = _extract_forget_term(graph, forget_term_value_id) - input_term = _extract_input_term(graph, input_term_value_id) - if forget_term is None or input_term is None: - continue - - input_gate = input_term["input_gate"] - cell_gate = input_term["cell_gate"] - forget_gate = forget_term["forget_gate"] - chunk_node_id = output_gate["chunk_node_id"] - if ( - input_gate["chunk_node_id"] != chunk_node_id - or forget_gate["chunk_node_id"] != chunk_node_id - or cell_gate["chunk_node_id"] != chunk_node_id - ): - continue - if input_gate["index"] != 0 or forget_gate["index"] != 1 or cell_gate["index"] != 2: - continue - - chunk_node = graph.nodes.get(chunk_node_id) - if chunk_node is None or len(chunk_node.inputs) != 1: - continue - - gates = _extract_lstm_gate_sum(graph, chunk_node.inputs[0]) - if gates is None: - continue - - return LSTMCellMatch( - anchor_node_id=node.id, - node_ids=collect_node_ids( - node, - c_tanh_node, - c_next_node, - output_gate["node_ids"], - forget_gate["node_ids"], - input_gate["node_ids"], - cell_gate["node_ids"], - forget_term["node_ids"], - input_term["node_ids"], - chunk_node, - gates["node_ids"], - ), - x_value_id=gates["x_value_id"], - h_prev_value_id=gates["h_prev_value_id"], - c_prev_value_id=forget_term["c_prev_value_id"], - weight_ih_value_id=gates["weight_ih_value_id"], - weight_hh_value_id=gates["weight_hh_value_id"], - bias_ih_value_id=gates["bias_value_ids"][0], - bias_hh_value_id=gates["bias_value_ids"][1] if len(gates["bias_value_ids"]) > 1 else None, - h_output_value_id=node.outputs[0], - c_output_value_id=c_next_node.outputs[0], - ) - - return None - - -def _extract_gate_activation( - graph: IRGraph, - value_id: str, - *, - expected_activation: str, -) -> dict[str, object] | None: - activation_node = producer(graph, strip_passthrough(graph, value_id)) - if activation_node is None or activation_node.op != expected_activation or len(activation_node.inputs) != 1: - return None - - getitem_node = producer(graph, strip_passthrough(graph, activation_node.inputs[0])) - if getitem_node is None or getitem_node.op != "getitem" or len(getitem_node.inputs) != 1: - return None - - chunk_node = producer(graph, strip_passthrough(graph, getitem_node.inputs[0])) - if chunk_node is None or chunk_node.op != "chunk": - return None - if int(chunk_node.attrs.get("chunks", 0)) != 4: - return None - - return { - "chunk_node_id": chunk_node.id, - "index": int(getitem_node.attrs.get("index", -1)), - "preactivation_value_id": getitem_node.inputs[0], - "node_ids": (activation_node.id, getitem_node.id), - } - - -def _extract_forget_term(graph: IRGraph, value_id: str) -> dict[str, object] | None: - mul_node = producer(graph, strip_passthrough(graph, value_id)) - if mul_node is None or mul_node.op != "multiply" or len(mul_node.inputs) != 2: - return None - - for gate_value_id, c_prev_value_id in ( - (mul_node.inputs[0], mul_node.inputs[1]), - (mul_node.inputs[1], mul_node.inputs[0]), - ): - forget_gate = _extract_gate_activation(graph, gate_value_id, expected_activation="sigmoid") - if forget_gate is None: - continue - return { - "forget_gate": forget_gate, - "c_prev_value_id": strip_passthrough(graph, c_prev_value_id), - "node_ids": (mul_node.id,), - } - return None - - -def _extract_input_term(graph: IRGraph, value_id: str) -> dict[str, object] | None: - mul_node = producer(graph, strip_passthrough(graph, value_id)) - if mul_node is None or mul_node.op != "multiply" or len(mul_node.inputs) != 2: - return None - - for input_gate_value_id, cell_gate_value_id in ( - (mul_node.inputs[0], mul_node.inputs[1]), - (mul_node.inputs[1], mul_node.inputs[0]), - ): - input_gate = _extract_gate_activation(graph, input_gate_value_id, expected_activation="sigmoid") - if input_gate is None: - continue - cell_tanh_node = producer(graph, strip_passthrough(graph, cell_gate_value_id)) - if cell_tanh_node is None or cell_tanh_node.op != "tanh" or len(cell_tanh_node.inputs) != 1: - continue - cell_gate = _extract_gate_activation(graph, cell_gate_value_id, expected_activation="tanh") - if cell_gate is None: - continue - return { - "input_gate": input_gate, - "cell_gate": cell_gate, - "node_ids": (mul_node.id, cell_tanh_node.id), - } - return None - - -def _extract_lstm_gate_sum(graph: IRGraph, value_id: str) -> dict[str, object] | None: - terms, node_ids = _flatten_add_terms(graph, strip_passthrough(graph, value_id)) - linear_nodes: list[IRNode] = [] - bias_value_ids: list[str] = [] - for term_value_id in terms: - term_value_id = strip_passthrough(graph, term_value_id) - term_node = producer(graph, term_value_id) - if term_node is not None and term_node.op == "linear" and len(term_node.inputs) >= 2: - linear_nodes.append(term_node) - continue - if term_value_id in graph.constants: - bias_value_ids.append(term_value_id) - continue - return None - - if len(linear_nodes) != 2 or len(bias_value_ids) not in {1, 2}: - return None - - first_linear, second_linear = linear_nodes - return { - "x_value_id": first_linear.inputs[0], - "weight_ih_value_id": first_linear.inputs[1], - "h_prev_value_id": second_linear.inputs[0], - "weight_hh_value_id": second_linear.inputs[1], - "bias_value_ids": tuple(bias_value_ids), - "node_ids": (*node_ids, first_linear.id, second_linear.id), - } - - -def _flatten_add_terms(graph: IRGraph, value_id: str) -> tuple[list[str], tuple[str, ...]]: - node = producer(graph, strip_passthrough(graph, value_id)) - if node is None or node.op != "add" or len(node.inputs) != 2: - return [value_id], () - left_terms, left_nodes = _flatten_add_terms(graph, node.inputs[0]) - right_terms, right_nodes = _flatten_add_terms(graph, node.inputs[1]) - return [*left_terms, *right_terms], (*left_nodes, *right_nodes, node.id) diff --git a/python/cactus/transpile/fusion/mlp.py b/python/cactus/transpile/fusion/mlp.py deleted file mode 100644 index 56a219514..000000000 --- a/python/cactus/transpile/fusion/mlp.py +++ /dev/null @@ -1,53 +0,0 @@ -from __future__ import annotations - -from dataclasses import dataclass - -from cactus.transpile.graph_ir import IRGraph -from cactus.transpile.graph_ir import IRNode -from cactus.transpile.fusion.common import strip_passthrough -from cactus.transpile.fusion.linear import match_linear - - -@dataclass(frozen=True) -class GatedMLPMatch: - input_value_id: str - gate_weight_value_id: str - up_weight_value_id: str - down_weight_value_id: str - activation: str - bias_value_ids: tuple[str | None, str | None, str | None] - node_ids: tuple[str, ...] - - -def match_gated_mlp(graph: IRGraph, node: IRNode) -> GatedMLPMatch | None: - down = match_linear(graph, node) - if down is None: - return None - - mul_node = graph.nodes.get(graph.values[down.input_value_id].producer) if graph.values.get(down.input_value_id) is not None else None - if mul_node is None or mul_node.op != "multiply" or len(mul_node.inputs) != 2: - return None - - for activated_value_id, up_value_id in ((mul_node.inputs[0], mul_node.inputs[1]), (mul_node.inputs[1], mul_node.inputs[0])): - activated_node = graph.nodes.get(graph.values[strip_passthrough(graph, activated_value_id)].producer) if graph.values.get(strip_passthrough(graph, activated_value_id)) is not None else None - if activated_node is None or activated_node.op not in {"gelu", "silu"}: - continue - - gate = match_linear(graph, graph.nodes[graph.values[strip_passthrough(graph, activated_node.inputs[0])].producer]) if graph.values.get(strip_passthrough(graph, activated_node.inputs[0])) is not None and graph.values[strip_passthrough(graph, activated_node.inputs[0])].producer in graph.nodes else None - up = match_linear(graph, graph.nodes[graph.values[strip_passthrough(graph, up_value_id)].producer]) if graph.values.get(strip_passthrough(graph, up_value_id)) is not None and graph.values[strip_passthrough(graph, up_value_id)].producer in graph.nodes else None - if gate is None or up is None: - continue - if strip_passthrough(graph, gate.input_value_id) != strip_passthrough(graph, up.input_value_id): - continue - - return GatedMLPMatch( - input_value_id=strip_passthrough(graph, gate.input_value_id), - gate_weight_value_id=gate.weight_value_id, - up_weight_value_id=up.weight_value_id, - down_weight_value_id=down.weight_value_id, - activation=activated_node.op, - bias_value_ids=(gate.bias_value_id, up.bias_value_id, down.bias_value_id), - node_ids=tuple(sorted({*gate.node_ids, activated_node.id, *up.node_ids, mul_node.id, *down.node_ids})), - ) - - return None diff --git a/python/cactus/transpile/fusion/rel_pos_bias.py b/python/cactus/transpile/fusion/rel_pos_bias.py deleted file mode 100644 index a0e709c9c..000000000 --- a/python/cactus/transpile/fusion/rel_pos_bias.py +++ /dev/null @@ -1,179 +0,0 @@ -from __future__ import annotations - -from dataclasses import dataclass - -from cactus.transpile.fusion.common import producer -from cactus.transpile.fusion.common import strip_passthrough -from cactus.transpile.graph_ir import IRGraph -from cactus.transpile.graph_ir import IRNode - - -@dataclass(frozen=True) -class RelPosBiasMatch: - query_value_id: str - relative_key_value_id: str - scale: float - node_ids: tuple[str, ...] - - -def match_rel_pos_bias(graph: IRGraph, node: IRNode) -> RelPosBiasMatch | None: - if node.op != "scalar_multiply" or len(node.inputs) != 1: - return None - - gather_match = _match_gather_rel_pos_bias(graph, node) - if gather_match is not None: - return gather_match - - final_slice = producer(graph, strip_passthrough(graph, node.inputs[0])) - if final_slice is None or final_slice.op != "slice": - return None - if int(final_slice.attrs.get("axis", -1)) != 3 or int(final_slice.attrs.get("start", -1)) != 0: - return None - - reshape_after_shift = producer(graph, strip_passthrough(graph, final_slice.inputs[0])) - if reshape_after_shift is None or reshape_after_shift.op not in {"reshape", "view"} or len(reshape_after_shift.inputs) != 1: - return None - - shift_slice = producer(graph, strip_passthrough(graph, reshape_after_shift.inputs[0])) - if shift_slice is None or shift_slice.op != "slice": - return None - if int(shift_slice.attrs.get("axis", -1)) != 2 or int(shift_slice.attrs.get("start", -1)) != 1: - return None - - reshape_for_shift = producer(graph, strip_passthrough(graph, shift_slice.inputs[0])) - if reshape_for_shift is None or reshape_for_shift.op not in {"reshape", "view"} or len(reshape_for_shift.inputs) != 1: - return None - - pad = producer(graph, strip_passthrough(graph, reshape_for_shift.inputs[0])) - if pad is None or pad.op != "pad" or len(pad.inputs) != 1: - return None - if tuple(int(v) for v in pad.attrs.get("pads", ())) != (1, 0): - return None - if str(pad.attrs.get("mode", "constant")) != "constant": - return None - if float(pad.attrs.get("value", 0.0)) != 0.0: - return None - - matmul = producer(graph, strip_passthrough(graph, pad.inputs[0])) - if matmul is None or matmul.op != "matmul" or len(matmul.inputs) != 2: - return None - - rhs_permute = producer(graph, strip_passthrough(graph, matmul.inputs[1])) - if rhs_permute is None or rhs_permute.op != "permute" or len(rhs_permute.inputs) != 1: - return None - if tuple(int(v) for v in rhs_permute.attrs.get("permutation", ())) != (0, 2, 3, 1): - return None - - query_value_id = strip_passthrough(graph, matmul.inputs[0]) - relative_key_value_id = strip_passthrough(graph, rhs_permute.inputs[0]) - - query_shape = graph.values.get(query_value_id).shape if query_value_id in graph.values else None - relative_key_shape = graph.values.get(relative_key_value_id).shape if relative_key_value_id in graph.values else None - output_shape = graph.values.get(node.outputs[0]).shape if node.outputs and node.outputs[0] in graph.values else None - if query_shape is None or relative_key_shape is None or output_shape is None: - return None - if len(query_shape) != 4 or len(relative_key_shape) != 4 or len(output_shape) != 4: - return None - - batch, heads, seq_len, head_dim = (int(v) for v in query_shape) - rel_batch, rel_len, rel_heads, rel_dim = (int(v) for v in relative_key_shape) - out_batch, out_heads, out_t, out_s = (int(v) for v in output_shape) - if rel_batch not in {1, batch}: - return None - if rel_heads != heads or rel_dim != head_dim: - return None - if rel_len < 2 * seq_len - 1: - return None - if (out_batch, out_heads, out_t, out_s) != (batch, heads, seq_len, seq_len): - return None - - return RelPosBiasMatch( - query_value_id=query_value_id, - relative_key_value_id=relative_key_value_id, - scale=float(node.attrs.get("value", 1.0)), - node_ids=tuple( - sorted( - { - node.id, - final_slice.id, - reshape_after_shift.id, - shift_slice.id, - reshape_for_shift.id, - pad.id, - matmul.id, - rhs_permute.id, - } - ) - ), - ) - - -def _match_gather_rel_pos_bias(graph: IRGraph, node: IRNode) -> RelPosBiasMatch | None: - gather = producer(graph, strip_passthrough(graph, node.inputs[0])) - if gather is None or gather.op != "gather" or len(gather.inputs) < 2: - return None - if int(gather.attrs.get("axis", -1)) not in {-1, 3}: - return None - - matmul = producer(graph, strip_passthrough(graph, gather.inputs[0])) - if matmul is None or matmul.op != "matmul" or len(matmul.inputs) != 2: - return None - - query_heads = producer(graph, strip_passthrough(graph, matmul.inputs[0])) - rhs_transpose = producer(graph, strip_passthrough(graph, matmul.inputs[1])) - if query_heads is None or query_heads.op != "permute" or len(query_heads.inputs) != 1: - return None - if rhs_transpose is None or rhs_transpose.op != "permute" or len(rhs_transpose.inputs) != 1: - return None - if tuple(int(v) for v in query_heads.attrs.get("permutation", ())) != (0, 2, 1, 3): - return None - if tuple(int(v) for v in rhs_transpose.attrs.get("permutation", ())) != (0, 1, 3, 2): - return None - - relative_key_heads = producer(graph, strip_passthrough(graph, rhs_transpose.inputs[0])) - if relative_key_heads is None or relative_key_heads.op != "permute" or len(relative_key_heads.inputs) != 1: - return None - if tuple(int(v) for v in relative_key_heads.attrs.get("permutation", ())) != (0, 2, 1, 3): - return None - - query_value_id = strip_passthrough(graph, query_heads.inputs[0]) - relative_key_value_id = strip_passthrough(graph, relative_key_heads.inputs[0]) - output_value_id = node.outputs[0] if node.outputs else "" - - query_shape = graph.values.get(query_value_id).shape if query_value_id in graph.values else None - relative_key_shape = graph.values.get(relative_key_value_id).shape if relative_key_value_id in graph.values else None - output_shape = graph.values.get(output_value_id).shape if output_value_id in graph.values else None - if query_shape is None or relative_key_shape is None or output_shape is None: - return None - if len(query_shape) != 4 or len(relative_key_shape) != 4 or len(output_shape) != 4: - return None - - batch, seq_len, heads, head_dim = (int(v) for v in query_shape) - rel_batch, rel_len, rel_heads, rel_dim = (int(v) for v in relative_key_shape) - out_batch, out_heads, out_t, out_s = (int(v) for v in output_shape) - if rel_batch not in {1, batch}: - return None - if rel_heads != heads or rel_dim != head_dim: - return None - if rel_len < 2 * seq_len - 1: - return None - if (out_batch, out_heads, out_t, out_s) != (batch, heads, seq_len, seq_len): - return None - - return RelPosBiasMatch( - query_value_id=query_value_id, - relative_key_value_id=relative_key_value_id, - scale=float(node.attrs.get("value", 1.0)), - node_ids=tuple( - sorted( - { - node.id, - gather.id, - matmul.id, - query_heads.id, - rhs_transpose.id, - relative_key_heads.id, - } - ) - ), - ) diff --git a/python/cactus/transpile/fusion/rms_norm.py b/python/cactus/transpile/fusion/rms_norm.py deleted file mode 100644 index 322833744..000000000 --- a/python/cactus/transpile/fusion/rms_norm.py +++ /dev/null @@ -1,133 +0,0 @@ -from __future__ import annotations - -from dataclasses import dataclass - -import torch - -from cactus.transpile.graph_ir import IRGraph -from cactus.transpile.graph_ir import IRNode -from cactus.transpile.fusion.common import producer -from cactus.transpile.fusion.common import strip_passthrough - - -@dataclass(frozen=True) -class RMSNormMatch: - input_value_id: str - weight_value_id: str | None - weight_offset: float - eps: float - node_ids: tuple[str, ...] - - -def match_rms_norm(graph: IRGraph, node: IRNode) -> RMSNormMatch | None: - anchor_value_id = node.inputs[0] if node.op == "type_as" and node.inputs else node.outputs[0] - anchor_value_id = strip_passthrough(graph, anchor_value_id) - anchor_node = producer(graph, anchor_value_id) - if anchor_node is None or anchor_node.op != "multiply" or len(anchor_node.inputs) != 2: - return None - - for normed_value_id, scale_value_id in ( - (anchor_node.inputs[0], anchor_node.inputs[1]), - (anchor_node.inputs[1], anchor_node.inputs[0]), - ): - scale = _extract_scale(graph, scale_value_id) - norm = _extract_normed_value(graph, normed_value_id) - if scale is None or norm is None: - continue - if strip_passthrough(graph, norm["input_value_id"]) != strip_passthrough(graph, norm["pow_input_value_id"]): - continue - return RMSNormMatch( - input_value_id=norm["input_value_id"], - weight_value_id=scale["weight_value_id"], - weight_offset=float(scale["offset"]), - eps=float(norm["eps"]), - node_ids=tuple(sorted({anchor_node.id, *norm["node_ids"], *scale["node_ids"]})), - ) - - norm = _extract_normed_value(graph, anchor_node.outputs[0]) - if norm is None: - return None - if strip_passthrough(graph, norm["input_value_id"]) != strip_passthrough(graph, norm["pow_input_value_id"]): - return None - return RMSNormMatch( - input_value_id=norm["input_value_id"], - weight_value_id=None, - weight_offset=0.0, - eps=float(norm["eps"]), - node_ids=tuple(sorted({anchor_node.id, *norm["node_ids"]})), - ) - - -def _extract_scale(graph: IRGraph, value_id: str) -> dict[str, object] | None: - current = value_id - offset = 0.0 - node_ids: set[str] = set() - producer_node = producer(graph, current) - if producer_node is not None and producer_node.op == "scalar_add": - add_value = float(producer_node.attrs.get("value", 0.0)) - if add_value in (0.0, 1.0): - offset = add_value - node_ids.add(producer_node.id) - current = producer_node.inputs[0] - current = strip_passthrough(graph, current) - if current not in graph.constants or not isinstance(graph.constants[current], torch.Tensor): - return None - return {"weight_value_id": current, "offset": offset, "node_ids": tuple(sorted(node_ids))} - - -def _extract_normed_value(graph: IRGraph, value_id: str) -> dict[str, object] | None: - mul = producer(graph, strip_passthrough(graph, value_id)) - if mul is None or mul.op != "multiply" or len(mul.inputs) != 2: - return None - for x_candidate, rsqrt_candidate in ((mul.inputs[0], mul.inputs[1]), (mul.inputs[1], mul.inputs[0])): - rsqrt = _extract_rsqrt_chain(graph, rsqrt_candidate) - if rsqrt is None: - continue - if strip_passthrough(graph, x_candidate) != strip_passthrough(graph, rsqrt["pow_input_value_id"]): - continue - return { - "input_value_id": strip_passthrough(graph, x_candidate), - "pow_input_value_id": rsqrt["pow_input_value_id"], - "eps": rsqrt["eps"], - "node_ids": tuple(sorted({mul.id, *rsqrt["node_ids"]})), - } - return None - - -def _extract_rsqrt_chain(graph: IRGraph, value_id: str) -> dict[str, object] | None: - pow_node = producer(graph, strip_passthrough(graph, value_id)) - if pow_node is None or pow_node.op != "pow" or float(pow_node.attrs.get("exponent", 0.0)) != -0.5: - return None - add_node = producer(graph, pow_node.inputs[0]) - if add_node is None or add_node.op != "scalar_add": - return None - mean_node = producer(graph, add_node.inputs[0]) - if mean_node is None or mean_node.op != "mean" or not bool(mean_node.attrs.get("keepdim", False)): - return None - pow2_node = producer(graph, mean_node.inputs[0]) - if pow2_node is None or pow2_node.op != "pow" or float(pow2_node.attrs.get("exponent", 0.0)) != 2.0: - return None - if not _reduces_last_dim(graph, mean_node, pow2_node.inputs[0]): - return None - return { - "pow_input_value_id": pow2_node.inputs[0], - "eps": float(add_node.attrs.get("value", 0.0)), - "node_ids": (pow_node.id, add_node.id, mean_node.id, pow2_node.id), - } - - -def _reduces_last_dim(graph: IRGraph, mean_node: IRNode, input_value_id: str) -> bool: - value = graph.values.get(input_value_id) - if value is None or value.shape is None: - return False - rank = len(value.shape) - if rank == 0: - return False - axis = mean_node.attrs.get("axis") - if isinstance(axis, (list, tuple)): - if len(axis) != 1: - return False - axis = axis[0] - if not isinstance(axis, int): - return False - return axis % rank == rank - 1 diff --git a/python/cactus/transpile/fusion/rope.py b/python/cactus/transpile/fusion/rope.py deleted file mode 100644 index 25118b868..000000000 --- a/python/cactus/transpile/fusion/rope.py +++ /dev/null @@ -1,252 +0,0 @@ -from __future__ import annotations - -from dataclasses import dataclass -from typing import Any - -import torch - -from cactus.transpile.graph_ir import IRGraph -from cactus.transpile.graph_ir import IRNode -from cactus.transpile.fusion.common import producer -from cactus.transpile.fusion.common import strip_layout_passthrough -from cactus.transpile.fusion.common import strip_passthrough - - -@dataclass(frozen=True) -class RoPEMatch: - input_value_id: str - theta: float - position_offset: int - partial: bool - node_ids: tuple[str, ...] - - -def match_rope(graph: IRGraph, value_id: str) -> RoPEMatch | None: - desc = _extract_rope_descriptor(graph, value_id) - if desc is None: - return None - return RoPEMatch( - input_value_id=strip_passthrough(graph, desc["input_value_id"]), - theta=float(desc["theta"]), - position_offset=int(desc["position_offset"]), - partial=bool(desc.get("partial_rope", False)), - node_ids=tuple(sorted(set(desc.get("node_ids", ())))), - ) - - -def _extract_rope_descriptor(graph: IRGraph, value_id: str) -> dict[str, object] | None: - current = strip_layout_passthrough(graph, value_id) - node = producer(graph, current) - if node is None: - return None - if node.op == "rope": - return { - "input_value_id": node.inputs[0], - "theta": float(node.attrs.get("theta", 0.0)), - "position_offset": int(node.attrs.get("position_offset", 0)), - "partial_rope": False, - "node_ids": (node.id,), - } - classic = _extract_classic_rope_descriptor(graph, current) - if classic is not None: - classic["partial_rope"] = False - return classic - if node.op == "cat" and len(node.inputs) == 2: - for rope_branch_value_id, passthrough_value_id in ((node.inputs[0], node.inputs[1]), (node.inputs[1], node.inputs[0])): - rope_branch = _extract_rope_descriptor(graph, rope_branch_value_id) - if rope_branch is None: - continue - passthrough_node = producer(graph, strip_layout_passthrough(graph, passthrough_value_id)) - if passthrough_node is None or passthrough_node.op != "slice": - continue - if strip_passthrough(graph, passthrough_node.inputs[0]) != strip_passthrough(graph, rope_branch["input_value_id"]): - continue - rope_branch["partial_rope"] = True - rope_branch["node_ids"] = tuple(sorted({node.id, passthrough_node.id, *rope_branch["node_ids"]})) - return rope_branch - return None - - -def _extract_classic_rope_descriptor(graph: IRGraph, value_id: str) -> dict[str, object] | None: - node = producer(graph, strip_layout_passthrough(graph, value_id)) - if node is None or node.op != "add" or len(node.inputs) != 2: - return None - for direct_mult_id, rotated_mult_id in ((node.inputs[0], node.inputs[1]), (node.inputs[1], node.inputs[0])): - direct_mult = producer(graph, direct_mult_id) - rotated_mult = producer(graph, rotated_mult_id) - if direct_mult is None or direct_mult.op != "multiply" or rotated_mult is None or rotated_mult.op != "multiply": - continue - direct_match = _extract_direct_rope_branch(graph, direct_mult) - rotated_match = _extract_rotated_rope_branch(graph, rotated_mult) - if direct_match is None or rotated_match is None: - continue - input_value_id = strip_passthrough(graph, direct_match["input_value_id"]) - if input_value_id != strip_passthrough(graph, rotated_match["input_value_id"]): - continue - cos_info = _extract_rope_trig(graph, direct_match["trig_value_id"], expected="scalar_cos") - sin_info = _extract_rope_trig(graph, rotated_match["trig_value_id"], expected="scalar_sin") - if cos_info is None or sin_info is None: - continue - if abs(float(cos_info["theta"]) - float(sin_info["theta"])) > 1e-2: - continue - if int(cos_info["position_offset"]) != int(sin_info["position_offset"]): - continue - return { - "input_value_id": input_value_id, - "theta": cos_info["theta"], - "position_offset": cos_info["position_offset"], - "node_ids": tuple(sorted({node.id, *direct_match["node_ids"], *rotated_match["node_ids"], *cos_info["node_ids"], *sin_info["node_ids"]})), - } - return None - - -def _extract_direct_rope_branch(graph: IRGraph, node: IRNode) -> dict[str, object] | None: - for input_value_id, trig_value_id in ((node.inputs[0], node.inputs[1]), (node.inputs[1], node.inputs[0])): - if _unwrap_trig_value(graph, trig_value_id, expected="scalar_cos") is not None: - return {"input_value_id": input_value_id, "trig_value_id": trig_value_id, "node_ids": (node.id,)} - return None - - -def _extract_rotated_rope_branch(graph: IRGraph, node: IRNode) -> dict[str, object] | None: - for rotated_value_id, trig_value_id in ((node.inputs[0], node.inputs[1]), (node.inputs[1], node.inputs[0])): - input_value_id = _extract_rotate_half_source(graph, rotated_value_id) - if input_value_id is None: - continue - if _unwrap_trig_value(graph, trig_value_id, expected="scalar_sin") is None: - continue - return {"input_value_id": input_value_id, "trig_value_id": trig_value_id, "node_ids": (node.id,)} - return None - - -def _extract_rotate_half_source(graph: IRGraph, value_id: str) -> str | None: - cat_node = producer(graph, strip_passthrough(graph, value_id)) - if cat_node is None or cat_node.op != "cat" or len(cat_node.inputs) != 2: - return None - neg_node = producer(graph, cat_node.inputs[0]) - right_slice = producer(graph, cat_node.inputs[1]) - if neg_node is None or neg_node.op != "scalar_multiply" or float(neg_node.attrs.get("value", 0.0)) != -1.0: - return None - left_slice = producer(graph, neg_node.inputs[0]) - if left_slice is None or right_slice is None or left_slice.op != "slice" or right_slice.op != "slice": - return None - if strip_passthrough(graph, left_slice.inputs[0]) != strip_passthrough(graph, right_slice.inputs[0]): - return None - return left_slice.inputs[0] - - -def _extract_rope_trig(graph: IRGraph, value_id: str, *, expected: str) -> dict[str, object] | None: - trig_info = _unwrap_trig_value(graph, value_id, expected=expected) - if trig_info is None: - return None - trig_node = trig_info["trig_node"] - cat_node = producer(graph, trig_node.inputs[0]) - if cat_node is None or cat_node.op != "cat" or len(cat_node.inputs) != 2 or cat_node.inputs[0] != cat_node.inputs[1]: - return None - angle_info = _extract_rope_angle_source(graph, cat_node.inputs[0]) - if angle_info is None: - return None - matmul_node = angle_info["matmul_node"] - inv_freq_const_id = None - arange_node = None - for input_id in matmul_node.inputs: - if inv_freq_const_id is None: - inv_freq_const_id = _find_constant_ancestor(graph, input_id) - if arange_node is None: - arange_node = _find_arange_ancestor(graph, input_id) - if inv_freq_const_id is None or arange_node is None: - return None - theta = _infer_rope_theta(graph.constants[inv_freq_const_id]) - if theta is None: - return None - return { - "theta": theta, - "position_offset": int(arange_node.attrs.get("start", 0)), - "node_ids": tuple(sorted({trig_node.id, cat_node.id, matmul_node.id})), - } - - -def _unwrap_trig_value(graph: IRGraph, value_id: str, *, expected: str) -> dict[str, object] | None: - current = value_id - while True: - node = producer(graph, current) - if node is None: - return None - if node.op in {"unsqueeze", "precision_cast", "reshape", "view", "expand"}: - current = node.inputs[0] - continue - if node.op == "type_as": - current = node.inputs[0] - continue - if node.op == "scalar_multiply": - current = node.inputs[0] - continue - if node.op != expected: - return None - return {"trig_node": node} - - -def _extract_rope_angle_source(graph: IRGraph, value_id: str) -> dict[str, object] | None: - current = value_id - while True: - current = strip_layout_passthrough(graph, current) - node = producer(graph, current) - if node is None: - return None - if node.op == "permute": - matmul_node = producer(graph, node.inputs[0]) - if matmul_node is None or matmul_node.op != "matmul": - return None - return {"matmul_node": matmul_node} - if node.op == "matmul": - return {"matmul_node": node} - if node.op in {"index", "slice"} and len(node.inputs) == 1: - current = node.inputs[0] - continue - return None - - -def _find_constant_ancestor(graph: IRGraph, value_id: str) -> str | None: - current = value_id - visited: set[str] = set() - while current not in visited: - visited.add(current) - current = strip_layout_passthrough(graph, current) - if current in graph.constants and isinstance(graph.constants[current], torch.Tensor): - return current - node = producer(graph, current) - if node is None or len(node.inputs) != 1: - return None - current = node.inputs[0] - return None - - -def _find_arange_ancestor(graph: IRGraph, value_id: str) -> IRNode | None: - current = value_id - visited: set[str] = set() - while current not in visited: - visited.add(current) - current = strip_layout_passthrough(graph, current) - node = producer(graph, current) - if node is None: - return None - if node.op == "arange": - return node - if current in graph.constants and isinstance(graph.constants[current], torch.Tensor): - return None - if len(node.inputs) != 1: - return None - current = node.inputs[0] - return None - - -def _infer_rope_theta(value: Any) -> float | None: - if not isinstance(value, torch.Tensor): - return None - flat = value.detach().cpu().float().reshape(-1) - if flat.numel() < 2: - return None - second = float(flat[1].item()) - if second <= 0.0: - return None - inv_count = flat.numel() - return float((1.0 / second) ** inv_count) diff --git a/python/cactus/transpile/graph_ir.py b/python/cactus/transpile/graph_ir.py deleted file mode 100644 index 632585c57..000000000 --- a/python/cactus/transpile/graph_ir.py +++ /dev/null @@ -1,201 +0,0 @@ -import copy -from dataclasses import dataclass -from dataclasses import field - -import torch - - -def _clone_ir_payload(value, memo): - if isinstance(value, (str, int, float, bool)) or value is None: - return value - if isinstance(value, torch.Tensor): - # IR constants may be real tensors or export-time tensor wrappers. - # They are treated as immutable graph payloads, so preserve them by - # reference instead of triggering Tensor.__deepcopy__. - return value - if isinstance(value, tuple): - return tuple(_clone_ir_payload(item, memo) for item in value) - if isinstance(value, list): - return [_clone_ir_payload(item, memo) for item in value] - if isinstance(value, dict): - return { - copy.deepcopy(key, memo): _clone_ir_payload(inner, memo) - for key, inner in value.items() - } - try: - return copy.deepcopy(value, memo) - except Exception: - return value - - -@dataclass -class IRValue: - id: str - shape: tuple[int, ...] | None = None - dtype: str | None = None - producer: str | None = None - users: list[str] = field(default_factory=list) - meta: dict[str, object] = field(default_factory=dict) - - def __deepcopy__(self, memo): - existing = memo.get(id(self)) - if existing is not None: - return existing - cloned = IRValue( - id=self.id, - shape=_clone_ir_payload(self.shape, memo), - dtype=self.dtype, - producer=self.producer, - users=list(self.users), - meta=_clone_ir_payload(self.meta, memo), - ) - memo[id(self)] = cloned - return cloned - -@dataclass -class IRNode: - id: str - op: str - inputs: list[str] - outputs: list[str] - attrs: dict[str, object] = field(default_factory=dict) - meta: dict[str, object] = field(default_factory=dict) - kind: str = "generic" - - def __deepcopy__(self, memo): - existing = memo.get(id(self)) - if existing is not None: - return existing - cloned = IRNode( - id=self.id, - op=self.op, - inputs=list(self.inputs), - outputs=list(self.outputs), - attrs=_clone_ir_payload(self.attrs, memo), - meta=_clone_ir_payload(self.meta, memo), - kind=self.kind, - ) - memo[id(self)] = cloned - return cloned - -@dataclass -class IRGraph: - values: dict[str, IRValue] - nodes: dict[str, IRNode] - order: list[str] - inputs: list[str] - outputs: list[str] - constants: dict[str, object] = field(default_factory=dict) - meta: dict[str, object] = field(default_factory=dict) - - def __deepcopy__(self, memo): - existing = memo.get(id(self)) - if existing is not None: - return existing - cloned = IRGraph( - values={value_id: copy.deepcopy(value, memo) for value_id, value in self.values.items()}, - nodes={node_id: copy.deepcopy(node, memo) for node_id, node in self.nodes.items()}, - order=list(self.order), - inputs=list(self.inputs), - outputs=list(self.outputs), - constants={ - value_id: _clone_ir_payload(constant, memo) - for value_id, constant in self.constants.items() - }, - meta=_clone_ir_payload(self.meta, memo), - ) - memo[id(self)] = cloned - return cloned - - def add_node(self, node: IRNode) -> None: - if node.id in self.nodes: - raise ValueError(f"duplicate IR node id: {node.id}") - self.nodes[node.id] = node - for output in node.outputs: - if output in self.values: - raise ValueError(f"duplicate IR value id: {output}") - self.values[output] = IRValue(id=output, producer=node.id) - - def add_value(self, value: IRValue) -> None: - if value.id in self.values: - raise ValueError(f"duplicate IR value id: {value.id}") - self.values[value.id] = value - - -def verify_ir(graph: IRGraph) -> None: - order_set = set(graph.order) - node_keys = set(graph.nodes.keys()) - if len(order_set) != len(graph.order): - raise ValueError("IR order contains duplicate node ids") - if order_set != node_keys: - missing_in_order = node_keys - order_set - missing_in_nodes = order_set - node_keys - raise ValueError( - "IR order/nodes mismatch: " - f"missing_in_order={sorted(missing_in_order)} " - f"missing_in_nodes={sorted(missing_in_nodes)}" - ) - - seen_nodes: set[str] = set() - seen_outputs: set[str] = set() - expected_users: dict[str, list[str]] = {} - - for value_id in graph.inputs: - if value_id not in graph.values: - raise ValueError(f"IR input missing value entry: {value_id}") - expected_users.setdefault(value_id, []) - - for value_id in graph.constants: - if value_id not in graph.values: - raise ValueError(f"IR constant missing value entry: {value_id}") - expected_users.setdefault(value_id, []) - - for node_id in graph.order: - if node_id in seen_nodes: - raise ValueError(f"duplicate node in IR order: {node_id}") - seen_nodes.add(node_id) - - node = graph.nodes[node_id] - if not node.outputs: - raise ValueError(f"IR node has no outputs: {node_id}") - - for input_id in node.inputs: - if input_id not in graph.values: - raise ValueError(f"IR node {node_id} references missing input value {input_id}") - expected_users.setdefault(input_id, []).append(node_id) - producer = graph.values[input_id].producer - if producer is not None and producer not in graph.nodes: - raise ValueError( - f"IR value {input_id} for node {node_id} has missing producer {producer}" - ) - - for output_id in node.outputs: - if output_id in seen_outputs: - raise ValueError(f"IR value produced more than once: {output_id}") - seen_outputs.add(output_id) - if output_id not in graph.values: - raise ValueError(f"IR node {node_id} missing output value entry {output_id}") - value = graph.values[output_id] - if value.producer != node_id: - raise ValueError( - f"IR value {output_id} producer mismatch: expected {node_id}, got {value.producer}" - ) - expected_users.setdefault(output_id, []) - - for output_id in graph.outputs: - if output_id not in graph.values: - raise ValueError(f"IR graph output missing value entry: {output_id}") - - for value_id, value in graph.values.items(): - expected = expected_users.get(value_id, []) - if value.users != expected: - raise ValueError( - f"IR value {value_id} user list mismatch: expected {expected}, got {value.users}" - ) - if value.producer is None and value_id not in graph.inputs and value_id not in graph.constants: - if value_id in seen_outputs: - raise ValueError(f"IR value {value_id} is produced but has no producer") - if value.producer is not None and value.producer not in graph.nodes: - raise ValueError( - f"IR value {value_id} references unknown producer {value.producer}" - ) diff --git a/python/cactus/transpile/hf_model.py b/python/cactus/transpile/hf_model.py deleted file mode 100644 index 6719aae30..000000000 --- a/python/cactus/transpile/hf_model.py +++ /dev/null @@ -1,3406 +0,0 @@ -from __future__ import annotations - -import argparse -import copy -import gc -import hashlib -import itertools -import json -import os -import re -import shutil -import sys -from collections.abc import Mapping -from collections import Counter -from dataclasses import dataclass -from dataclasses import fields -from pathlib import Path -from typing import Any - -import numpy as np -import torch - -PYTHON_ROOT = Path(__file__).resolve().parents[2] -if str(PYTHON_ROOT) not in sys.path: - sys.path.insert(0, str(PYTHON_ROOT)) - -from cactus.transpile.runtime_compat import Graph -from cactus.convert.cactus_adapters.tensor_io import save_tensor_with_header -from cactus.convert.model_adapters.nemo import ensure_parakeet_tdt_nemo_source -from cactus.transpile.audio_preprocess import generic_log_mel_features as _generic_log_mel_features -from cactus.transpile.audio_preprocess import load_audio_waveform as _load_audio_waveform -from cactus.transpile.audio_preprocess import prepare_cactus_audio_features -from cactus.transpile.capture_pytorch import capture_model -from cactus.transpile.canonicalize.cleanup import canonicalize_exported_graph -from cactus.transpile.component_partition import extract_component_subgraphs -from cactus.transpile.component_partition import summarize_ir_components -from cactus.transpile.component_plan import infer_component_plan_from_config -from cactus.transpile.component_pipeline import capture_component_spec -from cactus.transpile.component_pipeline import execute_component_pipeline -from cactus.transpile.multimodal_runtime import prepare_gemma4_multimodal_inputs as _shared_prepare_gemma4_multimodal_inputs -from cactus.transpile.graph_ir import IRGraph -from cactus.transpile.graph_ir import verify_ir -from cactus.transpile.lower import TranspiledGraph -from cactus.transpile.lower import _lower_constant_value -from cactus.transpile.lower import _lower_input_value -from cactus.transpile.lower import _lower_ir_node -from cactus.transpile.lower import _lookup_weight_binding -from cactus.transpile.media_limits import resize_static_image -from cactus.transpile.model_adapters import UnsupportedComponentsError -from cactus.transpile.model_adapters import build_component_module_specs -from cactus.transpile.model_adapters import canonicalize_model_interface -from cactus.transpile.model_profiles import multimodal_context_tokens_for_model_type -from cactus.transpile.optimize_graph import FusionConfig -from cactus.transpile.optimize_graph import optimize_graph -from cactus.transpile.runtime_support import ensure_transformers_supports_model_type as _ensure_transformers_supports_profiled_model_type -from cactus.transpile.runtime_support import patch_torch_flex_attention_compat as _patch_torch_flex_attention_compat -from cactus.transpile.runtime_support import patch_transformers_torchvision_probe as _patch_transformers_torchvision_probe -from cactus.transpile.tdt_runtime import greedy_decode_parakeet_tdt_token_ids -from cactus.transpile.tdt_runtime import load_tdt_local_model -from cactus.transpile.tdt_runtime import prepare_parakeet_tdt_audio_features -from cactus.transpile.weight_compat import ensure_binding_compatible - -_DEFAULT_CAUSAL_PROMPT = "The capital of France is" -_DEFAULT_MULTIMODAL_CONTEXT_TOKENS = 2048 - - -@dataclass -class PreparedInputs: - names: tuple[str, ...] - tensors: tuple[torch.Tensor, ...] - metadata: dict[str, object] - - -def _ensure_transformers_supports_model_type(model_type: str) -> str | None: - return _ensure_transformers_supports_profiled_model_type(model_type) - - -def _resolve_local_snapshot(model_id_or_path: str) -> str | None: - explicit = Path(model_id_or_path) - if explicit.exists(): - nemo_export = ensure_parakeet_tdt_nemo_source(model_id_or_path) - if nemo_export is not None: - return nemo_export - return str(explicit) - - snapshots_dir = ( - Path.home() - / ".cache" - / "huggingface" - / "hub" - / ("models--" + model_id_or_path.replace("/", "--")) - / "snapshots" - ) - if not snapshots_dir.exists(): - return None - snapshots = sorted(path for path in snapshots_dir.iterdir() if path.is_dir()) - if not snapshots: - return None - snapshot = str(snapshots[-1]) - nemo_export = ensure_parakeet_tdt_nemo_source(snapshot) - if nemo_export is not None: - return nemo_export - return snapshot - - -def _snapshot_has_model_weights(path: str | Path) -> bool: - root = Path(path) - if not root.exists() or not root.is_dir(): - return False - candidates = ( - "model.safetensors", - "model.safetensors.index.json", - "pytorch_model.bin", - "pytorch_model.bin.index.json", - ) - return any((root / name).exists() for name in candidates) - - -def _validate_weights_dir(weights_dir: str | None, *, model_id: str) -> Path | None: - if not weights_dir: - return None - - root = Path(weights_dir).resolve() - if not root.exists(): - raise RuntimeError( - f"weights_dir does not exist: {root}\n" - "\n" - f"Create the folder first with:\n" - f" cactus convert {model_id} {root}\n" - ) - - manifest_path = root / "weights_manifest.json" - if not manifest_path.exists(): - raise RuntimeError( - f"weights_dir is missing weights_manifest.json: {manifest_path}\n" - "\n" - "The transpiler binds weights only through the converted CQ/Cactus " - "manifest; it no longer guesses filenames from model-specific layer names.\n" - "\n" - f"Re-convert with the current converter:\n" - f" cactus convert {model_id} {root}\n" - ) - - return root - - -def _serialize_json_compatible(value: Any) -> Any: - if isinstance(value, dict): - return {str(key): _serialize_json_compatible(inner) for key, inner in value.items()} - if isinstance(value, (list, tuple)): - return [_serialize_json_compatible(inner) for inner in value] - if isinstance(value, Path): - return str(value) - if isinstance(value, (str, int, float, bool)) or value is None: - return value - if isinstance(value, torch.dtype): - return str(value) - if isinstance(value, torch.Tensor): - payload = { - "type": "torch.Tensor", - "dtype": str(value.dtype), - "shape": list(value.shape), - } - if value.numel() == 1: - payload["data"] = value.detach().cpu().reshape(-1)[0].item() - return payload - if isinstance(value, np.ndarray): - payload = { - "type": "numpy.ndarray", - "dtype": str(value.dtype), - "shape": list(value.shape), - } - if value.size == 1: - payload["data"] = np.asarray(value).reshape(-1)[0].item() - return payload - if hasattr(value, "__dataclass_fields__"): - return { - field.name: _serialize_json_compatible(getattr(value, field.name)) - for field in fields(value) - } - try: - return repr(value) - except Exception: - return f"<{type(value).__module__}.{type(value).__name__}>" - - -def _slugify_model_artifact_name(value: str) -> str: - return re.sub(r"[^A-Za-z0-9._-]+", "_", value).strip("._-") or "model" - - -def _default_artifact_dir_for_model(model_id: str) -> Path: - return (Path.cwd() / "transpiled" / _slugify_model_artifact_name(model_id)).resolve() - - -def _graph_to_dict(graph) -> dict[str, object]: - return { - "meta": _serialize_json_compatible(graph.meta), - "inputs": list(graph.inputs), - "outputs": list(graph.outputs), - "constants": { - value_id: _serialize_json_compatible(constant) - for value_id, constant in graph.constants.items() - }, - "values": { - value_id: _serialize_json_compatible(value) - for value_id, value in graph.values.items() - }, - "nodes": [ - _serialize_json_compatible(graph.nodes[node_id]) - for node_id in graph.order - ], - } - - -def _write_json(path: Path, payload: dict[str, object]) -> None: - path.parent.mkdir(parents=True, exist_ok=True) - path.write_text(json.dumps(payload, indent=2, sort_keys=True) + "\n") - - -def _component_artifact_name(prefix: str, component: str) -> str: - return f"{prefix}_{component}.json" - - -def _component_graphs_to_payload(component_graphs: dict[str, IRGraph]) -> dict[str, object]: - return { - component: _graph_to_dict(graph) - for component, graph in component_graphs.items() - } - - -def _binding_entries_by_node_id( - bindings: list[dict[str, object]], -) -> dict[int, dict[str, object]]: - result: dict[int, dict[str, object]] = {} - for binding in bindings: - try: - result[int(binding["node_id"])] = binding - except Exception: - continue - return result - - -def _externalize_shared_rope_tables( - transpiled_graph: TranspiledGraph, - *, - weights_root: Path, - registry: dict[str, str], -) -> int: - value_ids = getattr(transpiled_graph, "bound_constant_value_ids", {}) or {} - already_bound = { - int(binding["node_id"]) - for binding in transpiled_graph.bound_constant_bindings - if isinstance(binding, dict) and "node_id" in binding - } - externalized = 0 - for constant in transpiled_graph.bound_constants: - node_id = int(constant.id) - if node_id in already_bound: - continue - value_id = str(value_ids.get(node_id, "")) - if not value_id.startswith("c_rope_table_"): - continue - data = np.ascontiguousarray(constant.numpy()) - digest = hashlib.sha1(data.tobytes()).hexdigest()[:16] - filename = registry.get(digest) - if filename is None: - filename = f"rope_table_{digest}.weights" - save_tensor_with_header(data, weights_root / filename, precision="FP16") - registry[digest] = filename - transpiled_graph.bound_constant_bindings.append( - { - "node_id": node_id, - "value_id": value_id, - "path": filename, - "kind": "weight", - "source_name": value_id, - } - ) - externalized += 1 - return externalized - - -def _embed_materialized_bound_constants(transpiled_graph: TranspiledGraph) -> int: - """Serialize small graph-local constants into the .cactus graph itself. - - CQ model weights remain external mmap files. Only materialized constants - without an existing weight binding are embedded. - """ - existing_bindings = list(_serialize_json_compatible(transpiled_graph.bound_constant_bindings)) - external_node_ids = set(_binding_entries_by_node_id(existing_bindings)) - mark_embedded_input = getattr(transpiled_graph.graph, "mark_embedded_input", None) - if not callable(mark_embedded_input): - raise RuntimeError( - "Cactus runtime does not support embedded graph constants; rebuild with cactus build --python" - ) - - embedded_count = 0 - for constant_tensor in transpiled_graph.bound_constants: - node_id = int(constant_tensor.id) - if node_id in external_node_ids: - continue - mark_embedded_input(constant_tensor) - embedded_count += 1 - return embedded_count - - -def _write_graph_binding_manifest( - *, - artifact_dir: Path, - filename: str, - model_id: str, - model_source: str, - task: str, - family: str, - inputs_metadata: dict[str, object], - transpiled_graph: TranspiledGraph, -) -> Path: - manifest_path = artifact_dir / filename - _write_json( - manifest_path, - { - "model_id": model_id, - "model_source": model_source, - "task": task, - "family": family, - "inputs": _serialize_json_compatible(inputs_metadata), - "runtime_input_node_ids": [int(tensor.id) for tensor in transpiled_graph.runtime_inputs], - "output_node_ids": [int(tensor.id) for tensor in transpiled_graph.outputs], - "bound_constant_bindings": _serialize_json_compatible(transpiled_graph.bound_constant_bindings), - }, - ) - return manifest_path - - -def _export_lfm2_vl_position_embedding_grid( - *, - model, - artifact_dir: Path, - component_metadata: dict[str, dict[str, object]], -) -> None: - root = getattr(model, "model", None) - vision_tower = getattr(root, "vision_tower", None) - vision_model = getattr(vision_tower, "vision_model", None) - embeddings = getattr(vision_model, "embeddings", None) - if embeddings is None: - return - pos_size = int(embeddings.position_embedding_size) - embed_dim = int(embeddings.embed_dim) - grid = ( - embeddings.position_embedding.weight.detach().cpu().to(torch.float32) - .reshape(pos_size, pos_size, embed_dim) - .contiguous() - .numpy() - ) - rel_path = Path("components") / "vision_encoder" / "position_embedding_grid.f32" - out_path = artifact_dir / rel_path - out_path.parent.mkdir(parents=True, exist_ok=True) - grid.tofile(out_path) - meta = component_metadata.setdefault("vision_encoder", {}) - meta["position_embedding_grid_path"] = str(rel_path) - meta["position_embedding_grid_shape"] = f"{pos_size},{pos_size},{embed_dim}" - print(f"saved_vision_position_embedding_grid={out_path}") - - -def _write_component_bundle( - *, - artifact_dir: Path, - model_id: str, - model_source: str, - task: str, - family: str, - inputs_metadata: dict[str, object], - raw_component_graphs: dict[str, IRGraph], - optimized_component_graphs: dict[str, IRGraph], - transpiled_component_graphs: dict[str, TranspiledGraph] | None = None, - component_io_signatures: dict[str, dict[str, tuple[str, ...]]] | None = None, - component_metadata: dict[str, dict[str, object]] | None = None, - graph_filename: str = "graph.cactus", - npu_encoder_mlpackages: dict[str, str] | None = None, -) -> Path: - bundle_dir = artifact_dir / "components" - component_order = [ - component - for component in ( - "source_encoder", - "audio_encoder", - "vision_encoder", - "vision_projector", - "text_embedding", - "lm_encoder", - "lm_encoder_text_chunk", - "decoder_prefill_chunk", - "decoder_embed_chunk", - "decoder", - "decoder_cross_kv", - "lm_encoder_step", - "lm_encoder_media_step", - "decoder_media_step", - "decoder_step", - "unspecified", - ) - if component in raw_component_graphs or component in optimized_component_graphs - ] - extra_components = sorted( - component - for component in set(raw_component_graphs) | set(optimized_component_graphs) - if component not in component_order - ) - component_order.extend(extra_components) - - manifest_components: list[dict[str, object]] = [] - - rope_table_registry: dict[str, str] = {} - for component in component_order: - raw_graph = raw_component_graphs.get(component) - optimized_graph = optimized_component_graphs.get(component) - transpiled_graph = None if transpiled_component_graphs is None else transpiled_component_graphs.get(component) - component_dir = bundle_dir / component - raw_relpath = None - optimized_relpath = None - graph_relpath = None - if raw_graph is not None: - raw_relpath = Path(component) / "raw_ir.json" - _write_json( - bundle_dir / raw_relpath, - { - "model_id": model_id, - "model_source": model_source, - "task": task, - "family": family, - "component": component, - "inputs": _serialize_json_compatible(inputs_metadata), - "graph": _graph_to_dict(raw_graph), - }, - ) - if optimized_graph is not None: - optimized_relpath = Path(component) / "optimized_ir.json" - _write_json( - bundle_dir / optimized_relpath, - { - "model_id": model_id, - "model_source": model_source, - "task": task, - "family": family, - "component": component, - "inputs": _serialize_json_compatible(inputs_metadata), - "graph": _graph_to_dict(optimized_graph), - }, - ) - if transpiled_graph is not None: - graph_relpath = Path(component) / graph_filename - component_dir.mkdir(parents=True, exist_ok=True) - shutil.rmtree(component_dir / "bound_constants", ignore_errors=True) - _externalize_shared_rope_tables( - transpiled_graph, - weights_root=artifact_dir, - registry=rope_table_registry, - ) - _embed_materialized_bound_constants(transpiled_graph) - transpiled_graph.graph.save(bundle_dir / graph_relpath) - - graph_for_signature = optimized_graph or raw_graph - if graph_for_signature is None: - continue - manifest_components.append( - { - "component": component, - "directory": str(component_dir.relative_to(artifact_dir)), - "raw_ir": None if raw_relpath is None else str((bundle_dir / raw_relpath).relative_to(artifact_dir)), - "optimized_ir": None if optimized_relpath is None else str((bundle_dir / optimized_relpath).relative_to(artifact_dir)), - "graph": None if graph_relpath is None else str((bundle_dir / graph_relpath).relative_to(artifact_dir)), - "inputs": list(graph_for_signature.inputs), - "outputs": list(graph_for_signature.outputs), - "logical_inputs": list((component_io_signatures or {}).get(component, {}).get("input_keys", ())), - "logical_outputs": list((component_io_signatures or {}).get(component, {}).get("output_keys", ())), - "metadata": _serialize_json_compatible((component_metadata or {}).get(component, {})), - "node_count": len(graph_for_signature.order), - "weight_binding_count": _count_weight_bindings(graph_for_signature), - "runtime_input_node_ids": [] if transpiled_graph is None else [int(tensor.id) for tensor in transpiled_graph.runtime_inputs], - "output_node_ids": [] if transpiled_graph is None else [int(tensor.id) for tensor in transpiled_graph.outputs], - "cache_state_node_ids": [] if transpiled_graph is None else [ - { - "layer_key": str(layer_key), - "key": int(key_tensor.id), - "value": int(value_tensor.id), - } - for layer_key, key_tensor, value_tensor in getattr(transpiled_graph, "cache_state_tensors", []) - ], - "bound_constant_bindings": [] if transpiled_graph is None else ( - list(_serialize_json_compatible(transpiled_graph.bound_constant_bindings)) - ), - } - ) - - manifest_path = bundle_dir / "manifest.json" - manifest_payload: dict[str, object] = { - "model_id": model_id, - "model_source": model_source, - "task": task, - "family": family, - "component_order": component_order, - "inputs": _serialize_json_compatible(inputs_metadata), - "components": manifest_components, - } - for manifest_key, mlpackage_path in (npu_encoder_mlpackages or {}).items(): - if mlpackage_path: - manifest_payload[manifest_key] = mlpackage_path - if family == "parakeet_tdt" and "npu_audio_encoder" in manifest_payload: - manifest_payload["npu_audio_compute_units"] = "CPU_AND_NE" - _write_json(manifest_path, manifest_payload) - return manifest_path - - -def _named_tensor_store(prepared: PreparedInputs) -> dict[str, torch.Tensor]: - return { - name: tensor - for name, tensor in zip(prepared.names, prepared.tensors, strict=True) - } - - -def _aggregate_component_counts(component_graphs: dict[str, IRGraph]) -> dict[str, int]: - return { - component: len(graph.order) - for component, graph in sorted(component_graphs.items()) - } - - -def _aggregate_component_op_counts(component_graphs: dict[str, IRGraph]) -> Counter: - counter: Counter[str] = Counter() - for graph in component_graphs.values(): - counter.update(graph.nodes[node_id].op for node_id in graph.order) - return counter - - -def _count_component_weight_bindings(component_graphs: dict[str, IRGraph]) -> int: - return sum(_count_weight_bindings(graph) for graph in component_graphs.values()) - - -def _execute_gemma4_component_pipeline( - *, - component_graphs, - prepared: PreparedInputs, -) -> np.ndarray: - store, _ = execute_component_pipeline( - [component_graphs[name] for name in ("vision_encoder", "audio_encoder", "lm_encoder", "decoder")], - initial_store=_named_tensor_store(prepared), - ) - logits = store["logits"] - return np.asarray(logits, dtype=np.float32) - - -def _run_parakeet_tdt_component_decode( - *, - component_graphs, - model, - prepared: PreparedInputs, -) -> dict[str, object]: - store, _ = execute_component_pipeline( - [component_graphs["audio_encoder"]], - initial_store=_named_tensor_store(prepared), - ) - encoder_hidden_states = np.asarray(store["encoder_hidden_states"]) - batch_size = int(encoder_hidden_states.shape[0]) - if batch_size != 1: - raise ValueError("Parakeet TDT component decode currently expects batch size 1") - - hidden_dtype = prepared.tensors[0].dtype - initial_states = model.initial_decoder_state( - batch_size=batch_size, - device=torch.device("cpu"), - dtype=hidden_dtype, - ) - initial_state_arrays = tuple(state.detach().cpu().numpy() for state in initial_states) - decoder_component = component_graphs["decoder"] - - def _step( - frame: np.ndarray, - token_id: int, - state_values: tuple[np.ndarray, ...], - ) -> tuple[np.ndarray, tuple[np.ndarray, ...]]: - input_store: dict[str, object] = { - "encoder_frame": np.ascontiguousarray(frame), - "token_ids": np.full((batch_size,), token_id, dtype=np.int64), - } - for index in range(model.config.predictor_num_layers): - input_store[f"state_h_{index}"] = np.ascontiguousarray(state_values[index * 2]) - input_store[f"state_c_{index}"] = np.ascontiguousarray(state_values[index * 2 + 1]) - - runtime_inputs = [input_store[key] for key in decoder_component.input_keys] - decoder_component.transpiled_graph.set_inputs(runtime_inputs) - outputs = decoder_component.transpiled_graph.execute() - logits = outputs[0].numpy().astype(np.float32, copy=False) - next_states = tuple(output.numpy() for output in outputs[1:]) - return logits, next_states - - emitted = greedy_decode_parakeet_tdt_token_ids( - config=model.config, - encoder_hidden_states=encoder_hidden_states, - initial_states=initial_state_arrays, - step=_step, - ) - - return { - "token_ids": emitted, - "transcript": model.decode_token_ids(emitted), - "encoder_hidden_shape": list(encoder_hidden_states.shape), - } - - -def _run_component_pipeline_transpile( - *, - args, - task: str, - family: str, - model_source: str, - model, - prepared: PreparedInputs, - component_specs, - fusion_config: FusionConfig, - weights_dir: str | None, - artifact_dir: Path | None, - processor_or_tokenizer, - canonical, -) -> int: - print(f"model_id={args.model_id}") - print(f"model_source={model_source}") - print(f"task={task}") - print(f"adapter_family={family}") - print("adapter_module=component_pipeline") - print(f"input_names={','.join(prepared.names)}") - for name, tensor in zip(prepared.names, prepared.tensors, strict=True): - print(f"input_{name}_shape={list(tensor.shape)}") - if weights_dir: - print(f"weights_dir={weights_dir}") - - npu_enabled = bool(getattr(args, "npu", False)) - npu_quantize = getattr(args, "npu_quantize", None) - npu_audio_quantize = getattr(args, "npu_audio_quantize", None) - npu_vision_quantize = getattr(args, "npu_vision_quantize", None) - npu_encoder_mlpackages: dict[str, str] = {} - if npu_enabled and artifact_dir is not None: - from .npu import run_encoder_pipeline - npu_encoder_mlpackages = run_encoder_pipeline( - component_specs, - artifact_dir, - enabled=True, - quantize_bits=npu_quantize, - audio_quantize_bits=npu_audio_quantize, - vision_quantize_bits=npu_vision_quantize, - ) - import gc as _gc - _gc.collect() - - print("capture_begin=true", flush=True) - captured_components = {} - for spec in component_specs: - print(f"capture_component_begin={spec.component}", flush=True) - captured_components[spec.component] = capture_component_spec(spec, fusion_config=fusion_config) - print(f"capture_component_done={spec.component}", flush=True) - print("capture_done=true", flush=True) - - raw_component_graphs = { - name: captured.raw_ir_graph - for name, captured in captured_components.items() - } - optimized_component_graphs = { - name: captured.optimized_ir_graph - for name, captured in captured_components.items() - } - transpiled_component_graphs = { - name: captured.transpiled_graph - for name, captured in captured_components.items() - } - component_io_signatures = { - name: { - "input_keys": tuple(captured.input_keys), - "output_keys": tuple(captured.output_keys), - } - for name, captured in captured_components.items() - } - component_metadata = { - name: dict(captured.metadata) - for name, captured in captured_components.items() - } - - raw_ir_nodes = sum(len(graph.order) for graph in raw_component_graphs.values()) - optimized_ir_nodes = sum(len(graph.order) for graph in optimized_component_graphs.values()) - binding_count = _count_component_weight_bindings(optimized_component_graphs) - raw_component_counts = _aggregate_component_counts(raw_component_graphs) - optimized_component_counts = _aggregate_component_counts(optimized_component_graphs) - op_counts = _aggregate_component_op_counts(optimized_component_graphs) - - print(f"raw_ir_nodes={raw_ir_nodes}") - print(f"optimized_ir_nodes={optimized_ir_nodes}") - print(f"weight_bindings={binding_count}") - if raw_component_counts: - print( - "raw_components=" - + ",".join(f"{name}:{count}" for name, count in sorted(raw_component_counts.items())) - ) - if optimized_component_counts: - print( - "optimized_components=" - + ",".join(f"{name}:{count}" for name, count in sorted(optimized_component_counts.items())) - ) - print( - "ops=" - f"attention:{op_counts.get('attention', 0)} " - f"conv1d:{op_counts.get('conv1d', 0)} " - f"conv2d:{op_counts.get('conv2d', 0)} " - f"batch_norm:{op_counts.get('batch_norm', 0)} " - f"layer_norm:{op_counts.get('layer_norm', 0)} " - f"rms_norm:{op_counts.get('rms_norm', 0)} " - f"rope:{op_counts.get('rope', 0)} " - f"linear:{op_counts.get('linear', 0)}" - ) - - if weights_dir and binding_count == 0: - raise RuntimeError( - f"No weight bindings were resolved from {weights_dir}\n" - "\n" - "The weights folder exists, but none of the component graphs matched entries in weights_manifest.json.\n" - "\n" - f"Recommended fix:\n" - f" cactus convert {args.model_id} {weights_dir}\n" - ) - - if artifact_dir is not None: - artifact_dir.mkdir(parents=True, exist_ok=True) - _write_json( - artifact_dir / "raw_ir.json", - { - "model_id": args.model_id, - "model_source": model_source, - "task": task, - "family": family, - "inputs": _serialize_json_compatible(prepared.metadata), - "component_order": [spec.component for spec in component_specs], - "components": _component_graphs_to_payload(raw_component_graphs), - }, - ) - _write_json( - artifact_dir / "optimized_ir.json", - { - "model_id": args.model_id, - "model_source": model_source, - "task": task, - "family": family, - "inputs": _serialize_json_compatible(prepared.metadata), - "component_order": [spec.component for spec in component_specs], - "components": _component_graphs_to_payload(optimized_component_graphs), - }, - ) - for component, component_graph in raw_component_graphs.items(): - _write_json( - artifact_dir / _component_artifact_name("raw_ir", component), - { - "model_id": args.model_id, - "model_source": model_source, - "task": task, - "family": family, - "component": component, - "inputs": _serialize_json_compatible(prepared.metadata), - "graph": _graph_to_dict(component_graph), - }, - ) - for component, component_graph in optimized_component_graphs.items(): - _write_json( - artifact_dir / _component_artifact_name("optimized_ir", component), - { - "model_id": args.model_id, - "model_source": model_source, - "task": task, - "family": family, - "component": component, - "inputs": _serialize_json_compatible(prepared.metadata), - "graph": _graph_to_dict(component_graph), - }, - ) - print(f"saved_raw_ir={artifact_dir / 'raw_ir.json'}") - print(f"saved_optimized_ir={artifact_dir / 'optimized_ir.json'}") - for component in raw_component_graphs: - print(f"saved_raw_component_ir_{component}={artifact_dir / _component_artifact_name('raw_ir', component)}") - for component in optimized_component_graphs: - print( - f"saved_optimized_component_ir_{component}=" - f"{artifact_dir / _component_artifact_name('optimized_ir', component)}" - ) - - if family == "lfm2_vl" and "vision_encoder" in component_metadata: - _export_lfm2_vl_position_embedding_grid( - model=model, - artifact_dir=artifact_dir, - component_metadata=component_metadata, - ) - - component_manifest_path = _write_component_bundle( - artifact_dir=artifact_dir, - model_id=args.model_id, - model_source=model_source, - task=task, - family=family, - inputs_metadata=prepared.metadata, - raw_component_graphs=raw_component_graphs, - optimized_component_graphs=optimized_component_graphs, - transpiled_component_graphs=transpiled_component_graphs, - component_io_signatures=component_io_signatures, - component_metadata=component_metadata, - graph_filename=args.graph_filename, - npu_encoder_mlpackages=npu_encoder_mlpackages, - ) - print(f"saved_component_bundle_manifest={component_manifest_path}") - for component in transpiled_component_graphs: - print( - f"saved_component_graph_{component}=" - f"{artifact_dir / 'components' / component / args.graph_filename}" - ) - print("note=saved split component graphs; use the component bundle manifest to rebind mmap weights/embeddings when loading") - - if args.skip_execute: - return 0 - - if task == "multimodal_causal_lm_logits": - print("execute_begin=true") - transpiled_output = _execute_gemma4_component_pipeline( - component_graphs=captured_components, - prepared=prepared, - ) - print("execute_done=true") - tokenizer_like = getattr(processor_or_tokenizer, "tokenizer", processor_or_tokenizer) - if args.skip_reference_compare: - result_payload = { - "model_id": args.model_id, - "model_source": model_source, - "task": task, - "family": family, - "inputs": _serialize_json_compatible(prepared.metadata), - "output_shape": list(transpiled_output.shape), - "raw_ir_nodes": raw_ir_nodes, - "optimized_ir_nodes": optimized_ir_nodes, - "weight_bindings": binding_count, - "reference_compare_skipped": True, - } - print(f"output_shape={list(transpiled_output.shape)}") - transpiled_next = int(np.argmax(transpiled_output[0, -1])) - print(f"transpiled_next_token_id={transpiled_next}") - result_payload["transpiled_next_token_id"] = transpiled_next - if hasattr(tokenizer_like, "decode"): - transpiled_token = tokenizer_like.decode([transpiled_next]) - print(f"transpiled_next_token={transpiled_token!r}") - result_payload["transpiled_next_token"] = transpiled_token - if artifact_dir is not None: - _write_json(artifact_dir / "result.json", result_payload) - print(f"saved_result={artifact_dir / 'result.json'}") - return 0 - - print("reference_begin=true") - with torch.no_grad(): - reference_output = canonical.module(*prepared.tensors).detach().float().cpu().numpy() - print("reference_done=true") - max_abs_diff = float(np.max(np.abs(reference_output - transpiled_output))) - mean_abs_diff = float(np.mean(np.abs(reference_output - transpiled_output))) - hf_next = int(np.argmax(reference_output[0, -1])) - transpiled_next = int(np.argmax(transpiled_output[0, -1])) - print(f"hf_next_token_id={hf_next}") - print(f"transpiled_next_token_id={transpiled_next}") - print(f"logits_max_abs_diff={max_abs_diff:.6f}") - print(f"logits_mean_abs_diff={mean_abs_diff:.6f}") - if hasattr(tokenizer_like, "decode"): - print(f"hf_next_token={tokenizer_like.decode([hf_next])!r}") - print(f"transpiled_next_token={tokenizer_like.decode([transpiled_next])!r}") - result_payload = { - "model_id": args.model_id, - "model_source": model_source, - "task": task, - "family": family, - "inputs": _serialize_json_compatible(prepared.metadata), - "output_shape": list(reference_output.shape), - "raw_ir_nodes": raw_ir_nodes, - "optimized_ir_nodes": optimized_ir_nodes, - "weight_bindings": binding_count, - "max_abs_diff": max_abs_diff, - "mean_abs_diff": mean_abs_diff, - "hf_next_token_id": hf_next, - "transpiled_next_token_id": transpiled_next, - } - if hasattr(tokenizer_like, "decode"): - result_payload["hf_next_token"] = tokenizer_like.decode([hf_next]) - result_payload["transpiled_next_token"] = tokenizer_like.decode([transpiled_next]) - if artifact_dir is not None: - _write_json(artifact_dir / "result.json", result_payload) - print(f"saved_result={artifact_dir / 'result.json'}") - return 0 - - if task == "causal_lm_logits": - tokenizer_like = getattr(processor_or_tokenizer, "tokenizer", processor_or_tokenizer) - result_payload = { - "model_id": args.model_id, - "model_source": model_source, - "task": task, - "family": family, - "inputs": _serialize_json_compatible(prepared.metadata), - "raw_ir_nodes": raw_ir_nodes, - "optimized_ir_nodes": optimized_ir_nodes, - "weight_bindings": binding_count, - } - if "decoder" in captured_components: - print("execute_begin=true") - initial_store = { - name: tensor - for name, tensor in zip(prepared.names, prepared.tensors, strict=True) - } - store, _ = execute_component_pipeline( - [captured_components["decoder"]], - initial_store=initial_store, - ) - print("execute_done=true") - transpiled_output = np.asarray(store["logits"]) - transpiled_next = int(np.argmax(transpiled_output[0, -1])) - print(f"output_shape={list(transpiled_output.shape)}") - print(f"transpiled_next_token_id={transpiled_next}") - result_payload["output_shape"] = list(transpiled_output.shape) - result_payload["transpiled_next_token_id"] = transpiled_next - if hasattr(tokenizer_like, "decode"): - transpiled_token = tokenizer_like.decode([transpiled_next]) - print(f"transpiled_next_token={transpiled_token!r}") - result_payload["transpiled_next_token"] = transpiled_token - if not args.skip_reference_compare and canonical is not None: - print("reference_begin=true") - with torch.no_grad(): - reference_output = canonical.module(*prepared.tensors).detach().float().cpu().numpy() - print("reference_done=true") - max_abs_diff = float(np.max(np.abs(reference_output - transpiled_output))) - mean_abs_diff = float(np.mean(np.abs(reference_output - transpiled_output))) - hf_next = int(np.argmax(reference_output[0, -1])) - print(f"hf_next_token_id={hf_next}") - print(f"logits_max_abs_diff={max_abs_diff:.6f}") - print(f"logits_mean_abs_diff={mean_abs_diff:.6f}") - result_payload["hf_next_token_id"] = hf_next - result_payload["max_abs_diff"] = max_abs_diff - result_payload["mean_abs_diff"] = mean_abs_diff - if hasattr(tokenizer_like, "decode"): - result_payload["hf_next_token"] = tokenizer_like.decode([hf_next]) - else: - result_payload["reference_compare_skipped"] = True - result_payload["note"] = "causal LM component bundle contains only cached decode step components" - if artifact_dir is not None: - _write_json(artifact_dir / "result.json", result_payload) - print(f"saved_result={artifact_dir / 'result.json'}") - return 0 - - if task == "tdt_transcription": - print("execute_begin=true") - transpiled_decode = _run_parakeet_tdt_component_decode( - component_graphs=captured_components, - model=model, - prepared=prepared, - ) - print("execute_done=true") - print(f"transpiled_transcript={transpiled_decode['transcript']!r}") - result_payload = { - "model_id": args.model_id, - "model_source": model_source, - "task": task, - "family": family, - "inputs": _serialize_json_compatible(prepared.metadata), - "raw_ir_nodes": raw_ir_nodes, - "optimized_ir_nodes": optimized_ir_nodes, - "weight_bindings": binding_count, - "transpiled_token_ids": transpiled_decode["token_ids"], - "transpiled_transcript": transpiled_decode["transcript"], - "encoder_hidden_shape": transpiled_decode["encoder_hidden_shape"], - } - if args.skip_reference_compare: - result_payload["reference_compare_skipped"] = True - if artifact_dir is not None: - _write_json(artifact_dir / "result.json", result_payload) - print(f"saved_result={artifact_dir / 'result.json'}") - return 0 - - print("reference_begin=true") - reference_token_ids = model.greedy_decode_token_ids(prepared.tensors[0]) - reference_transcript = model.decode_token_ids(reference_token_ids) - print("reference_done=true") - print(f"reference_transcript={reference_transcript!r}") - print(f"transcript_match={reference_transcript == transpiled_decode['transcript']}") - result_payload["reference_token_ids"] = reference_token_ids - result_payload["reference_transcript"] = reference_transcript - result_payload["transcript_match"] = bool(reference_transcript == transpiled_decode["transcript"]) - result_payload["token_id_match"] = bool(reference_token_ids == transpiled_decode["token_ids"]) - if artifact_dir is not None: - _write_json(artifact_dir / "result.json", result_payload) - print(f"saved_result={artifact_dir / 'result.json'}") - return 0 - - raise NotImplementedError(f"component pipeline execution is not implemented for task={task}") - - -def _parse_dtype(name: str) -> torch.dtype: - normalized = name.strip().lower() - mapping = { - "float16": torch.float16, - "fp16": torch.float16, - "float32": torch.float32, - "fp32": torch.float32, - "bfloat16": torch.bfloat16, - "bf16": torch.bfloat16, - } - if normalized not in mapping: - raise ValueError(f"unsupported torch dtype: {name}") - return mapping[normalized] - - -class TranspileWrapper(torch.nn.Module): - def __init__(self, adapter_module: torch.nn.Module, *, weights_dir: str | None = None): - super().__init__() - self.adapter = adapter_module - self.weights_dir = weights_dir - - def forward(self, *bound_inputs: torch.Tensor) -> torch.Tensor: - return self.adapter(*bound_inputs) - - def get_transpile_metadata(self) -> dict[str, object]: - metadata: dict[str, object] = {} - provider = getattr(self.adapter, "get_transpile_metadata", None) - if callable(provider): - provided = provider() - if isinstance(provided, dict): - metadata.update(provided) - graph_meta: dict[str, object] = {} - base_graph = metadata.get("graph", {}) - if isinstance(base_graph, dict): - graph_meta.update(base_graph) - if self.weights_dir: - graph_meta["weights_dir"] = self.weights_dir - metadata["graph"] = graph_meta - return metadata - - -def _tie_lfm2_vl_lm_head_if_needed(model: torch.nn.Module) -> str | None: - if str(getattr(getattr(model, "config", None), "model_type", "") or "").lower() != "lfm2_vl": - return None - lm_head = getattr(model, "lm_head", None) - model_root = getattr(model, "model", None) - language_model = getattr(model_root, "language_model", None) - embed_tokens = getattr(language_model, "embed_tokens", None) - if not isinstance(lm_head, torch.nn.Linear) or not isinstance(embed_tokens, torch.nn.Embedding): - return None - if tuple(lm_head.weight.shape) != tuple(embed_tokens.weight.shape): - return None - if lm_head.weight.data_ptr() == embed_tokens.weight.data_ptr(): - return None - lm_head.weight = embed_tokens.weight - return "tied LFM2-VL lm_head.weight to language_model.embed_tokens.weight" - - -def _load_local_torch_state_dict(model_source: str) -> dict[str, torch.Tensor] | None: - root = Path(model_source) - if not root.exists() or not root.is_dir(): - return None - - safetensors_path = root / "model.safetensors" - if safetensors_path.exists(): - try: - from safetensors.torch import load_file # type: ignore - - return load_file(str(safetensors_path)) - except Exception: - return None - - pytorch_path = root / "pytorch_model.bin" - if pytorch_path.exists(): - try: - loaded = torch.load(str(pytorch_path), map_location="cpu", weights_only=True) - except TypeError: - loaded = torch.load(str(pytorch_path), map_location="cpu") - return loaded if isinstance(loaded, dict) else None - - return None - - -def _remap_gemma4_checkpoint_state_dict(state_dict: dict[str, torch.Tensor]) -> dict[str, torch.Tensor]: - remapped: dict[str, torch.Tensor] = {} - for key, value in state_dict.items(): - new_key = key - if "audio_tower" in new_key: - new_key = re.sub(r"subsample_conv_projection\.layer(\d+)\.", r"subsample_conv_projection.conv_\1.", new_key) - new_key = new_key.replace("audio_tower.layers.", "audio_tower.conformer.") - new_key = new_key.replace(".feed_forward1.", ".ffw_layer_start.") - new_key = new_key.replace(".feed_forward2.", ".ffw_layer_end.") - new_key = re.sub(r"\.self_attn\.(q_proj|k_proj|v_proj)\.", r".attention.attn.\1.", new_key) - new_key = new_key.replace(".self_attn.per_dim_scale", ".attention.attn.per_dim_scale") - new_key = new_key.replace(".self_attn.relative_k_proj.", ".attention.attn.relative_position_embedding.pos_proj.") - new_key = new_key.replace(".self_attn.post.", ".attention.post.") - new_key = new_key.replace(".norm_pre_attn.", ".attention.pre_attn_norm.") - new_key = new_key.replace(".norm_post_attn.", ".attention.post_norm.") - new_key = new_key.replace(".norm_out.", ".norm.") - new_key = new_key.replace(".linear.weight", ".weight") - remapped[new_key] = value - if new_key.endswith(".attention.attn.per_dim_scale"): - remapped[new_key.replace(".per_dim_scale", ".per_dim_key_scale")] = value - if "lm_head.weight" not in remapped: - tied_embedding = remapped.get("model.language_model.embed_tokens.weight") - if tied_embedding is not None: - remapped["lm_head.weight"] = tied_embedding - return remapped - - -def _repair_gemma4_checkpoint_weights(model: torch.nn.Module, model_source: str) -> dict[str, object]: - multimodal_backbone = getattr(model, "model", model) - audio_tower = getattr(multimodal_backbone, "audio_tower", None) - # Older/local Gemma4 variants used Cactus-style module names such as - # audio_tower.conformer and bare clippable-linear weights. Current HF - # Gemma4 classes use audio_tower.layers and *.linear.weight, and - # from_pretrained already loads them correctly. Do not reload a legacy - # remap into the current layout, because that leaves real weights missing. - if hasattr(audio_tower, "layers"): - return {"applied": False, "reason": "current HF Gemma4 layout uses audio_tower.layers"} - if not hasattr(audio_tower, "conformer"): - return {"applied": False, "reason": "current HF Gemma4 layout does not need legacy key remap"} - - raw_state_dict = _load_local_torch_state_dict(model_source) - if raw_state_dict is None: - return {"applied": False, "reason": "no local checkpoint state_dict"} - - remapped_state_dict = _remap_gemma4_checkpoint_state_dict(raw_state_dict) - load_result = model.load_state_dict(remapped_state_dict, strict=False) - return { - "applied": True, - "missing_keys": list(load_result.missing_keys), - "unexpected_keys": list(load_result.unexpected_keys), - } - - -def _load_config_json(model_id_or_path: str) -> dict[str, object]: - local_snapshot = _resolve_local_snapshot(model_id_or_path) - config_source = Path(local_snapshot) / "config.json" if local_snapshot else None - if config_source is not None and config_source.exists(): - return json.loads(config_source.read_text()) - explicit = Path(model_id_or_path) / "config.json" - if explicit.exists(): - return json.loads(explicit.read_text()) - return {} - - -def _load_optional_json(model_id_or_path: str, filename: str) -> dict[str, object]: - local_snapshot = _resolve_local_snapshot(model_id_or_path) - candidate = Path(local_snapshot) / filename if local_snapshot else None - if candidate is not None and candidate.exists(): - return json.loads(candidate.read_text()) - explicit = Path(model_id_or_path) / filename - if explicit.exists(): - return json.loads(explicit.read_text()) - return {} - - -def _load_gemma4_tokenizer_fallback(source_candidates: list[str]) -> object | None: - try: - from transformers import PreTrainedTokenizerFast # type: ignore - except Exception: - return None - - for source in source_candidates: - root = Path(source) - if not root.exists() or not root.is_dir(): - continue - tokenizer_json = root / "tokenizer.json" - tokenizer_config_path = root / "tokenizer_config.json" - if not tokenizer_json.exists() or not tokenizer_config_path.exists(): - continue - - tokenizer_config = json.loads(tokenizer_config_path.read_text()) - tokenizer = PreTrainedTokenizerFast( - tokenizer_file=str(tokenizer_json), - bos_token=tokenizer_config.get("bos_token"), - eos_token=tokenizer_config.get("eos_token"), - unk_token=tokenizer_config.get("unk_token"), - pad_token=tokenizer_config.get("pad_token"), - mask_token=tokenizer_config.get("mask_token"), - padding_side=str(tokenizer_config.get("padding_side", "right")), - additional_special_tokens=list(tokenizer_config.get("extra_special_tokens", []) or []), - ) - model_max_length = tokenizer_config.get("model_max_length") - if isinstance(model_max_length, int): - tokenizer.model_max_length = model_max_length - - chat_template_path = root / "chat_template.jinja" - if chat_template_path.exists(): - tokenizer.chat_template = chat_template_path.read_text() - - for token_attr in ( - "image_token", - "audio_token", - "boi_token", - "eoi_token", - "boa_token", - "eoa_token", - ): - token_value = tokenizer_config.get(token_attr) - if isinstance(token_value, str): - setattr(tokenizer, token_attr, token_value) - setattr(tokenizer, f"{token_attr}_id", tokenizer.convert_tokens_to_ids(token_value)) - return tokenizer - - return None - - -def _load_gemma4_processor_fallback( - *, - source_candidates: list[str], - common_kwargs: dict[str, object], -) -> object | None: - try: - from transformers.models.gemma4.feature_extraction_gemma4 import Gemma4AudioFeatureExtractor # type: ignore - from transformers.models.gemma4.image_processing_gemma4 import Gemma4ImageProcessor # type: ignore - from transformers.models.gemma4.processing_gemma4 import Gemma4Processor # type: ignore - except Exception: - return None - - tokenizer = None - processor_config: dict[str, object] = {} - for source in source_candidates: - try: - try: - from transformers import AutoTokenizer # type: ignore - - tokenizer = AutoTokenizer.from_pretrained(source, **common_kwargs) - except Exception: - tokenizer = _load_gemma4_tokenizer_fallback([source]) - processor_config = _load_optional_json(source, "processor_config.json") - if tokenizer is not None and processor_config: - break - except Exception: - continue - if tokenizer is None or not processor_config: - return None - - feature_config = dict(processor_config.get("feature_extractor", {}) or {}) - image_config = dict(processor_config.get("image_processor", {}) or {}) - feature_config.pop("feature_extractor_type", None) - image_config.pop("image_processor_type", None) - - feature_extractor = Gemma4AudioFeatureExtractor(**feature_config) - image_processor = Gemma4ImageProcessor(**image_config) - - processor_kwargs = {} - for key in ("image_seq_length", "audio_seq_length", "audio_ms_per_token"): - if key in processor_config: - processor_kwargs[key] = processor_config[key] - - return Gemma4Processor( - feature_extractor=feature_extractor, - image_processor=image_processor, - tokenizer=tokenizer, - **processor_kwargs, - ) - - -def _infer_task_from_config(model_id_or_path: str) -> str: - config = _load_config_json(model_id_or_path) - plan = infer_component_plan_from_config(config, model_id=model_id_or_path) - if plan is not None: - return plan.task - architectures = [str(value) for value in config.get("architectures", []) if isinstance(value, str)] - model_type = str(config.get("model_type", "") or "").lower() - decoding_cfg = config.get("decoding") - if isinstance(decoding_cfg, dict) and str(decoding_cfg.get("model_type", "") or "").lower() == "tdt": - return "tdt_transcription" - loss_cfg = config.get("loss") - if isinstance(loss_cfg, dict) and str(loss_cfg.get("loss_name", "") or "").lower() == "tdt": - return "tdt_transcription" - - if "nomic" in model_type or any("NomicBert" in value for value in architectures): - return "text_embedding" - if any("CausalLM" in value for value in architectures): - return "causal_lm_logits" - if any("CTC" in value for value in architectures): - return "ctc_logits" - if model_type == "whisper": - return "seq2seq_transcription" - if any("ConditionalGeneration" in value and "Whisper" in value for value in architectures): - return "seq2seq_transcription" - - lowered_id = model_id_or_path.lower() - if "nomic-embed" in lowered_id: - return "text_embedding" - if "parakeet-tdt" in lowered_id: - return "tdt_transcription" - if "whisper" in lowered_id: - return "seq2seq_transcription" - if "ctc" in lowered_id: - return "ctc_logits" - if any(token in lowered_id for token in ("qwen", "gemma", "llama", "mistral", "lfm")): - return "causal_lm_logits" - - raise RuntimeError( - f"Could not infer transpile task for {model_id_or_path}.\n" - "\n" - "Pass one explicitly with --task, for example:\n" - " --task causal_lm_logits\n" - " --task tdt_transcription\n" - " --task ctc_logits\n" - " --task encoder_hidden_states\n" - " --task seq2seq_transcription\n" - ) - - -def _resolve_audio_sample_rate(processor: object) -> int: - for attr_name in ("feature_extractor", "tokenizer"): - child = getattr(processor, attr_name, None) - sample_rate = getattr(child, "sampling_rate", None) - if isinstance(sample_rate, int) and sample_rate > 0: - return sample_rate - sample_rate = getattr(processor, "sampling_rate", None) - if isinstance(sample_rate, int) and sample_rate > 0: - return sample_rate - return 16000 - - -def _infer_fallback_audio_input_names(config: dict[str, object], task: str) -> tuple[str, ...]: - model_type = str(config.get("model_type", "") or "").lower() - if model_type == "whisper": - return ("input_features",) - if task == "seq2seq_transcription": - return ("input_features",) - if task == "encoder_hidden_states": - return ("input_features",) - if task == "tdt_transcription": - return ("input_features",) - audio_cfg = config.get("audio_config") - if isinstance(audio_cfg, dict) and any(key in audio_cfg for key in ("features", "input_feat_size", "num_mel_bins")): - return ("input_features", "attention_mask") if task == "ctc_logits" else ("input_features",) - encoder_cfg = config.get("encoder") - if isinstance(encoder_cfg, dict) and any(key in encoder_cfg for key in ("feat_in", "num_mel_bins")): - return ("input_features", "attention_mask") if task == "ctc_logits" else ("input_features",) - return ("input_values", "attention_mask") if task == "ctc_logits" else ("input_values",) - - -def _resolve_encoder_module(model: torch.nn.Module) -> torch.nn.Module | None: - get_encoder = getattr(model, "get_encoder", None) - if callable(get_encoder): - encoder = get_encoder() - if isinstance(encoder, torch.nn.Module): - return encoder - encoder = getattr(model, "encoder", None) - if isinstance(encoder, torch.nn.Module): - return encoder - model_attr = getattr(model, "model", None) - if model_attr is not None: - encoder = getattr(model_attr, "encoder", None) - if isinstance(encoder, torch.nn.Module): - return encoder - return None - - -def _infer_expected_input_feature_frames(model: torch.nn.Module) -> int | None: - config = getattr(model, "config", None) - max_source_positions = getattr(config, "max_source_positions", None) - if not isinstance(max_source_positions, int) or max_source_positions <= 0: - return None - - encoder = _resolve_encoder_module(model) - if encoder is None: - return None - - stride_product = 1 - found_conv = False - for child in encoder.children(): - if isinstance(child, torch.nn.Conv1d): - stride = child.stride[0] if isinstance(child.stride, tuple) else child.stride - stride_product *= int(stride) - found_conv = True - - if not found_conv: - return None - return int(max_source_positions) * stride_product - - -def _prepare_fallback_audio_inputs( - *, - input_names: tuple[str, ...], - config: dict[str, object], - preprocessor_config: dict[str, object], - model: torch.nn.Module, - task: str, - audio_file: str, - torch_dtype: torch.dtype, -) -> PreparedInputs: - if not input_names: - input_names = _infer_fallback_audio_input_names(config, task) - preprocessor_root = config.get("preprocessor") - preprocessor_root = preprocessor_root if isinstance(preprocessor_root, dict) else {} - encoder_root = config.get("encoder") - encoder_root = encoder_root if isinstance(encoder_root, dict) else {} - - sample_rate = int( - preprocessor_config.get( - "sampling_rate", - preprocessor_root.get("sample_rate", config.get("sampling_rate", 16000)), - ) - or 16000 - ) - num_mels = int( - preprocessor_config.get( - "feature_size", - preprocessor_root.get( - "features", - encoder_root.get("feat_in", config.get("num_mel_bins", config.get("feature_size", 80))), - ), - ) - or 80 - ) - encoder_cfg = config.get("encoder_config") - if isinstance(encoder_cfg, dict): - num_mels = int(encoder_cfg.get("num_mel_bins", encoder_cfg.get("feat_in", num_mels)) or num_mels) - raw_hop_length = preprocessor_config.get( - "hop_length", - preprocessor_root.get("window_stride", config.get("hop_length", 160)), - ) - if raw_hop_length is None: - raw_hop_length = 160 - if isinstance(raw_hop_length, float) and raw_hop_length > 0.0 and raw_hop_length < 1.0: - hop_length = max(1, int(round(float(sample_rate) * raw_hop_length))) - else: - hop_length = int(raw_hop_length) - - n_fft = int(preprocessor_config.get("n_fft", config.get("n_fft", preprocessor_root.get("n_fft", 400))) or 400) - raw_frame_length = preprocessor_config.get( - "win_length", - preprocessor_config.get( - "frame_length", - preprocessor_root.get("window_size", preprocessor_root.get("n_window_size", config.get("frame_length", n_fft))), - ), - ) - if raw_frame_length is None: - raw_frame_length = n_fft - if isinstance(raw_frame_length, float) and raw_frame_length > 0.0 and raw_frame_length < 1.0: - frame_length = max(1, int(round(float(sample_rate) * raw_frame_length))) - else: - frame_length = int(raw_frame_length) - preemphasis = preprocessor_config.get("preemphasis", preprocessor_root.get("preemph")) - if preemphasis is not None: - preemphasis = float(preemphasis) - waveform = _load_audio_waveform(audio_file, target_sample_rate=sample_rate) - - tensors: list[torch.Tensor] = [] - if input_names and input_names[0] == "input_features": - model_type = str(config.get("model_type", "") or "").lower() - is_parakeet_like = ( - task == "tdt_transcription" - or "parakeet" in model_type - or "parakeet" in type(model).__module__.lower() - ) - if is_parakeet_like: - expected_frames = _infer_expected_input_feature_frames(model) - parakeet_features, _ = prepare_parakeet_tdt_audio_features( - audio_file, - expected_frames=expected_frames, - expected_mels=num_mels, - torch_dtype=torch_dtype, - ) - tensors.append(parakeet_features) - return PreparedInputs( - names=input_names[: len(tensors)], - tensors=tuple(tensors), - metadata={ - "audio_file": str(Path(audio_file).resolve()), - "sample_rate": sample_rate, - "fallback_audio_preprocessor": True, - "native_parakeet_audio_preprocessor": True, - "input_shapes": { - name: list(tensor.shape) - for name, tensor in zip(input_names, tensors) - }, - }, - ) - is_whisper_like = ( - str(config.get("model_type", "") or "").lower() == "whisper" - or "transformers.models.whisper." in type(model).__module__ - ) - if is_whisper_like: - expected_frames = _infer_expected_input_feature_frames(model) or 3000 - try: - whisper_features, active_frames = prepare_cactus_audio_features( - audio_file, - model_type="whisper", - expected_frames=expected_frames, - expected_mels=num_mels, - torch_dtype=torch_dtype, - layout="mels_frames", - ) - tensors.append(whisper_features) - return PreparedInputs( - names=input_names[: len(tensors)], - tensors=tuple(tensors), - metadata={ - "audio_file": str(Path(audio_file).resolve()), - "sample_rate": sample_rate, - "fallback_audio_preprocessor": True, - "native_cactus_audio_preprocessor": True, - "active_feature_frames": active_frames, - "input_shapes": { - name: list(tensor.shape) - for name, tensor in zip(input_names, tensors) - }, - }, - ) - except Exception: - pass - features, feature_length = _generic_log_mel_features( - waveform, - sample_rate=sample_rate, - num_mels=num_mels, - n_fft=n_fft, - hop_length=hop_length, - frame_length=frame_length, - preemphasis=preemphasis, - ) - expected_frames = _infer_expected_input_feature_frames(model) - attention_mask: np.ndarray | None = None - target_frames = feature_length - if isinstance(expected_frames, int) and expected_frames > 0: - target_frames = expected_frames - - active_frames = min(feature_length, target_frames) - features = features[:active_frames, :] - if target_frames > active_frames: - pad_width = target_frames - active_frames - features = np.pad(features, ((0, pad_width), (0, 0)), mode="constant") - attention_mask = np.zeros((target_frames,), dtype=np.bool_) - attention_mask[:active_frames] = True - - if is_whisper_like: - features = np.ascontiguousarray(features.T) - - tensors.append(torch.from_numpy(features).unsqueeze(0).to(dtype=torch_dtype)) - if attention_mask is not None and len(input_names) > 1 and input_names[1] == "attention_mask": - tensors.append(torch.from_numpy(attention_mask).unsqueeze(0)) - else: - input_values = torch.from_numpy(waveform).unsqueeze(0).to(dtype=torch_dtype) - tensors.append(input_values) - if len(input_names) > 1 and input_names[1] == "attention_mask": - tensors.append(torch.ones_like(input_values, dtype=torch.float32)) - - return PreparedInputs( - names=input_names[: len(tensors)], - tensors=tuple(tensors), - metadata={ - "audio_file": str(Path(audio_file).resolve()), - "sample_rate": sample_rate, - "fallback_audio_preprocessor": True, - "input_shapes": { - name: list(tensor.shape) - for name, tensor in zip(input_names, tensors) - }, - }, - ) - - -def _prepare_audio_inputs( - processor: object | None, - *, - input_names: tuple[str, ...], - config: dict[str, object], - preprocessor_config: dict[str, object], - model: torch.nn.Module, - task: str, - audio_file: str, - torch_dtype: torch.dtype, -) -> PreparedInputs: - if processor is None: - return _prepare_fallback_audio_inputs( - input_names=input_names, - config=config, - preprocessor_config=preprocessor_config, - model=model, - task=task, - audio_file=audio_file, - torch_dtype=torch_dtype, - ) - - model_type = str(config.get("model_type", "") or "").lower() - if model_type == "whisper" or "transformers.models.whisper." in type(model).__module__: - return _prepare_fallback_audio_inputs( - input_names=input_names, - config=config, - preprocessor_config=preprocessor_config, - model=model, - task=task, - audio_file=audio_file, - torch_dtype=torch_dtype, - ) - - sample_rate = _resolve_audio_sample_rate(processor) - waveform = _load_audio_waveform(audio_file, target_sample_rate=sample_rate) - batch = processor( - waveform, - sampling_rate=sample_rate, - return_tensors="pt", - ) - - preferred_keys = tuple(input_names) + tuple( - key for key in ("input_features", "input_values", "attention_mask") if key not in input_names - ) - tensor_keys = [key for key, value in batch.items() if isinstance(value, torch.Tensor)] - ordered_keys = [key for key in preferred_keys if key in tensor_keys] - ordered_keys.extend(key for key in tensor_keys if key not in ordered_keys) - if not ordered_keys: - raise RuntimeError(f"processor did not return tensor inputs for audio file: {audio_file}") - - tensors: list[torch.Tensor] = [] - for key in ordered_keys: - value = batch[key] - if not isinstance(value, torch.Tensor): - continue - if torch.is_floating_point(value): - value = value.to(dtype=torch_dtype) - tensors.append(value) - - return PreparedInputs( - names=tuple(ordered_keys[: len(tensors)]), - tensors=tuple(tensors), - metadata={ - "audio_file": str(Path(audio_file).resolve()), - "sample_rate": sample_rate, - "fallback_audio_preprocessor": False, - "input_shapes": { - name: list(tensor.shape) - for name, tensor in zip(ordered_keys, tensors) - }, - }, - ) - - -def _contains_token_subsequence(token_ids: list[int], subsequence: list[int]) -> bool: - if not subsequence: - return False - width = len(subsequence) - return any(token_ids[index : index + width] == subsequence for index in range(len(token_ids) - width + 1)) - - -def _resolve_whisper_decoder_prompt_token_ids( - tokenizer_or_processor: object, - *, - prompt: str | None, - decoder_start_token_id: int | None, - forced_decoder_ids: list[list[int]] | tuple[tuple[int, int], ...] | None = None, -) -> list[int]: - tokenizer = getattr(tokenizer_or_processor, "tokenizer", tokenizer_or_processor) - encode = getattr(tokenizer, "encode", None) - if not callable(encode): - encode = None - - normalized_prompt = (prompt or "").strip() - if normalized_prompt == _DEFAULT_CAUSAL_PROMPT: - normalized_prompt = "" - - prompt_token_ids: list[int] = [] - if decoder_start_token_id is not None: - prompt_token_ids.append(int(decoder_start_token_id)) - - normalized_forced_ids: list[tuple[int, int]] = [] - if forced_decoder_ids is not None: - for item in forced_decoder_ids: - if ( - isinstance(item, (list, tuple)) - and len(item) == 2 - and isinstance(item[0], int) - and isinstance(item[1], int) - ): - normalized_forced_ids.append((int(item[0]), int(item[1]))) - normalized_forced_ids.sort(key=lambda pair: pair[0]) - - next_position = 1 - for position, token_id in normalized_forced_ids: - if position != next_position: - break - prompt_token_ids.append(int(token_id)) - next_position += 1 - - if normalized_prompt: - if encode is None: - raise RuntimeError("Whisper transpile requires a tokenizer to encode a non-empty prompt") - try: - encoded_prompt = encode(normalized_prompt, add_special_tokens=False) - except TypeError: - encoded_prompt = encode(normalized_prompt) - prompt_token_ids.extend(int(value) for value in encoded_prompt) - - if not prompt_token_ids: - if encode is None: - raise RuntimeError("Whisper transpile requires a tokenizer or decoder_start_token_id") - try: - prompt_token_ids = [int(value) for value in encode("<|startoftranscript|>", add_special_tokens=False)] - except TypeError: - prompt_token_ids = [int(value) for value in encode("<|startoftranscript|>")] - if not prompt_token_ids and decoder_start_token_id is not None: - prompt_token_ids = [int(decoder_start_token_id)] - - return prompt_token_ids - - -def _augment_whisper_seq2seq_metadata( - prepared: PreparedInputs, - *, - tokenizer_or_processor: object, - model: torch.nn.Module, - prompt: str | None, - max_new_tokens: int, -) -> PreparedInputs: - config = getattr(model, "config", None) - decoder_start_token_id = getattr(config, "decoder_start_token_id", None) - eos_token_id = getattr(config, "eos_token_id", None) - pad_token_id = getattr(config, "pad_token_id", None) - max_target_positions = int(getattr(config, "max_target_positions", 0) or 0) - suppress_tokens = [int(value) for value in (getattr(config, "suppress_tokens", None) or [])] - begin_suppress_tokens = [int(value) for value in (getattr(config, "begin_suppress_tokens", None) or [])] - - decoder_input_ids = _resolve_whisper_decoder_prompt_token_ids( - tokenizer_or_processor, - prompt=prompt, - decoder_start_token_id=int(decoder_start_token_id) if isinstance(decoder_start_token_id, int) else None, - forced_decoder_ids=getattr(config, "forced_decoder_ids", None), - ) - if not decoder_input_ids: - raise RuntimeError("Whisper transpile could not resolve decoder prompt token ids") - - target_token_count = len(decoder_input_ids) + max(1, int(max_new_tokens)) - if max_target_positions > 0: - target_token_count = min(target_token_count, max_target_positions) - target_token_count = max(target_token_count, len(decoder_input_ids)) - - metadata = dict(prepared.metadata) - metadata.update( - { - "decoder_input_ids": [int(value) for value in decoder_input_ids], - "decoder_start_token_id": None if decoder_start_token_id is None else int(decoder_start_token_id), - "eos_token_id": None if eos_token_id is None else int(eos_token_id), - "pad_token_id": None if pad_token_id is None else int(pad_token_id), - "target_token_count": int(target_token_count), - "max_target_positions": int(max_target_positions), - "suppress_tokens": suppress_tokens, - "begin_suppress_tokens": begin_suppress_tokens, - } - ) - return PreparedInputs( - names=prepared.names, - tensors=prepared.tensors, - metadata=metadata, - ) - - -def _tokenize_text_prompt( - tokenizer: object, - prompt: str, - *, - enable_thinking_if_supported: bool = False, -) -> torch.Tensor: - apply_chat_template = getattr(tokenizer, "apply_chat_template", None) - if callable(apply_chat_template): - try: - encoded = apply_chat_template( - [{"role": "user", "content": prompt}], - tokenize=True, - add_generation_prompt=True, - return_tensors="pt", - enable_thinking=bool(enable_thinking_if_supported), - ) - if isinstance(encoded, torch.Tensor) and encoded.ndim == 2: - return encoded.to(dtype=torch.long) - if isinstance(encoded, Mapping): - input_ids = encoded.get("input_ids") - if isinstance(input_ids, torch.Tensor) and input_ids.ndim == 2: - return input_ids.to(dtype=torch.long) - except Exception: - pass - encoded = tokenizer(prompt, return_tensors="pt") - return encoded["input_ids"].to(dtype=torch.long) - - -def _resolve_text_padding_token_id(tokenizer: object | None) -> int: - for attr_name in ("pad_token_id", "eos_token_id", "bos_token_id"): - token_id = getattr(tokenizer, attr_name, None) if tokenizer is not None else None - if isinstance(token_id, int) and token_id >= 0: - return int(token_id) - return 0 - - -def _resolve_graph_safe_text_padding_token_id( - tokenizer: object | None, - prompt_input_ids: torch.Tensor, -) -> int: - # Keep the tokenizer pad ID in the captured graph metadata. The lowerer keeps - # large token-id comparisons in FP32, so Qwen-style high pad IDs are safe and - # runtime padding stays consistent with the traced mask. - return _resolve_text_padding_token_id(tokenizer) - - -def _prepare_text_inputs( - tokenizer: object, - *, - prompt: str, - input_ids_text: str | None, - max_new_tokens: int, - enable_thinking_if_supported: bool = False, -) -> PreparedInputs: - if input_ids_text: - token_ids = [int(part.strip()) for part in input_ids_text.split(",") if part.strip()] - if not token_ids: - raise ValueError("--input-ids was provided but no ids were parsed") - prompt_input_ids = torch.tensor([token_ids], dtype=torch.long) - else: - prompt_input_ids = _tokenize_text_prompt( - tokenizer, - prompt, - enable_thinking_if_supported=enable_thinking_if_supported, - ) - if prompt_input_ids.ndim != 2 or int(prompt_input_ids.shape[0]) != 1: - raise ValueError( - "causal_lm_logits transpile currently expects prompt input ids with shape [1, T], " - f"got {tuple(int(dim) for dim in prompt_input_ids.shape)}" - ) - if int(max_new_tokens) < 0: - raise ValueError("--max-new-tokens must be non-negative") - - prompt_token_count = int(prompt_input_ids.shape[1]) - target_token_count = prompt_token_count + int(max_new_tokens) - padding_token_id = _resolve_graph_safe_text_padding_token_id(tokenizer, prompt_input_ids) - if target_token_count > prompt_token_count: - input_ids = torch.full((1, target_token_count), padding_token_id, dtype=torch.long) - input_ids[:, :prompt_token_count] = prompt_input_ids - else: - input_ids = prompt_input_ids - return PreparedInputs( - names=("input_ids",), - tensors=(input_ids,), - metadata={ - "prompt": prompt, - "prompt_input_ids": prompt_input_ids.tolist(), - "input_ids": input_ids.tolist(), - "prompt_token_count": prompt_token_count, - "target_token_count": target_token_count, - "max_new_tokens": int(max_new_tokens), - "padding_token_id": int(padding_token_id), - "enable_thinking": bool(enable_thinking_if_supported), - }, - ) - - -_NOMIC_EMBED_SEQ_LEN = 256 - - -def _prepare_text_embedding_inputs( - tokenizer: object, - *, - prompt: str, - input_ids_text: str | None, - seq_len: int = _NOMIC_EMBED_SEQ_LEN, -) -> PreparedInputs: - base = _prepare_text_inputs( - tokenizer, - prompt=prompt, - input_ids_text=input_ids_text, - max_new_tokens=0, - ) - prompt_input_ids = base.tensors[0] - prompt_token_count = int(prompt_input_ids.shape[1]) - padding_token_id = int(base.metadata.get("padding_token_id", 0)) - target = max(1, int(seq_len)) - input_ids = torch.full((1, target), padding_token_id, dtype=torch.long) - attention_mask = torch.zeros((1, target), dtype=torch.long) - keep = min(prompt_token_count, target) - input_ids[:, :keep] = prompt_input_ids[:, :keep] - attention_mask[:, :keep] = 1 - return PreparedInputs( - names=("input_ids", "attention_mask"), - tensors=(input_ids, attention_mask), - metadata={ - "prompt": prompt, - "prompt_token_count": prompt_token_count, - "target_token_count": target, - "padding_token_id": padding_token_id, - }, - ) - - -def _add_multimodal_generation_headroom( - prepared: PreparedInputs, - *, - tokenizer: object | None, - max_new_tokens: int, - min_context_tokens: int | None = None, -) -> PreparedInputs: - requested = max(0, int(max_new_tokens)) - if "input_ids" not in prepared.names: - return prepared - - tensor_by_name = dict(zip(prepared.names, prepared.tensors, strict=True)) - input_ids = tensor_by_name.get("input_ids") - if not isinstance(input_ids, torch.Tensor) or input_ids.ndim != 2: - return prepared - - prompt_token_count = int(input_ids.shape[1]) - min_context = max(0, int(min_context_tokens or 0)) - target_token_count = max(prompt_token_count + requested, min_context) - if target_token_count <= prompt_token_count: - return prepared - padding_token_id = _resolve_graph_safe_text_padding_token_id(tokenizer, input_ids) - - padded_tensors: list[torch.Tensor] = [] - for name, tensor in zip(prepared.names, prepared.tensors, strict=True): - if ( - name == "position_ids" - and tensor.ndim == 3 - and int(tensor.shape[1]) == 1 - and int(tensor.shape[2]) == prompt_token_count - ): - padded = torch.zeros( - (int(tensor.shape[0]), 1, target_token_count), - dtype=tensor.dtype, - device=tensor.device, - ) - padded[:, :, :prompt_token_count] = tensor - padded_tensors.append(padded) - continue - if ( - name not in {"input_ids", "attention_mask", "token_type_ids", "mm_token_type_ids"} - or tensor.ndim != 2 - or int(tensor.shape[0]) != 1 - or int(tensor.shape[1]) != prompt_token_count - ): - padded_tensors.append(tensor) - continue - pad_value = int(padding_token_id) if name == "input_ids" else 0 - padded = torch.full( - (1, target_token_count), - pad_value, - dtype=tensor.dtype, - device=tensor.device, - ) - padded[:, :prompt_token_count] = tensor - padded_tensors.append(padded) - - metadata = dict(prepared.metadata) - metadata["prompt_token_count"] = prompt_token_count - metadata["target_token_count"] = target_token_count - metadata["max_new_tokens"] = max(requested, target_token_count - prompt_token_count) - metadata["padding_token_id"] = int(padding_token_id) - input_shapes = dict(metadata.get("input_shapes") or {}) - for name, tensor in zip(prepared.names, padded_tensors, strict=True): - input_shapes[name] = [int(dim) for dim in tensor.shape] - metadata["input_shapes"] = input_shapes - return PreparedInputs( - names=prepared.names, - tensors=tuple(padded_tensors), - metadata=metadata, - ) - - -def _multimodal_context_token_floor(model_type: str = "") -> int: - default_context = multimodal_context_tokens_for_model_type( - model_type, - _DEFAULT_MULTIMODAL_CONTEXT_TOKENS, - ) - raw = os.environ.get( - "CACTUS_TRANSPILER_MULTIMODAL_CONTEXT_TOKENS", - str(default_context), - ) - try: - return max(0, int(raw)) - except (TypeError, ValueError): - return default_context - - -_LFM2_VL_MULTIMODAL_INPUT_ORDER = ( - "input_ids", - "attention_mask", - "pixel_values", - "spatial_shapes", - "pixel_attention_mask", -) - -_QWEN3_5_MULTIMODAL_INPUT_ORDER = ( - "input_ids", - "attention_mask", - "position_ids", - "pixel_values", - "image_grid_thw", -) - - -def _load_image_inputs(image_files: tuple[str, ...]) -> list[object]: - if not image_files: - return [] - try: - from PIL import Image # type: ignore - except Exception as exc: # pragma: no cover - raise RuntimeError(f"Pillow is required for --image-file: {exc}") from exc - - images: list[object] = [] - for image_file in image_files: - path = Path(image_file).resolve() - if not path.exists(): - raise RuntimeError(f"image_file does not exist: {path}") - with Image.open(path) as image: - images.append(resize_static_image(image.convert("RGB")).copy()) - return images - - -def _prepare_gemma4_multimodal_inputs( - processor: object | None, - *, - prompt: str, - image_files: tuple[str, ...], - audio_file: str | None, - torch_dtype: torch.dtype, - system_prompt: str = "", - enable_thinking_if_supported: bool = False, - use_gemma4_chat_template: bool = False, -) -> PreparedInputs: - if processor is None: - raise RuntimeError("multimodal Gemma4 transpile requires an AutoProcessor with image and audio support") - shared = _shared_prepare_gemma4_multimodal_inputs( - processor, - prompt=prompt, - image_files=image_files, - audio_file=audio_file, - torch_dtype=torch_dtype, - system_prompt=system_prompt, - enable_thinking_if_supported=enable_thinking_if_supported, - use_gemma4_chat_template=use_gemma4_chat_template, - ) - return PreparedInputs( - names=shared.names, - tensors=shared.tensors, - metadata=shared.metadata, - ) - - -def _prepare_lfm2_vl_multimodal_inputs( - processor: object | None, - *, - prompt: str, - image_files: tuple[str, ...], - torch_dtype: torch.dtype, - system_prompt: str = "", - enable_thinking_if_supported: bool = False, -) -> PreparedInputs: - if processor is None: - raise RuntimeError("LFM2-VL multimodal transpile requires an AutoProcessor with image support") - images = _load_image_inputs(image_files) - if not images: - raise RuntimeError("LFM2-VL multimodal transpile requires at least one --image-file") - - user_content: list[dict[str, object]] = [{"type": "image", "image": image} for image in images] - user_content.append({"type": "text", "text": prompt.strip()}) - - messages: list[dict[str, object]] = [] - normalized_system = system_prompt.strip() - if normalized_system: - messages.append({"role": "system", "content": normalized_system}) - messages.append({"role": "user", "content": user_content}) - - batch: Mapping[str, object] - apply_chat_template = getattr(processor, "apply_chat_template", None) - thinking_prefix = "\n" if enable_thinking_if_supported else "\n\n\n\n" - image_placeholders = "".join("<|vision_start|><|image_pad|><|vision_end|>" for _ in images) - fallback_text = ( - f"<|im_start|>user\n{image_placeholders}{prompt.strip()}<|im_end|>\n" - f"<|im_start|>assistant\n{thinking_prefix}" - ) - if callable(apply_chat_template): - try: - batch = apply_chat_template( - messages, - add_generation_prompt=True, - tokenize=True, - return_dict=True, - return_tensors="pt", - ) - except ValueError as exc: - if "chat template" not in str(exc).lower(): - raise - batch = processor(text=fallback_text, images=images, return_tensors="pt") - else: - batch = processor(text=fallback_text, images=images, return_tensors="pt") - - tensors: list[torch.Tensor] = [] - names: list[str] = [] - input_shapes: dict[str, list[int]] = {} - for key in _LFM2_VL_MULTIMODAL_INPUT_ORDER: - value = batch.get(key) if hasattr(batch, "get") else None - if not isinstance(value, torch.Tensor): - raise RuntimeError(f"LFM2-VL processor did not return required tensor input: {key}") - if torch.is_floating_point(value): - value = value.to(dtype=torch_dtype) - elif key == "pixel_attention_mask": - value = value.to(dtype=torch.int64) - else: - value = value.to(dtype=torch.long) - names.append(key) - tensors.append(value) - input_shapes[key] = [int(dim) for dim in value.shape] - - return PreparedInputs( - names=tuple(names), - tensors=tuple(tensors), - metadata={ - "prompt": prompt, - "system_prompt": system_prompt, - "image_files": [str(Path(path).resolve()) for path in image_files], - "input_shapes": input_shapes, - "enable_thinking": bool(enable_thinking_if_supported), - }, - ) - - -def _compute_qwen3_5_position_ids( - input_ids: torch.Tensor, - mm_token_type_ids: torch.Tensor, - image_grid_thw: torch.Tensor, - attention_mask: torch.Tensor | None, - *, - spatial_merge_size: int = 2, -) -> torch.Tensor: - position_ids = torch.zeros( - 3, - int(input_ids.shape[0]), - int(input_ids.shape[1]), - dtype=input_ids.dtype, - device=input_ids.device, - ) - image_grid_iter = iter(image_grid_thw) - for batch_idx in range(int(input_ids.shape[0])): - token_types = mm_token_type_ids[batch_idx] - valid_mask = attention_mask[batch_idx].bool() if attention_mask is not None else None - if valid_mask is not None: - token_types = token_types[valid_mask] - - current_pos = 0 - position_parts: list[torch.Tensor] = [] - for modality_type, group in itertools.groupby(enumerate(token_types.tolist()), lambda item: int(item[1])): - group_items = list(group) - token_count = group_items[-1][0] - group_items[0][0] + 1 - if int(modality_type) == 0: - part = torch.arange(token_count, dtype=input_ids.dtype, device=input_ids.device) - position_parts.append(part.view(1, -1).expand(3, -1) + current_pos) - current_pos += token_count - continue - - grid_thw = next(image_grid_iter) - grid_t = int(grid_thw[0].item()) - grid_h = int(grid_thw[1].item()) - grid_w = int(grid_thw[2].item()) - llm_grid_t = grid_t - llm_grid_h = grid_h // int(spatial_merge_size) - llm_grid_w = grid_w // int(spatial_merge_size) - image_seq_length = llm_grid_t * llm_grid_h * llm_grid_w - position_width = torch.arange( - current_pos, - current_pos + llm_grid_w, - dtype=input_ids.dtype, - device=input_ids.device, - ).repeat(llm_grid_h * llm_grid_t) - position_height = torch.arange( - current_pos, - current_pos + llm_grid_h, - dtype=input_ids.dtype, - device=input_ids.device, - ).repeat_interleave(llm_grid_w * llm_grid_t) - position_temporal = torch.full( - (image_seq_length,), - current_pos, - dtype=input_ids.dtype, - device=input_ids.device, - ) - position_parts.append(torch.stack([position_temporal, position_height, position_width], dim=0)) - current_pos += max(grid_h, grid_w) // int(spatial_merge_size) - - if not position_parts: - continue - positions = torch.cat(position_parts, dim=1) - if valid_mask is not None: - position_ids[:, batch_idx, valid_mask] = positions.to(position_ids.device) - else: - position_ids[:, batch_idx, : positions.shape[1]] = positions.to(position_ids.device) - return position_ids - - -def _prepare_qwen3_5_multimodal_inputs( - processor: object | None, - *, - prompt: str, - image_files: tuple[str, ...], - torch_dtype: torch.dtype, - system_prompt: str = "", - enable_thinking_if_supported: bool = False, -) -> PreparedInputs: - if processor is None: - raise RuntimeError("Qwen3.5 multimodal transpile requires an AutoProcessor with image support") - images = _load_image_inputs(image_files) - if not images: - raise RuntimeError("Qwen3.5 multimodal transpile requires at least one --image-file") - - user_content: list[dict[str, object]] = [{"type": "image", "image": image} for image in images] - user_content.append({"type": "text", "text": prompt.strip()}) - - messages: list[dict[str, object]] = [] - normalized_system = system_prompt.strip() - if normalized_system: - messages.append({"role": "system", "content": normalized_system}) - messages.append({"role": "user", "content": user_content}) - - apply_chat_template = getattr(processor, "apply_chat_template", None) - if callable(apply_chat_template): - batch = apply_chat_template( - messages, - add_generation_prompt=True, - tokenize=True, - return_dict=True, - return_tensors="pt", - ) - else: - batch = processor(text=prompt, images=images, return_tensors="pt") - - input_ids = batch.get("input_ids") if hasattr(batch, "get") else None - attention_mask = batch.get("attention_mask") if hasattr(batch, "get") else None - mm_token_type_ids = batch.get("mm_token_type_ids") if hasattr(batch, "get") else None - image_grid_thw = batch.get("image_grid_thw") if hasattr(batch, "get") else None - if not all(isinstance(value, torch.Tensor) for value in (input_ids, attention_mask, mm_token_type_ids, image_grid_thw)): - raise RuntimeError("Qwen3.5 processor did not return required multimodal text/grid tensors") - spatial_merge_size = int(getattr(getattr(processor, "image_processor", None), "merge_size", 2) or 2) - position_ids = _compute_qwen3_5_position_ids( - input_ids.to(dtype=torch.long), - mm_token_type_ids.to(dtype=torch.long), - image_grid_thw.to(dtype=torch.long), - attention_mask.to(dtype=torch.long), - spatial_merge_size=spatial_merge_size, - ) - batch["position_ids"] = position_ids - - tensors: list[torch.Tensor] = [] - names: list[str] = [] - input_shapes: dict[str, list[int]] = {} - for key in _QWEN3_5_MULTIMODAL_INPUT_ORDER: - value = batch.get(key) if hasattr(batch, "get") else None - if not isinstance(value, torch.Tensor): - raise RuntimeError(f"Qwen3.5 processor did not return required tensor input: {key}") - if torch.is_floating_point(value): - value = value.to(dtype=torch_dtype) - else: - value = value.to(dtype=torch.long) - names.append(key) - tensors.append(value) - input_shapes[key] = [int(dim) for dim in value.shape] - - return PreparedInputs( - names=tuple(names), - tensors=tuple(tensors), - metadata={ - "prompt": prompt, - "system_prompt": system_prompt, - "image_files": [str(Path(path).resolve()) for path in image_files], - "input_shapes": input_shapes, - "enable_thinking": bool(enable_thinking_if_supported), - "spatial_merge_size": spatial_merge_size, - }, - ) - - -def _load_model_source(model_id: str, *, local_files_only: bool) -> str: - local_snapshot = _resolve_local_snapshot(model_id) - if local_snapshot and _snapshot_has_model_weights(local_snapshot): - return local_snapshot - if local_snapshot and local_files_only: - raise RuntimeError( - f"Found local snapshot for {model_id}, but it is incomplete and has no model weights:\n" - f" {local_snapshot}\n" - "\n" - "Re-run without --local-files-only to let transformers download the missing weights." - ) - return model_id - - -def _load_transformers_bundle( - *, - model_id: str, - task: str, - torch_dtype: torch.dtype, - token: str | None, - trust_remote_code: bool, - local_files_only: bool, -): - config = _load_config_json(model_id) - config_model_type = str(config.get("model_type", "") or "").lower() - if not config_model_type: - normalized_model_id = model_id.lower().replace("_", "-") - if "qwen3-vl" in normalized_model_id: - config_model_type = "qwen3_vl" - elif "lfm2.5-vl" in normalized_model_id or "lfm2-vl" in normalized_model_id: - config_model_type = "lfm2_vl" - if config_model_type: - config = {"model_type": config_model_type} - external_transformers_site_packages = _ensure_transformers_supports_model_type(config_model_type) - patch_note = _patch_transformers_torchvision_probe() - if patch_note: - print(f"note={patch_note}") - flex_patch_note = _patch_torch_flex_attention_compat() - if flex_patch_note: - print(f"note={flex_patch_note}") - if external_transformers_site_packages: - print(f"note=using external transformers install for {config_model_type}: {external_transformers_site_packages}") - - try: - from transformers import AutoFeatureExtractor # type: ignore - from transformers import AutoModel # type: ignore - from transformers import AutoModelForCTC # type: ignore - from transformers import AutoModelForCausalLM # type: ignore - from transformers import AutoModelForImageTextToText # type: ignore - from transformers import AutoModelForSeq2SeqLM # type: ignore - from transformers import AutoModelForSpeechSeq2Seq # type: ignore - from transformers import AutoProcessor # type: ignore - from transformers import AutoTokenizer # type: ignore - except Exception as exc: # pragma: no cover - raise RuntimeError(f"transformers is not available: {exc}") from exc - - model_source = _load_model_source(model_id, local_files_only=local_files_only) - source_candidates = [] - for candidate in (model_source, model_id): - if candidate not in source_candidates: - source_candidates.append(candidate) - common_kwargs: dict[str, object] = { - "trust_remote_code": trust_remote_code, - "local_files_only": local_files_only, - } - if token: - common_kwargs["token"] = token - - if task == "tdt_transcription": - model = load_tdt_local_model(model_source, torch_dtype=torch_dtype).eval() - return model_source, None, model, config - - if task == "text_embedding": - # nomic-bert weights are produced from the model's remote modeling code - # (fused Wqkv / experts.mlp.w1,w2); the transformers-native nomic_bert has a - # different architecture, so the remote code is required to match the bundle. - remote_kwargs = {**common_kwargs, "trust_remote_code": True} - tokenizer = None - tokenizer_errors: list[str] = [] - for source in source_candidates: - try: - tokenizer = AutoTokenizer.from_pretrained(source, **remote_kwargs) - break - except Exception as exc: - tokenizer_errors.append(f"{source}: {exc}") - if tokenizer is None: - raise RuntimeError( - f"Could not load tokenizer for {model_id}:\n" - + "\n".join(tokenizer_errors) - ) - model = AutoModel.from_pretrained( - model_source, - dtype=torch_dtype, - device_map=None, - low_cpu_mem_usage=True, - **remote_kwargs, - ).eval() - return model_source, tokenizer, model, config - - if task == "causal_lm_logits": - tokenizer = None - tokenizer_errors: list[str] = [] - for source in source_candidates: - try: - tokenizer = AutoTokenizer.from_pretrained(source, **common_kwargs) - break - except Exception as exc: - tokenizer_errors.append(f"{source}: {exc}") - if tokenizer is None: - raise RuntimeError( - f"Could not load tokenizer for {model_id}:\n" - + "\n".join(tokenizer_errors) - ) - if config_model_type in {"lfm2_vl", "qwen3_5", "qwen3_vl"}: - model = AutoModelForImageTextToText.from_pretrained( - model_source, - dtype=torch_dtype, - device_map=None, - low_cpu_mem_usage=True, - **common_kwargs, - ).eval() - tie_note = _tie_lfm2_vl_lm_head_if_needed(model) - if tie_note: - print(f"note={tie_note}") - elif config_model_type in {"qwen3_5", "qwen3_vl"} and isinstance(config.get("text_config"), dict): - model = AutoModelForImageTextToText.from_pretrained( - model_source, - dtype=torch_dtype, - device_map=None, - low_cpu_mem_usage=True, - **common_kwargs, - ).eval() - else: - model = AutoModelForCausalLM.from_pretrained( - model_source, - dtype=torch_dtype, - device_map=None, - low_cpu_mem_usage=True, - **common_kwargs, - ).eval() - return model_source, tokenizer, model, config - if task == "multimodal_causal_lm_logits": - processor = None - processor_errors: list[str] = [] - for source in source_candidates: - try: - processor = AutoProcessor.from_pretrained(source, **common_kwargs) - break - except Exception as exc: - processor_errors.append(f"{source}: {exc}") - if processor is None and config_model_type == "gemma4": - processor = _load_gemma4_processor_fallback( - source_candidates=source_candidates, - common_kwargs=common_kwargs, - ) - if processor is not None: - print("note=using manual gemma4 processor fallback") - if processor is None: - processor_config_hint = "" - if config_model_type == "gemma4": - processor_config_hint = ( - "\n" - "Gemma4 multimodal transpile needs a processor bundle, not just tokenizer/model weights.\n" - "Your local snapshot may be missing files such as `processor_config.json` or modality-specific\n" - "preprocessor configs. Use an official Gemma4 snapshot that includes the processor, or let\n" - "transformers download one by re-running without `--local-files-only`.\n" - ) - raise RuntimeError( - f"Could not load processor for {model_id}:\n" - + "\n".join(processor_errors) - + processor_config_hint - ) - - if config_model_type in {"lfm2_vl", "qwen3_5", "qwen3_vl"}: - model = AutoModelForImageTextToText.from_pretrained( - model_source, - dtype=torch_dtype, - device_map=None, - low_cpu_mem_usage=True, - **common_kwargs, - ).eval() - tie_note = _tie_lfm2_vl_lm_head_if_needed(model) - if tie_note: - print(f"note={tie_note}") - else: - model = AutoModelForCausalLM.from_pretrained( - model_source, - dtype=torch_dtype, - device_map=None, - low_cpu_mem_usage=True, - **common_kwargs, - ).eval() - if config_model_type == "gemma4": - repair_result = _repair_gemma4_checkpoint_weights(model, model_source) - if repair_result.get("applied"): - missing = repair_result.get("missing_keys", []) - unexpected = repair_result.get("unexpected_keys", []) - print( - "note=applied gemma4 checkpoint key remap" - f" missing_after={len(missing)} unexpected_after={len(unexpected)}" - ) - return model_source, processor, model, config - - processor = None - processor_errors: list[str] = [] - missing_optional_audio_dep: str | None = None - for source in source_candidates: - for loader in (AutoProcessor, AutoFeatureExtractor): - try: - processor = loader.from_pretrained(source, **common_kwargs) - break - except Exception as exc: - processor_errors.append(f"{loader.__name__}@{source}: {exc}") - if isinstance(exc, ImportError) and "requires the librosa library" in str(exc): - missing_optional_audio_dep = "librosa" - break - if processor is not None: - break - if missing_optional_audio_dep is not None: - break - if processor is None: - if missing_optional_audio_dep == "librosa": - print("note=falling back to built-in audio preprocessing because the HF feature extractor requires librosa") - else: - print("note=falling back to built-in audio preprocessing because no HF processor/feature extractor was available") - - if task == "ctc_logits": - model_loaders = (AutoModelForCTC, AutoModel) - elif task in {"encoder_hidden_states", "seq2seq_transcription"}: - model_loaders = (AutoModelForSpeechSeq2Seq, AutoModelForSeq2SeqLM, AutoModel) - else: - raise NotImplementedError(f"unsupported generic HF task: {task}") - - load_errors: list[str] = [] - for loader in model_loaders: - try: - model = loader.from_pretrained( - model_source, - dtype=torch_dtype, - device_map=None, - low_cpu_mem_usage=True, - **common_kwargs, - ).eval() - return model_source, processor, model, config - except Exception as exc: - load_errors.append(f"{loader.__name__}: {exc}") - - raise RuntimeError( - f"Could not load model for task={task} from {model_source}.\n" - "\n".join(load_errors) - ) - - -def _load_optional_tokenizer( - *, - model_id: str, - model_source: str, - token: str | None, - trust_remote_code: bool, - local_files_only: bool, -): - try: - from transformers import AutoTokenizer # type: ignore - except Exception: - return None - - source_candidates = [] - for candidate in (model_source, model_id): - if candidate not in source_candidates: - source_candidates.append(candidate) - - common_kwargs: dict[str, object] = { - "trust_remote_code": trust_remote_code, - "local_files_only": local_files_only, - } - if token: - common_kwargs["token"] = token - - for source in source_candidates: - try: - return AutoTokenizer.from_pretrained(source, **common_kwargs) - except Exception: - continue - return None - - -def _ctc_greedy_decode_token_ids(logits: np.ndarray, *, blank_token_id: int | None) -> list[int]: - if logits.ndim != 3 or logits.shape[0] < 1: - raise ValueError(f"expected CTC logits with shape [batch, time, vocab], got {list(logits.shape)}") - - raw_ids = np.argmax(logits[0], axis=-1).tolist() - collapsed: list[int] = [] - previous: int | None = None - for token_id in raw_ids: - if token_id != previous: - collapsed.append(int(token_id)) - previous = int(token_id) - - if blank_token_id is None: - return collapsed - return [token_id for token_id in collapsed if int(token_id) != int(blank_token_id)] - - -def _decode_token_ids(tokenizer: object, token_ids: list[int]) -> str: - decode = getattr(tokenizer, "decode", None) - if not callable(decode): - raise TypeError(f"tokenizer does not expose decode(): {type(tokenizer).__name__}") - try: - return str(decode(token_ids, skip_special_tokens=True)) - except TypeError: - return str(decode(token_ids)) - - -def _count_weight_bindings(ir_graph: IRGraph) -> int: - count = 0 - for value in ir_graph.values.values(): - if isinstance(value.meta, dict) and isinstance(value.meta.get("path"), str): - count += 1 - return count - - -def _lower_preoptimized_ir(ir: IRGraph) -> TranspiledGraph: - verify_ir(ir) - graph = Graph() - graph._transpile_materialized_constants = [] # type: ignore[attr-defined] - env: dict[str, Any] = {} - runtime_inputs = [] - bound_constants = [] - bound_constant_bindings = [] - bound_constant_value_ids: dict[int, str] = {} - - for value_id in ir.inputs: - value = ir.values[value_id] - tensor = _lower_input_value(graph, value) - env[value_id] = tensor - runtime_inputs.append(tensor) - - for value_id, const in ir.constants.items(): - value = ir.values[value_id] - binding = _lookup_weight_binding(value) - if binding is not None: - binding = ensure_binding_compatible(binding, source_tensor=const) - lowered_const = _lower_constant_value(graph, value, const, binding=binding) - env[value_id] = lowered_const - if hasattr(lowered_const, "g") and hasattr(lowered_const, "id"): - bound_constants.append(lowered_const) - bound_constant_value_ids[int(lowered_const.id)] = str(value_id) - if binding is not None: - bound_constant_bindings.append( - { - "node_id": int(lowered_const.id), - "value_id": str(value_id), - "path": binding.path, - "kind": binding.kind, - "source_name": binding.source_name, - } - ) - - for node_id in ir.order: - node = ir.nodes[node_id] - outputs = _lower_ir_node(graph, node, env, ir) - if len(outputs) != len(node.outputs): - raise ValueError( - f"node {node.id} produced {len(outputs)} outputs, expected {len(node.outputs)}" - ) - for output_id, tensor in zip(node.outputs, outputs): - env[output_id] = tensor - - outputs = [env[value_id] for value_id in ir.outputs] - seen_bound_constant_ids = {int(tensor.id) for tensor in bound_constants} - for tensor in getattr(graph, "_transpile_materialized_constants", []): - tensor_id = int(tensor.id) - if tensor_id in seen_bound_constant_ids: - continue - bound_constants.append(tensor) - seen_bound_constant_ids.add(tensor_id) - return TranspiledGraph( - graph=graph, - runtime_inputs=runtime_inputs, - bound_constants=bound_constants, - bound_constant_bindings=bound_constant_bindings, - bound_constant_value_ids=bound_constant_value_ids, - outputs=outputs, - ) - - -def main() -> int: - parser = argparse.ArgumentParser( - description=( - "Load a Hugging Face model, canonicalize it into a generic transpile task, " - "capture it with the Cactus transpiler, lower it to a Cactus Graph, and " - "optionally save artifacts or run the lowered graph." - ) - ) - parser.add_argument( - "--model-id", - required=True, - help="Hugging Face model id or local snapshot path.", - ) - parser.add_argument( - "--task", - default="auto", - choices=( - "auto", - "causal_lm_logits", - "multimodal_causal_lm_logits", - "ctc_logits", - "encoder_hidden_states", - "seq2seq_transcription", - "tdt_transcription", - "text_embedding", - ), - help="Transpile task. Use auto to infer from config/model id.", - ) - parser.add_argument( - "--prompt", - default=_DEFAULT_CAUSAL_PROMPT, - help="Prompt used for causal_lm_logits or multimodal_causal_lm_logits when --input-ids is not set.", - ) - parser.add_argument( - "--system-prompt", - default="", - help="Optional system prompt used for Gemma4 multimodal chat-style prompt construction.", - ) - parser.add_argument( - "--enable-thinking", - action="store_true", - help="Enable model-specific thinking markers when the multimodal prompt format supports them.", - ) - parser.add_argument( - "--input-ids", - default="", - help="Optional comma-separated token ids for causal_lm_logits.", - ) - parser.add_argument( - "--max-new-tokens", - type=int, - default=32, - help=( - "For causal_lm_logits bundles, preallocate prompt-plus-generation context " - "so the saved bundle can do greedy autoregressive decoding." - ), - ) - parser.add_argument( - "--cache-context-length", - default="auto", - help="KV cache context length for cached decode graphs. Use auto to read the model config.", - ) - parser.add_argument( - "--audio-file", - default="", - help="Path to a WAV file for audio or multimodal tasks.", - ) - parser.add_argument( - "--image-file", - action="append", - default=[], - help="Path to an image file for multimodal tasks. Repeat to pass multiple images.", - ) - parser.add_argument( - "--torch-dtype", - default="float16", - help="Torch dtype for model loading: float16, float32, or bfloat16.", - ) - parser.add_argument( - "--token", - default=os.environ.get("HF_TOKEN"), - help="Optional Hugging Face token. Defaults to HF_TOKEN.", - ) - parser.add_argument( - "--trust-remote-code", - action="store_true", - help="Pass trust_remote_code=True to transformers loaders.", - ) - parser.add_argument( - "--local-files-only", - action="store_true", - help="Require the model/processor to already exist locally.", - ) - parser.add_argument( - "--weights-dir", - default="", - help="Converted Cactus CQ weights directory for mmap weight binding.", - ) - parser.add_argument( - "--allow-unconverted-weights", - action="store_true", - help="Debug only: allow transpiling without converted Cactus CQ weights.", - ) - parser.add_argument( - "--artifact-dir", - default="", - help="Directory where raw_ir.json, optimized_ir.json, graph.cactus, component graphs, and result.json are saved. Defaults to ./transpiled//.", - ) - parser.add_argument( - "--graph-filename", - default="graph.cactus", - help="Filename to use for Graph.save() inside --artifact-dir.", - ) - parser.add_argument( - "--component-pipeline", - default="auto", - choices=("auto", "on", "off"), - help="Use a split component pipeline when the model adapter supports it.", - ) - parser.add_argument( - "--components", - default=None, - help=( - "Optional comma-separated component subset for component-pipeline models " - "(for example: vision_encoder,audio_encoder,lm_encoder,decoder)." - ), - ) - parser.add_argument( - "--dynamic-batch", - action="store_true", - help=( - "Mark the Gemma4 decoder_step batch axis dynamic so the emitted graph runs any " - "batch at runtime (for batched/continuous decode). Default off." - ), - ) - parser.add_argument( - "--max-slots", - type=int, - default=1, - help=( - "KV-cache slot-pool capacity (max concurrent sequences) for batched decode. " - "Only used with --dynamic-batch; the decoder_step caches are sized for this many slots." - ), - ) - parser.add_argument( - "--npu", - action="store_true", - help="Also emit CoreML .mlpackage(s) for Apple Neural Engine audio + vision encoders.", - ) - parser.add_argument( - "--npu-quantize", - type=int, - default=None, - choices=[0, 4, 8], - help="Legacy override that forces BOTH audio and vision encoders to the same weight quant (0=fp16, 4=int4, 8=int8). When unset, per-component defaults apply: audio=int8, vision=fp16.", - ) - parser.add_argument( - "--npu-audio-quantize", - type=int, - default=None, - choices=[0, 4, 8], - help="Quantization for the audio encoder .mlpackage (0=fp16, 4=int4, 8=int8). Default int8: Conformer-style encoders absorb int8 quant noise well.", - ) - parser.add_argument( - "--npu-vision-quantize", - type=int, - default=None, - choices=[0, 4, 8], - help="Quantization for the vision encoder .mlpackage (0=fp16, 4=int4, 8=int8). Default fp16: ViT-style towers are small enough that fp16 is the safer correctness default — int4 visibly degrades on Gemma 4 vision.", - ) - parser.add_argument("--no-fuse-gated-deltanet", action="store_true") - parser.add_argument("--no-fuse-rms-norm", action="store_true") - parser.add_argument("--no-fuse-rope", action="store_true") - parser.add_argument("--no-fuse-attention", action="store_true") - parser.add_argument("--no-fuse-attention-block", action="store_true") - parser.add_argument("--no-fuse-add-clipped", action="store_true") - parser.add_argument( - "--skip-execute", - action="store_true", - help="Stop after lowering instead of running the transpiled graph.", - ) - parser.add_argument( - "--skip-reference-compare", - action="store_true", - help="Run the transpiled graph but skip the follow-up PyTorch reference pass.", - ) - args = parser.parse_args() - - validated_weights_dir = _validate_weights_dir(args.weights_dir.strip() or None, model_id=args.model_id) - if validated_weights_dir is None and not args.allow_unconverted_weights: - raise RuntimeError( - "Building the runtime graph requires converted Cactus CQ weights.\n" - "\n" - "Run:\n" - f" cactus convert {args.model_id} --bits 4\n" - "\n" - "For compiler-only debugging, pass --allow-unconverted-weights." - ) - - image_files = tuple(str(path) for path in args.image_file if str(path).strip()) - if args.task == "auto": - inferred_task = _infer_task_from_config(args.model_id) - config_for_auto = _load_config_json(args.model_id) - plan_for_auto = infer_component_plan_from_config(config_for_auto, model_id=args.model_id) - has_multimodal_config = ( - plan_for_auto is not None - and plan_for_auto.task == "multimodal_causal_lm_logits" - ) - if image_files or (args.audio_file and has_multimodal_config and inferred_task == "causal_lm_logits"): - task = plan_for_auto.task if plan_for_auto is not None else "multimodal_causal_lm_logits" - else: - task = inferred_task - else: - task = args.task - torch_dtype = _parse_dtype(args.torch_dtype) - weights_dir = str(validated_weights_dir) if validated_weights_dir is not None else None - artifact_dir = Path(args.artifact_dir).resolve() if args.artifact_dir else _default_artifact_dir_for_model(args.model_id) - - model_source, processor_or_tokenizer, model, model_config = _load_transformers_bundle( - model_id=args.model_id, - task=task, - torch_dtype=torch_dtype, - token=args.token, - trust_remote_code=args.trust_remote_code, - local_files_only=args.local_files_only, - ) - preprocessor_config = _load_optional_json(model_source, "preprocessor_config.json") - if not preprocessor_config: - preprocessor_config = _load_optional_json(args.model_id, "preprocessor_config.json") - auxiliary_tokenizer = None - if task == "ctc_logits": - auxiliary_tokenizer = _load_optional_tokenizer( - model_id=args.model_id, - model_source=model_source, - token=args.token, - trust_remote_code=args.trust_remote_code, - local_files_only=args.local_files_only, - ) - - canonical = None - if task == "causal_lm_logits": - prepared = _prepare_text_inputs( - processor_or_tokenizer, - prompt=args.prompt, - input_ids_text=args.input_ids.strip() or None, - max_new_tokens=int(args.max_new_tokens), - enable_thinking_if_supported=args.enable_thinking, - ) - canonical = canonicalize_model_interface( - model, - task=task, - input_names=prepared.names, - weights_dir=weights_dir, - inputs_metadata=prepared.metadata, - ) - elif task == "text_embedding": - prepared = _prepare_text_embedding_inputs( - processor_or_tokenizer, - prompt=args.prompt, - input_ids_text=args.input_ids.strip() or None, - ) - canonical = canonicalize_model_interface( - model, - task=task, - input_names=prepared.names, - weights_dir=weights_dir, - inputs_metadata=prepared.metadata, - ) - elif task == "multimodal_causal_lm_logits": - config_model_type = str(model_config.get("model_type", "") or "").lower() - if config_model_type == "lfm2_vl": - prepared = _prepare_lfm2_vl_multimodal_inputs( - processor_or_tokenizer, - prompt=args.prompt, - image_files=image_files, - torch_dtype=torch_dtype, - system_prompt=args.system_prompt, - enable_thinking_if_supported=args.enable_thinking, - ) - elif config_model_type in {"qwen3_5", "qwen3_vl"}: - prepared = _prepare_qwen3_5_multimodal_inputs( - processor_or_tokenizer, - prompt=args.prompt, - image_files=image_files, - torch_dtype=torch_dtype, - system_prompt=args.system_prompt, - enable_thinking_if_supported=args.enable_thinking, - ) - else: - prepared = _prepare_gemma4_multimodal_inputs( - processor_or_tokenizer, - prompt=args.prompt, - image_files=image_files, - audio_file=args.audio_file.strip() or None, - torch_dtype=torch_dtype, - system_prompt=args.system_prompt, - enable_thinking_if_supported=args.enable_thinking, - use_gemma4_chat_template=True, - ) - prepared = _add_multimodal_generation_headroom( - prepared, - tokenizer=getattr(processor_or_tokenizer, "tokenizer", processor_or_tokenizer), - max_new_tokens=int(args.max_new_tokens), - min_context_tokens=_multimodal_context_token_floor(config_model_type), - ) - canonical = canonicalize_model_interface( - model, - task=task, - input_names=prepared.names, - weights_dir=weights_dir, - ) - prime_static_features = getattr(canonical.module, "prime_static_multimodal_features", None) - if callable(prime_static_features): - prime_static_features(*prepared.tensors) - elif task == "tdt_transcription": - if not args.audio_file: - raise RuntimeError(f"--audio-file is required for task={task}") - prepared = _prepare_audio_inputs( - processor_or_tokenizer, - input_names=("input_features",), - config=model_config, - preprocessor_config=preprocessor_config, - model=model, - task=task, - audio_file=args.audio_file, - torch_dtype=torch_dtype, - ) - elif task == "seq2seq_transcription": - if not args.audio_file: - raise RuntimeError(f"--audio-file is required for task={task}") - prepared = _prepare_audio_inputs( - processor_or_tokenizer, - input_names=("input_features",), - config=model_config, - preprocessor_config=preprocessor_config, - model=model, - task=task, - audio_file=args.audio_file, - torch_dtype=torch_dtype, - ) - if getattr(model_config, "get", None) is not None or getattr(model, "config", None) is not None: - prepared = _augment_whisper_seq2seq_metadata( - prepared, - tokenizer_or_processor=processor_or_tokenizer, - model=model, - prompt=args.prompt, - max_new_tokens=int(args.max_new_tokens), - ) - else: - if not args.audio_file: - raise RuntimeError(f"--audio-file is required for task={task}") - canonical = canonicalize_model_interface(model, task=task) - prepared = _prepare_audio_inputs( - processor_or_tokenizer, - input_names=canonical.input_names, - config=model_config, - preprocessor_config=preprocessor_config, - model=model, - task=task, - audio_file=args.audio_file, - torch_dtype=torch_dtype, - ) - canonical = canonicalize_model_interface( - model, - task=task, - input_names=prepared.names, - ) - requested_components = None - if args.components: - requested_components = tuple( - component.strip() - for component in str(args.components).split(",") - if component.strip() - ) - plan_derived_components = False - if requested_components is None and task == "causal_lm_logits": - inferred_plan = infer_component_plan_from_config( - _load_config_json(args.model_id), model_id=args.model_id - ) - if ( - inferred_plan is not None - and inferred_plan.task == task - and inferred_plan.force_component_pipeline - and inferred_plan.components - ): - requested_components = tuple(inferred_plan.components) - plan_derived_components = True - spec_kwargs = dict( - task=task, - named_tensors=_named_tensor_store(prepared), - weights_dir=weights_dir, - inputs_metadata=prepared.metadata, - cache_context_length=args.cache_context_length, - dynamic_batch=getattr(args, "dynamic_batch", False), - max_slots=getattr(args, "max_slots", 1), - ) - try: - component_specs = build_component_module_specs( - model, components=requested_components, **spec_kwargs - ) - except UnsupportedComponentsError as exc: - if not plan_derived_components: - raise - print(f"warning: dropping inferred component plan ({exc}); using builder defaults", file=sys.stderr) - component_specs = build_component_module_specs(model, components=None, **spec_kwargs) - use_component_pipeline = False - if args.component_pipeline == "on": - if component_specs is None: - raise RuntimeError( - f"--component-pipeline=on was requested, but {type(model).__name__} does not expose component specs for task={task}" - ) - use_component_pipeline = True - elif args.component_pipeline == "auto" and component_specs is not None: - use_component_pipeline = True - - if use_component_pipeline: - family = canonical.family if canonical is not None else getattr(model, "family", type(model).__name__) - return _run_component_pipeline_transpile( - args=args, - task=task, - family=str(family), - model_source=model_source, - model=model, - prepared=prepared, - component_specs=component_specs, - fusion_config=FusionConfig( - enable_gated_deltanet=not args.no_fuse_gated_deltanet, - enable_rms_norm=not args.no_fuse_rms_norm, - enable_rope=not args.no_fuse_rope, - enable_attention=not args.no_fuse_attention, - enable_attention_block=not args.no_fuse_attention_block, - enable_add_clipped=not args.no_fuse_add_clipped, - ), - weights_dir=weights_dir, - artifact_dir=artifact_dir, - processor_or_tokenizer=processor_or_tokenizer, - canonical=canonical, - ) - - if canonical is None: - raise RuntimeError(f"task={task} requires a component pipeline but none was selected") - wrapper = TranspileWrapper(canonical.module, weights_dir=weights_dir).eval() - - print(f"model_id={args.model_id}") - print(f"model_source={model_source}") - print(f"task={task}") - print(f"adapter_family={canonical.family}") - print(f"adapter_module={type(canonical.module).__name__}") - print(f"input_names={','.join(prepared.names)}") - for name, tensor in zip(prepared.names, prepared.tensors): - print(f"input_{name}_shape={list(tensor.shape)}") - if weights_dir: - print(f"weights_dir={weights_dir}") - - print("capture_begin=true", flush=True) - prepare_cpu_float32_capture = getattr(canonical.module, "prepare_cpu_float32_capture", None) - restore_cpu_float32_capture = getattr(canonical.module, "restore_cpu_float32_capture", None) - if callable(prepare_cpu_float32_capture): - prepare_cpu_float32_capture() - try: - captured = capture_model(wrapper, prepared.tensors) - finally: - if callable(restore_cpu_float32_capture): - restore_cpu_float32_capture() - print("capture_done=true", flush=True) - captured.ir_graph.meta.setdefault("task", task) - captured.ir_graph.meta.setdefault("adapter_family", canonical.family) - raw_ir_graph = copy.deepcopy(captured.ir_graph) - - fusion_config = FusionConfig( - enable_gated_deltanet=not args.no_fuse_gated_deltanet, - enable_rms_norm=not args.no_fuse_rms_norm, - enable_rope=not args.no_fuse_rope, - enable_attention=not args.no_fuse_attention, - enable_attention_block=not args.no_fuse_attention_block, - enable_add_clipped=not args.no_fuse_add_clipped, - ) - - print("canonicalize_begin=true", flush=True) - canonicalize_exported_graph(captured.ir_graph) - print("canonicalize_done=true", flush=True) - print("optimize_begin=true", flush=True) - optimize_graph(captured.ir_graph, config=fusion_config) - print("optimize_done=true", flush=True) - print("lower_begin=true", flush=True) - tg = _lower_preoptimized_ir(captured.ir_graph) - print("lower_done=true", flush=True) - - optimized_ir_graph = copy.deepcopy(captured.ir_graph) - binding_count = _count_weight_bindings(optimized_ir_graph) - op_counts = Counter(optimized_ir_graph.nodes[node_id].op for node_id in optimized_ir_graph.order) - raw_component_counts = summarize_ir_components(raw_ir_graph) - optimized_component_counts = summarize_ir_components(optimized_ir_graph) - raw_component_graphs = extract_component_subgraphs(raw_ir_graph) - optimized_component_graphs = extract_component_subgraphs(optimized_ir_graph) - transpiled_component_graphs: dict[str, TranspiledGraph] | None = None - - print(f"raw_ir_nodes={len(raw_ir_graph.order)}") - print(f"optimized_ir_nodes={len(optimized_ir_graph.order)}") - print(f"weight_bindings={binding_count}") - if raw_component_counts: - print( - "raw_components=" - + ",".join(f"{name}:{count}" for name, count in sorted(raw_component_counts.items())) - ) - if optimized_component_counts: - print( - "optimized_components=" - + ",".join(f"{name}:{count}" for name, count in sorted(optimized_component_counts.items())) - ) - print( - "ops=" - f"attention:{op_counts.get('attention', 0)} " - f"conv1d:{op_counts.get('conv1d', 0)} " - f"conv2d:{op_counts.get('conv2d', 0)} " - f"batch_norm:{op_counts.get('batch_norm', 0)} " - f"layer_norm:{op_counts.get('layer_norm', 0)} " - f"rms_norm:{op_counts.get('rms_norm', 0)} " - f"rope:{op_counts.get('rope', 0)} " - f"linear:{op_counts.get('linear', 0)}" - ) - - if weights_dir and binding_count == 0: - raise RuntimeError( - f"No weight bindings were resolved from {weights_dir}\n" - "\n" - "The weights folder exists, but none of the captured constants matched entries in weights_manifest.json.\n" - "\n" - f"Recommended fix:\n" - f" cactus convert {args.model_id} {weights_dir}\n" - ) - - if artifact_dir is not None: - artifact_dir.mkdir(parents=True, exist_ok=True) - transpiled_component_graphs = { - component: _lower_preoptimized_ir(copy.deepcopy(component_graph)) - for component, component_graph in optimized_component_graphs.items() - } - _write_json( - artifact_dir / "raw_ir.json", - { - "model_id": args.model_id, - "model_source": model_source, - "task": task, - "family": canonical.family, - "inputs": _serialize_json_compatible(prepared.metadata), - "graph": _graph_to_dict(raw_ir_graph), - }, - ) - _write_json( - artifact_dir / "optimized_ir.json", - { - "model_id": args.model_id, - "model_source": model_source, - "task": task, - "family": canonical.family, - "inputs": _serialize_json_compatible(prepared.metadata), - "graph": _graph_to_dict(optimized_ir_graph), - }, - ) - for component, component_graph in raw_component_graphs.items(): - _write_json( - artifact_dir / _component_artifact_name("raw_ir", component), - { - "model_id": args.model_id, - "model_source": model_source, - "task": task, - "family": canonical.family, - "component": component, - "inputs": _serialize_json_compatible(prepared.metadata), - "graph": _graph_to_dict(component_graph), - }, - ) - for component, component_graph in optimized_component_graphs.items(): - _write_json( - artifact_dir / _component_artifact_name("optimized_ir", component), - { - "model_id": args.model_id, - "model_source": model_source, - "task": task, - "family": canonical.family, - "component": component, - "inputs": _serialize_json_compatible(prepared.metadata), - "graph": _graph_to_dict(component_graph), - }, - ) - print(f"saved_raw_ir={artifact_dir / 'raw_ir.json'}") - print(f"saved_optimized_ir={artifact_dir / 'optimized_ir.json'}") - for component in raw_component_graphs: - print(f"saved_raw_component_ir_{component}={artifact_dir / _component_artifact_name('raw_ir', component)}") - for component in optimized_component_graphs: - print( - f"saved_optimized_component_ir_{component}=" - f"{artifact_dir / _component_artifact_name('optimized_ir', component)}" - ) - - component_manifest_path = _write_component_bundle( - artifact_dir=artifact_dir, - model_id=args.model_id, - model_source=model_source, - task=task, - family=canonical.family, - inputs_metadata=prepared.metadata, - raw_component_graphs=raw_component_graphs, - optimized_component_graphs=optimized_component_graphs, - transpiled_component_graphs=transpiled_component_graphs, - graph_filename=args.graph_filename, - ) - print(f"saved_component_bundle_manifest={component_manifest_path}") - for component in optimized_component_graphs: - print( - f"saved_component_graph_{component}=" - f"{artifact_dir / 'components' / component / args.graph_filename}" - ) - graph_path = artifact_dir / args.graph_filename - _embed_materialized_bound_constants(tg) - tg.graph.save(graph_path) - print(f"saved_graph={graph_path}") - graph_binding_manifest_path = _write_graph_binding_manifest( - artifact_dir=artifact_dir, - filename="graph_bindings.json", - model_id=args.model_id, - model_source=model_source, - task=task, - family=canonical.family, - inputs_metadata=prepared.metadata, - transpiled_graph=tg, - ) - print(f"saved_graph_bindings={graph_binding_manifest_path}") - if binding_count > 0: - print( - "note=saved graph structure without embedded mmap bindings; use the saved component/full-graph manifests " - "to rebind weights/embeddings from the weights folder when loading" - ) - - if args.skip_execute: - return 0 - - tg.set_inputs([tensor.cpu().numpy() for tensor in prepared.tensors]) - print("execute_begin=true") - transpiled_output = tg.execute()[0].numpy().astype(np.float32) - print("execute_done=true") - - if args.skip_reference_compare: - result_payload = { - "model_id": args.model_id, - "model_source": model_source, - "task": task, - "family": canonical.family, - "inputs": _serialize_json_compatible(prepared.metadata), - "output_shape": list(transpiled_output.shape), - "raw_ir_nodes": len(raw_ir_graph.order), - "optimized_ir_nodes": len(optimized_ir_graph.order), - "weight_bindings": binding_count, - "reference_compare_skipped": True, - } - print(f"output_shape={list(transpiled_output.shape)}") - if task in {"causal_lm_logits", "multimodal_causal_lm_logits"}: - tokenizer_like = getattr(processor_or_tokenizer, "tokenizer", processor_or_tokenizer) - transpiled_next = int(np.argmax(transpiled_output[0, -1])) - print(f"transpiled_next_token_id={transpiled_next}") - result_payload["transpiled_next_token_id"] = transpiled_next - if hasattr(tokenizer_like, "decode"): - transpiled_token = tokenizer_like.decode([transpiled_next]) - print(f"transpiled_next_token={transpiled_token!r}") - result_payload["transpiled_next_token"] = transpiled_token - if artifact_dir is not None: - _write_json(artifact_dir / "result.json", result_payload) - print(f"saved_result={artifact_dir / 'result.json'}") - del model - del wrapper - gc.collect() - return 0 - - print("reference_begin=true") - with torch.no_grad(): - reference_output = wrapper(*prepared.tensors).detach().float().cpu().numpy() - print("reference_done=true") - - max_abs_diff = float(np.max(np.abs(reference_output - transpiled_output))) - mean_abs_diff = float(np.mean(np.abs(reference_output - transpiled_output))) - result_payload: dict[str, object] = { - "model_id": args.model_id, - "model_source": model_source, - "task": task, - "family": canonical.family, - "inputs": _serialize_json_compatible(prepared.metadata), - "output_shape": list(reference_output.shape), - "raw_ir_nodes": len(raw_ir_graph.order), - "optimized_ir_nodes": len(optimized_ir_graph.order), - "weight_bindings": binding_count, - "max_abs_diff": max_abs_diff, - "mean_abs_diff": mean_abs_diff, - } - - if task == "causal_lm_logits": - tokenizer = processor_or_tokenizer - hf_next = int(np.argmax(reference_output[0, -1])) - transpiled_next = int(np.argmax(transpiled_output[0, -1])) - print(f"hf_next_token_id={hf_next}") - print(f"transpiled_next_token_id={transpiled_next}") - print(f"logits_max_abs_diff={max_abs_diff:.6f}") - print(f"logits_mean_abs_diff={mean_abs_diff:.6f}") - print(f"hf_next_token={tokenizer.decode([hf_next])!r}") - print(f"transpiled_next_token={tokenizer.decode([transpiled_next])!r}") - result_payload.update( - { - "hf_next_token_id": hf_next, - "transpiled_next_token_id": transpiled_next, - "hf_next_token": tokenizer.decode([hf_next]), - "transpiled_next_token": tokenizer.decode([transpiled_next]), - } - ) - elif task == "ctc_logits": - blank_token_id = getattr(model.config, "pad_token_id", None) - if blank_token_id is None and auxiliary_tokenizer is not None: - blank_token_id = getattr(auxiliary_tokenizer, "pad_token_id", None) - print(f"output_shape={list(reference_output.shape)}") - print(f"output_max_abs_diff={max_abs_diff:.6f}") - print(f"output_mean_abs_diff={mean_abs_diff:.6f}") - if auxiliary_tokenizer is not None: - hf_token_ids = _ctc_greedy_decode_token_ids(reference_output, blank_token_id=blank_token_id) - transpiled_token_ids = _ctc_greedy_decode_token_ids(transpiled_output, blank_token_id=blank_token_id) - hf_transcript = _decode_token_ids(auxiliary_tokenizer, hf_token_ids) - transpiled_transcript = _decode_token_ids(auxiliary_tokenizer, transpiled_token_ids) - print(f"hf_transcript={hf_transcript!r}") - print(f"transpiled_transcript={transpiled_transcript!r}") - result_payload.update( - { - "blank_token_id": None if blank_token_id is None else int(blank_token_id), - "hf_transcript_token_ids": hf_token_ids, - "transpiled_transcript_token_ids": transpiled_token_ids, - "hf_transcript": hf_transcript, - "transpiled_transcript": transpiled_transcript, - } - ) - else: - print(f"output_shape={list(reference_output.shape)}") - print(f"output_max_abs_diff={max_abs_diff:.6f}") - print(f"output_mean_abs_diff={mean_abs_diff:.6f}") - - if artifact_dir is not None: - _write_json(artifact_dir / "result.json", result_payload) - print(f"saved_result={artifact_dir / 'result.json'}") - - return 0 - - -if __name__ == "__main__": - raise SystemExit(main()) diff --git a/python/cactus/transpile/import_ir.py b/python/cactus/transpile/import_ir.py deleted file mode 100644 index 19cd05112..000000000 --- a/python/cactus/transpile/import_ir.py +++ /dev/null @@ -1,277 +0,0 @@ -from __future__ import annotations - -from types import SimpleNamespace -from typing import Any - -import torch -from torch.export.graph_signature import ConstantArgument -from torch.export.graph_signature import InputKind - -from cactus.transpile.capture_pytorch import format_target -from cactus.transpile.capture_pytorch import get_dtype -from cactus.transpile.capture_pytorch import get_shape -from cactus.transpile.capture_pytorch import resolve_attr -from cactus.transpile.component_partition import annotate_ir_components -from cactus.transpile.graph_ir import IRGraph -from cactus.transpile.graph_ir import verify_ir -from cactus.transpile.importers import extract_literals -from cactus.transpile.importers import ImportContext -from cactus.transpile.importers import import_call_function -from cactus.transpile.importers import import_get_attr -from cactus.transpile.importers import import_output -from cactus.transpile.importers import import_placeholder -from cactus.transpile.importers import value_id -from cactus.transpile.normalize import dtype_to_ir -from cactus.transpile.import_semantics import apply_import_semantics -from cactus.transpile.weight_binding import resolve_transpile_weights_dir -from cactus.transpile.weight_binding import resolve_weight_binding - - -def import_captured_to_ir(captured: Any, *, strict: bool = True) -> IRGraph: - transpile_metadata = getattr(captured, "transpile_metadata", {}) or {} - graph_meta = transpile_metadata.get("graph", {}) - ir = IRGraph(values={}, nodes={}, order=[], inputs=[], outputs=[], constants={}, meta=dict(graph_meta)) - ctx = ImportContext(strict=strict, transpile_metadata=transpile_metadata) - weights_dir = resolve_transpile_weights_dir(ir.meta) - if weights_dir: - ir.meta["weights_dir"] = weights_dir - placeholder_specs: dict[str, Any] = {} - inline_counter = 0 - - def _try_register_weight_binding(value_id_str: str, target: str, value: Any) -> None: - if not isinstance(value, torch.Tensor): - return - binding = resolve_weight_binding( - weights_dir=weights_dir, - source_name=target, - ) - if binding is None: - return - bindings = ir.meta.setdefault("weight_bindings", {}) - if isinstance(bindings, dict): - bindings[value_id_str] = { - "path": binding.path, - "kind": binding.kind, - "source_name": binding.source_name, - } - if value_id_str in ir.values: - ir.values[value_id_str].meta.update( - { - "path": binding.path, - "kind": binding.kind, - "source_name": binding.source_name, - } - ) - - def _make_inlined_node(node: Any, *, inline_index: int) -> Any: - return SimpleNamespace( - name=f"inl_{inline_index}_{node.name}", - op=node.op, - target=node.target, - args=node.args, - kwargs=getattr(node, "kwargs", {}), - meta=getattr(node, "meta", {}), - ) - - def _resolve_target(target: str) -> Any: - candidates: list[str] = [] - parts = target.split(".") - for i in range(len(parts)): - candidate = ".".join(parts[i:]) - if candidate and candidate not in candidates: - candidates.append(candidate) - - exported_program = getattr(captured, "exported_program", None) - if exported_program is not None: - tensor_constants = getattr(exported_program, "tensor_constants", None) - if tensor_constants is not None: - for candidate in candidates: - if candidate in tensor_constants: - return tensor_constants[candidate] - constants = getattr(exported_program, "constants", None) - if constants is not None: - for candidate in candidates: - if candidate in constants: - return constants[candidate] - - for candidate in candidates: - try: - return resolve_attr(captured.graph_module, candidate) - except AttributeError: - pass - - for candidate in candidates: - if candidate in captured.state_dict: - return captured.state_dict[candidate] - - source_module = getattr(captured, "source_module", None) - if source_module is not None: - for candidate in candidates: - try: - return resolve_attr(source_module, candidate) - except AttributeError: - pass - - raise NotImplementedError(f"could not resolve get_attr target: {target}") from None - - def _collect_output_ids(arg: Any, local_alias: dict[str, str]) -> list[str]: - if hasattr(arg, "name") and hasattr(arg, "op"): - name = arg.name - if name in local_alias: - return [local_alias[name]] - return [ctx.resolve_value_id(arg)] - if isinstance(arg, (tuple, list)): - out: list[str] = [] - for item in arg: - out.extend(_collect_output_ids(item, local_alias)) - return out - return [] - - def _inline_wrap_with_set_grad_enabled(node: Any, shape: Any, dtype: Any) -> None: - nonlocal inline_counter - subgraph_ref = node.args[1] - if not hasattr(subgraph_ref, "op") or subgraph_ref.op != "get_attr": - raise NotImplementedError("wrap_with_set_grad_enabled without get_attr submodule is not supported") - - subgraph_module = _resolve_target(subgraph_ref.target) - if not hasattr(subgraph_module, "graph"): - raise NotImplementedError("wrap_with_set_grad_enabled target is not a GraphModule") - - inline_counter += 1 - inline_index = inline_counter - local_alias: dict[str, str] = {} - subgraph_inputs = list(node.args[2:]) - input_index = 0 - - ctx.push_local_aliases(local_alias) - try: - for sub_node in subgraph_module.graph.nodes: - sub_shape = get_shape(sub_node) - sub_dtype = dtype_to_ir(get_dtype(sub_node)) - sub_desc = ( - f"inlined:{node.name}:{sub_node.op}:{getattr(sub_node, 'name', '')}:" - f"{format_target(sub_node) if sub_node.op == 'call_function' else getattr(sub_node, 'target', '')}" - ) - ctx.push_node(sub_desc) - try: - if sub_node.op == "placeholder": - if input_index >= len(subgraph_inputs): - raise NotImplementedError("subgraph placeholder/input arity mismatch during inlining") - local_alias[sub_node.name] = ctx.resolve_value_id(subgraph_inputs[input_index]) - input_index += 1 - continue - - if sub_node.op == "get_attr": - prefixed = _make_inlined_node(sub_node, inline_index=inline_index) - value = _resolve_target(sub_node.target) - import_get_attr(ir, prefixed, ctx, value, shape=sub_shape, dtype=sub_dtype, source_name=str(sub_node.target)) - _try_register_weight_binding(f"v_{prefixed.name}", str(sub_node.target), value) - local_alias[sub_node.name] = f"v_{prefixed.name}" - continue - - if sub_node.op == "call_function": - prefixed = _make_inlined_node(sub_node, inline_index=inline_index) - import_call_function( - ir, - prefixed, - ctx, - shape=sub_shape, - dtype=sub_dtype, - torch_op=format_target(sub_node), - ) - local_alias[sub_node.name] = f"v_{prefixed.name}" - continue - - if sub_node.op == "output": - ctx.tuple_aliases[node.name] = _collect_output_ids(sub_node.args[0], local_alias) - return - - raise NotImplementedError(f"unsupported inlined subgraph node op: {sub_node.op}") - finally: - ctx.pop_node() - finally: - ctx.pop_local_aliases() - - raise NotImplementedError("inlined subgraph had no output node") - - for spec in captured.exported_program.graph_signature.input_specs: - arg = getattr(spec, "arg", None) - name = getattr(arg, "name", None) - if name is not None: - placeholder_specs[name] = spec - - for node in captured.graph.nodes: - shape = get_shape(node) - dtype = dtype_to_ir(get_dtype(node)) - node_desc = f"top:{node.op}:{getattr(node, 'name', '')}:{getattr(node, 'target', '')}" - ctx.push_node(node_desc) - try: - if node.op == "placeholder": - spec = placeholder_specs.get(node.name) - if spec is not None and isinstance(getattr(spec, "arg", None), ConstantArgument): - import_get_attr( - ir, - node, - ctx, - spec.arg.value, - shape=shape, - dtype=dtype, - source_name=node.name, - ) - continue - if spec is not None and spec.kind in {InputKind.PARAMETER, InputKind.BUFFER, InputKind.CONSTANT_TENSOR}: - target = getattr(spec, "target", None) - if target is None: - raise NotImplementedError(f"placeholder {node.name} has no target for lifted constant input") - value = _resolve_target(target) - import_get_attr(ir, node, ctx, value, shape=shape, dtype=dtype, source_name=str(target)) - _try_register_weight_binding(value_id(node, ctx), str(target), value) - continue - import_placeholder(ir, node, ctx, shape=shape, dtype=dtype) - continue - - if node.op == "get_attr": - if str(node.target).startswith("submod_"): - continue - value = _resolve_target(node.target) - import_get_attr(ir, node, ctx, value, shape=shape, dtype=dtype, source_name=str(node.target)) - _try_register_weight_binding(value_id(node, ctx), str(node.target), value) - continue - - if node.op == "call_function": - torch_op = format_target(node) - if torch_op.endswith(".wrap_with_set_grad_enabled") or torch_op == "wrap_with_set_grad_enabled": - _inline_wrap_with_set_grad_enabled(node, shape, dtype) - continue - if torch_op == "": - source = node.args[0] - if hasattr(source, "name") and source.name in ctx.tuple_aliases: - index_value = extract_literals(node.args[1]) - if not isinstance(index_value, int): - raise NotImplementedError(f"unsupported tuple getitem index: {node.args!r}") - alias_seq = ctx.tuple_aliases[source.name] - if index_value < 0 or index_value >= len(alias_seq): - raise NotImplementedError( - f"tuple getitem index {index_value} out of range for " - f"inlined subgraph {source.name!r} (size {len(alias_seq)})" - ) - ctx.alias_values[node.name] = alias_seq[index_value] - continue - import_call_function(ir, node, ctx, shape=shape, dtype=dtype, torch_op=torch_op) - continue - - if node.op == "call_module" and str(getattr(node, "target", "")) == "_guards_fn" and not node.users: - continue - - if node.op == "output": - import_output(ir, node, ctx) - continue - - raise NotImplementedError(f"unsupported FX node op: {node.op}") - finally: - ctx.pop_node() - - apply_import_semantics(ir) - annotate_ir_components(ir) - verify_ir(ir) - return ir diff --git a/python/cactus/transpile/import_semantics.py b/python/cactus/transpile/import_semantics.py deleted file mode 100644 index 0bf84d3e7..000000000 --- a/python/cactus/transpile/import_semantics.py +++ /dev/null @@ -1,337 +0,0 @@ -from __future__ import annotations - -from dataclasses import dataclass -from typing import Any - -import torch - -from cactus.transpile.fusion.common import producer -from cactus.transpile.fusion.common import strip_passthrough -from cactus.transpile.canonicalize.utils import rebuild_graph -from cactus.transpile.graph_ir import IRGraph -from cactus.transpile.graph_ir import IRNode -from cactus.transpile.graph_ir import verify_ir - - -@dataclass(frozen=True) -class EarlyRoPEMatch: - anchor_node_id: str - input_value_id: str - theta: float - position_offset: int - partial: bool - node_ids: tuple[str, ...] - - -def apply_import_semantics(graph: IRGraph) -> IRGraph: - _rewrite_explicit_attention(graph) - _tag_explicit_linear(graph) - _rewrite_early_rope(graph) - rebuild_graph(graph) - verify_ir(graph) - return graph - - -def _rewrite_explicit_attention(graph: IRGraph) -> None: - for node_id in list(graph.order): - node = graph.nodes.get(node_id) - if node is None or node.op != "scaled_dot_product_attention": - continue - additive_mask = bool(node.attrs.get("additive_mask", False)) - if not additive_mask and len(node.inputs) > 3: - mask_value = graph.values.get(node.inputs[3]) - if mask_value is not None and mask_value.dtype not in (None, "bool"): - additive_mask = True - node.op = "attention" - node.kind = "semantic" - node.meta["semantic_source"] = "import" - node.attrs = { - "scale": float(node.attrs.get("scale", 0.0)), - "is_causal": bool(node.attrs.get("is_causal", False)), - "window_size": int(node.attrs.get("window_size", 0)), - **({"enable_gqa": bool(node.attrs.get("enable_gqa", False))} if "enable_gqa" in node.attrs else {}), - **({"additive_mask": True} if additive_mask else {}), - } - - -def _tag_explicit_linear(graph: IRGraph) -> None: - for node in graph.nodes.values(): - if node.op not in {"linear", "addmm"}: - continue - node.meta.setdefault("semantic_source", "import") - - -def _rewrite_early_rope(graph: IRGraph) -> None: - for node_id in list(graph.order): - node = graph.nodes.get(node_id) - if node is None or not node.outputs: - continue - match = match_early_rope(graph, node.outputs[0]) - if match is None or match.anchor_node_id != node.id: - continue - if match.partial: - node.meta["semantic_hint"] = "partial_rope" - node.meta["rope_theta_hint"] = float(match.theta) - node.meta["rope_position_offset_hint"] = int(match.position_offset) - continue - node.op = "rope" - node.inputs = [match.input_value_id] - node.attrs = { - "theta": float(match.theta), - "position_offset": int(match.position_offset), - } - node.kind = "semantic" - node.meta["semantic_source"] = "import" - - -def match_early_rope(graph: IRGraph, value_id: str) -> EarlyRoPEMatch | None: - desc = _extract_rope_descriptor(graph, value_id) - if desc is None: - return None - return EarlyRoPEMatch( - anchor_node_id=str(desc["anchor_node_id"]), - input_value_id=str(desc["input_value_id"]), - theta=float(desc["theta"]), - position_offset=int(desc["position_offset"]), - partial=bool(desc.get("partial_rope", False)), - node_ids=tuple(sorted(set(desc.get("node_ids", ())))), - ) - - -def _extract_rope_descriptor(graph: IRGraph, value_id: str) -> dict[str, object] | None: - current = _strip_layout_passthrough(graph, value_id) - node = producer(graph, current) - if node is None: - return None - if node.op == "rope": - return { - "anchor_node_id": node.id, - "input_value_id": node.inputs[0], - "theta": float(node.attrs.get("theta", 0.0)), - "position_offset": int(node.attrs.get("position_offset", 0)), - "partial_rope": False, - "node_ids": (node.id,), - } - classic = _extract_classic_rope_descriptor(graph, current) - if classic is not None: - classic["partial_rope"] = False - return classic - if node.op == "cat" and len(node.inputs) == 2: - for rope_branch_value_id, passthrough_value_id in ((node.inputs[0], node.inputs[1]), (node.inputs[1], node.inputs[0])): - rope_branch = _extract_rope_descriptor(graph, rope_branch_value_id) - if rope_branch is None: - continue - passthrough = producer(graph, _strip_layout_passthrough(graph, passthrough_value_id)) - if passthrough is None or passthrough.op != "slice": - continue - if strip_passthrough(graph, passthrough.inputs[0]) != strip_passthrough(graph, rope_branch["input_value_id"]): - continue - rope_branch["anchor_node_id"] = node.id - rope_branch["partial_rope"] = True - rope_branch["node_ids"] = tuple(sorted({node.id, passthrough.id, *rope_branch["node_ids"]})) - return rope_branch - return None - - -def _extract_classic_rope_descriptor(graph: IRGraph, value_id: str) -> dict[str, object] | None: - node = producer(graph, _strip_layout_passthrough(graph, value_id)) - if node is None or node.op != "add" or len(node.inputs) != 2: - return None - for direct_mult_id, rotated_mult_id in ((node.inputs[0], node.inputs[1]), (node.inputs[1], node.inputs[0])): - direct_mult = producer(graph, direct_mult_id) - rotated_mult = producer(graph, rotated_mult_id) - if direct_mult is None or rotated_mult is None: - continue - if direct_mult.op != "multiply" or rotated_mult.op != "multiply": - continue - direct = _extract_direct_branch(graph, direct_mult) - rotated = _extract_rotated_branch(graph, rotated_mult) - if direct is None or rotated is None: - continue - input_value_id = strip_passthrough(graph, direct["input_value_id"]) - if input_value_id != strip_passthrough(graph, rotated["input_value_id"]): - continue - cos_info = _extract_rope_trig(graph, direct["trig_value_id"], expected=("cos", "scalar_cos")) - sin_info = _extract_rope_trig(graph, rotated["trig_value_id"], expected=("sin", "scalar_sin")) - if cos_info is None or sin_info is None: - continue - if abs(float(cos_info["theta"]) - float(sin_info["theta"])) > 1e-2: - continue - if int(cos_info["position_offset"]) != int(sin_info["position_offset"]): - continue - return { - "anchor_node_id": node.id, - "input_value_id": input_value_id, - "theta": float(cos_info["theta"]), - "position_offset": int(cos_info["position_offset"]), - "node_ids": tuple(sorted({node.id, *direct["node_ids"], *rotated["node_ids"], *cos_info["node_ids"], *sin_info["node_ids"]})), - } - return None - - -def _extract_direct_branch(graph: IRGraph, node: IRNode) -> dict[str, object] | None: - for input_value_id, trig_value_id in ((node.inputs[0], node.inputs[1]), (node.inputs[1], node.inputs[0])): - if _unwrap_trig_value(graph, trig_value_id, expected=("cos", "scalar_cos")) is not None: - return {"input_value_id": input_value_id, "trig_value_id": trig_value_id, "node_ids": (node.id,)} - return None - - -def _extract_rotated_branch(graph: IRGraph, node: IRNode) -> dict[str, object] | None: - for rotated_value_id, trig_value_id in ((node.inputs[0], node.inputs[1]), (node.inputs[1], node.inputs[0])): - input_value_id = _extract_rotate_half_source(graph, rotated_value_id) - if input_value_id is None: - continue - if _unwrap_trig_value(graph, trig_value_id, expected=("sin", "scalar_sin")) is None: - continue - return {"input_value_id": input_value_id, "trig_value_id": trig_value_id, "node_ids": (node.id,)} - return None - - -def _extract_rotate_half_source(graph: IRGraph, value_id: str) -> str | None: - cat_node = producer(graph, strip_passthrough(graph, value_id)) - if cat_node is None or cat_node.op != "cat" or len(cat_node.inputs) != 2: - return None - neg_node = producer(graph, cat_node.inputs[0]) - right_slice = producer(graph, cat_node.inputs[1]) - if neg_node is None or right_slice is None: - return None - if neg_node.op != "scalar_multiply" or float(neg_node.attrs.get("value", 0.0)) != -1.0: - return None - left_slice = producer(graph, neg_node.inputs[0]) - if left_slice is None or left_slice.op != "slice" or right_slice.op != "slice": - return None - if strip_passthrough(graph, left_slice.inputs[0]) != strip_passthrough(graph, right_slice.inputs[0]): - return None - return left_slice.inputs[0] - - -def _extract_rope_trig(graph: IRGraph, value_id: str, *, expected: tuple[str, ...]) -> dict[str, object] | None: - trig = _unwrap_trig_value(graph, value_id, expected=expected) - if trig is None: - return None - trig_node = trig["trig_node"] - cat_node = producer(graph, trig_node.inputs[0]) - if cat_node is None or cat_node.op != "cat" or len(cat_node.inputs) != 2 or cat_node.inputs[0] != cat_node.inputs[1]: - return None - angle_info = _extract_angle_source(graph, cat_node.inputs[0]) - if angle_info is None: - return None - matmul_node = angle_info["matmul_node"] - inv_freq_const_id = None - arange_node = None - for input_id in matmul_node.inputs: - if inv_freq_const_id is None: - inv_freq_const_id = _find_constant_ancestor(graph, input_id) - if arange_node is None: - arange_node = _find_arange_ancestor(graph, input_id) - if inv_freq_const_id is None or arange_node is None: - return None - theta = _infer_rope_theta(graph.constants[inv_freq_const_id]) - if theta is None: - return None - return { - "theta": theta, - "position_offset": int(arange_node.attrs.get("start", 0)), - "node_ids": tuple(sorted({trig_node.id, cat_node.id, matmul_node.id})), - } - - -def _unwrap_trig_value(graph: IRGraph, value_id: str, *, expected: tuple[str, ...]) -> dict[str, object] | None: - current = value_id - while True: - node = producer(graph, current) - if node is None: - return None - if node.op in {"unsqueeze", "reshape", "view", "expand", "precision_cast"}: - current = node.inputs[0] - continue - if node.op == "type_as": - current = node.inputs[0] - continue - if node.op == "scalar_multiply": - current = node.inputs[0] - continue - if node.op not in expected: - return None - return {"trig_node": node} - - -def _extract_angle_source(graph: IRGraph, value_id: str) -> dict[str, object] | None: - current = value_id - while True: - current = _strip_layout_passthrough(graph, current) - node = producer(graph, current) - if node is None: - return None - if node.op in {"transpose", "permute"}: - matmul_node = producer(graph, node.inputs[0]) - if matmul_node is None or matmul_node.op != "matmul": - return None - return {"matmul_node": matmul_node} - if node.op == "matmul": - return {"matmul_node": node} - if node.op in {"index", "slice"} and len(node.inputs) == 1: - current = node.inputs[0] - continue - return None - - -def _find_constant_ancestor(graph: IRGraph, value_id: str) -> str | None: - current = value_id - visited: set[str] = set() - while current not in visited: - visited.add(current) - current = _strip_layout_passthrough(graph, current) - if current in graph.constants and isinstance(graph.constants[current], torch.Tensor): - return current - node = producer(graph, current) - if node is None or len(node.inputs) != 1: - return None - current = node.inputs[0] - return None - - -def _find_arange_ancestor(graph: IRGraph, value_id: str) -> IRNode | None: - current = value_id - visited: set[str] = set() - while current not in visited: - visited.add(current) - current = _strip_layout_passthrough(graph, current) - node = producer(graph, current) - if node is None: - return None - if node.op == "arange": - return node - if current in graph.constants and isinstance(graph.constants[current], torch.Tensor): - return None - if len(node.inputs) != 1: - return None - current = node.inputs[0] - return None - - -def _strip_layout_passthrough(graph: IRGraph, value_id: str) -> str: - current = value_id - while True: - node = producer(graph, current) - if node is None: - return current - if node.op in {"precision_cast", "contiguous", "reshape", "view", "unsqueeze", "expand"} and len(node.inputs) == 1: - current = node.inputs[0] - continue - if node.op == "type_as" and len(node.inputs) >= 1: - current = node.inputs[0] - continue - return current - - -def _infer_rope_theta(value: Any) -> float | None: - if not isinstance(value, torch.Tensor): - return None - flat = value.detach().cpu().float().reshape(-1) - if flat.numel() < 2: - return None - second = float(flat[1].item()) - if second <= 0.0: - return None - return float((1.0 / second) ** flat.numel()) diff --git a/python/cactus/transpile/importers.py b/python/cactus/transpile/importers.py deleted file mode 100644 index 495b1dfbd..000000000 --- a/python/cactus/transpile/importers.py +++ /dev/null @@ -1,1869 +0,0 @@ -from __future__ import annotations - -import math -from dataclasses import dataclass -from dataclasses import field -from types import SimpleNamespace -from typing import Any - -import torch - -from cactus.transpile.aten_ops import canonical_torch_op -from cactus.transpile.graph_ir import IRGraph -from cactus.transpile.graph_ir import IRNode -from cactus.transpile.graph_ir import IRValue -from cactus.transpile.normalize import dtype_to_ir -from cactus.transpile.normalize import normalize_target - - -class UnsupportedImportError(NotImplementedError): - pass - - -@dataclass -class ImportContext: - strict: bool = True - alias_values: dict[str, str] = field(default_factory=dict) - tuple_aliases: dict[str, list[str]] = field(default_factory=dict) - local_alias_stack: list[dict[str, str]] = field(default_factory=list) - node_stack: list[str] = field(default_factory=list) - transpile_metadata: dict[str, object] = field(default_factory=dict) - - def push_local_aliases(self, aliases: dict[str, str]) -> None: - self.local_alias_stack.append(aliases) - - def pop_local_aliases(self) -> None: - if not self.local_alias_stack: - raise RuntimeError("local alias stack underflow") - self.local_alias_stack.pop() - - def resolve_value_id(self, node: Any) -> str: - if is_fx_node(node): - name = getattr(node, "name", "") - for aliases in reversed(self.local_alias_stack): - if name in aliases: - return aliases[name] - resolved = self.alias_values.get(name) - if resolved is not None: - return resolved - return f"v_{name}" - raise TypeError(f"cannot resolve value id for non-FX node: {type(node).__name__}") - - def fail(self, message: str) -> None: - if self.node_stack: - message = f"{message} [context: {' > '.join(self.node_stack)}]" - if self.strict: - raise UnsupportedImportError(message) - raise UnsupportedImportError(f"{message} [non-strict import has no generic fallback]") - - def push_node(self, description: str) -> None: - self.node_stack.append(description) - - def pop_node(self) -> None: - if self.node_stack: - self.node_stack.pop() - - -def value_id(node: Any, ctx: ImportContext) -> str: - return ctx.resolve_value_id(node) - - -def node_id(node: Any) -> str: - return f"n_{node.name}" - - -def is_fx_node(value: Any) -> bool: - return hasattr(value, "name") and hasattr(value, "op") - - -def extract_input_ids(arg: Any, ctx: ImportContext) -> list[str]: - if is_fx_node(arg): - return [value_id(arg, ctx)] - if isinstance(arg, (tuple, list)): - ids: list[str] = [] - for item in arg: - ids.extend(extract_input_ids(item, ctx)) - return ids - if isinstance(arg, dict): - ids: list[str] = [] - for item in arg.values(): - ids.extend(extract_input_ids(item, ctx)) - return ids - return [] - - -def extract_literals(value: Any) -> Any: - if is_fx_node(value): - meta = getattr(value, "meta", {}) or {} - meta_value = meta.get("val") - if meta_value is not None and not isinstance(meta_value, torch.Tensor): - return extract_literals(meta_value) - return None - if isinstance(value, torch.Size): - return tuple(extract_literals(v) for v in value) - if isinstance(value, (tuple, list)): - return type(value)(extract_literals(v) for v in value) - if isinstance(value, dict): - return {k: extract_literals(v) for k, v in value.items()} - if isinstance(value, torch.dtype): - return str(value) - if isinstance(value, (int, float, bool, str)) or value is None: - return value - return None - - -def make_prefixed_node(node: Any, prefix: str) -> Any: - return SimpleNamespace( - name=f"{prefix}{node.name}", - op=node.op, - target=node.target, - args=node.args, - kwargs=getattr(node, "kwargs", {}), - meta=getattr(node, "meta", {}), - ) - - -def add_value_if_missing(ir: IRGraph, value: IRValue) -> None: - if value.id not in ir.values: - ir.add_value(value) - - -def register_node( - ir: IRGraph, - ir_node: IRNode, - *, - shape: tuple[int, ...] | None, - dtype: str | None, -) -> None: - ir.add_node(ir_node) - ir.order.append(ir_node.id) - for output_id in ir_node.outputs: - ir.values[output_id].shape = shape - ir.values[output_id].dtype = dtype - for input_id in ir_node.inputs: - if input_id in ir.values: - ir.values[input_id].users.append(ir_node.id) - - -def _extract_numeric_literal(value: Any) -> float | int | None: - literal = extract_literals(value) - if isinstance(literal, bool): - return int(literal) - if isinstance(literal, (int, float)): - return literal - return None - - -def _extract_static_int(value: Any) -> int | None: - if is_fx_node(value): - meta = getattr(value, "meta", {}) or {} - meta_value = meta.get("val") - if meta_value is not None and not isinstance(meta_value, torch.Tensor): - return _extract_static_int(meta_value) - return None - - literal = extract_literals(value) - if isinstance(literal, bool): - return int(literal) - if isinstance(literal, int): - return int(literal) - if isinstance(literal, float) and float(literal).is_integer(): - return int(literal) - - try: - return int(value) - except Exception: - return None - - -def _extract_int_sequence_literal(value: Any) -> tuple[int, ...] | None: - if is_fx_node(value): - meta = getattr(value, "meta", {}) or {} - meta_value = meta.get("val") - if meta_value is not None and not isinstance(meta_value, torch.Tensor): - return _extract_int_sequence_literal(meta_value) - return None - - if isinstance(value, torch.Size): - value = tuple(value) - - if isinstance(value, (tuple, list)): - result: list[int] = [] - for item in value: - item_value = _extract_static_int(item) - if item_value is None: - return None - result.append(item_value) - return tuple(result) - - scalar_value = _extract_static_int(value) - if scalar_value is not None: - return (scalar_value,) - return None - - -def _try_materialize_shape(shape: tuple[Any, ...] | None) -> tuple[int, ...] | None: - if shape is None: - return None - result: list[int] = [] - for dim in shape: - dim_value = _extract_static_int(dim) - if dim_value is None: - return None - result.append(dim_value) - return tuple(result) - - -def _extract_fx_tensor_shape(value: Any) -> tuple[Any, ...] | None: - if not is_fx_node(value): - return None - meta = getattr(value, "meta", {}) or {} - meta_value = meta.get("val") - if isinstance(meta_value, torch.Tensor): - return tuple(meta_value.shape) - tensor_meta = meta.get("tensor_meta") - tensor_shape = getattr(tensor_meta, "shape", None) - if tensor_shape is not None: - return tuple(tensor_shape) - return None - - -def _base_meta(shape: tuple[int, ...] | None, dtype: str | None, torch_op: str, node: Any) -> dict[str, object]: - aten_op = canonical_torch_op(torch_op) - return { - "shape": shape, - "dtype": dtype, - "torch_op": aten_op, - "aten_op": aten_op, - "torch_name": node.name, - } - - -def _extract_fx_torch_dtype(node: Any) -> torch.dtype | None: - if not is_fx_node(node): - return None - meta = getattr(node, "meta", {}) or {} - value = meta.get("val") - dtype = getattr(value, "dtype", None) - if isinstance(dtype, torch.dtype): - return dtype - tensor_meta = meta.get("tensor_meta") - dtype = getattr(tensor_meta, "dtype", None) - if isinstance(dtype, torch.dtype): - return dtype - return None - - -def _import_scalar_binary( - ir: IRGraph, - node: Any, - ctx: ImportContext, - *, - shape: tuple[int, ...] | None, - dtype: str | None, - torch_op: str, - op_name: str, -) -> None: - lhs, rhs = node.args[0], node.args[1] - lhs_is_node = is_fx_node(lhs) - rhs_is_node = is_fx_node(rhs) - lhs_literal = _extract_numeric_literal(lhs) - rhs_literal = _extract_numeric_literal(rhs) - - if lhs_is_node and rhs_is_node: - ir_node = IRNode( - id=node_id(node), - op=op_name, - inputs=[value_id(lhs, ctx), value_id(rhs, ctx)], - outputs=[value_id(node, ctx)], - meta=_base_meta(shape, dtype, torch_op, node), - ) - register_node(ir, ir_node, shape=shape, dtype=dtype) - return - - if op_name == "add": - if lhs_is_node and rhs_literal is not None: - scalar_op = "scalar_add" - inputs = [value_id(lhs, ctx)] - attrs = {"value": rhs_literal} - elif rhs_is_node and lhs_literal is not None: - scalar_op = "scalar_add" - inputs = [value_id(rhs, ctx)] - attrs = {"value": lhs_literal} - else: - scalar_op = None - elif op_name == "multiply": - if lhs_is_node and rhs_literal is not None: - scalar_op = "scalar_multiply" - inputs = [value_id(lhs, ctx)] - attrs = {"value": rhs_literal} - elif rhs_is_node and lhs_literal is not None: - scalar_op = "scalar_multiply" - inputs = [value_id(rhs, ctx)] - attrs = {"value": lhs_literal} - else: - scalar_op = None - elif op_name == "subtract": - if lhs_is_node and rhs_literal is not None: - scalar_op = "scalar_subtract" - inputs = [value_id(lhs, ctx)] - attrs = {"value": rhs_literal} - elif rhs_is_node and lhs_literal is not None: - scalar_op = "scalar_subtract_reverse" - inputs = [value_id(rhs, ctx)] - attrs = {"value": lhs_literal} - else: - scalar_op = None - elif op_name == "divide": - if lhs_is_node and rhs_literal is not None: - scalar_op = "scalar_divide" - inputs = [value_id(lhs, ctx)] - attrs = {"value": rhs_literal} - elif rhs_is_node and lhs_literal is not None: - scalar_op = "scalar_divide_reverse" - inputs = [value_id(rhs, ctx)] - attrs = {"value": lhs_literal} - else: - scalar_op = None - else: - scalar_op = None - - if scalar_op is None: - ctx.fail(f"unsupported scalar form for {torch_op}: {node.args!r}") - - ir_node = IRNode( - id=node_id(node), - op=scalar_op, - inputs=inputs, - outputs=[value_id(node, ctx)], - attrs=attrs, - meta=_base_meta(shape, dtype, torch_op, node), - ) - register_node(ir, ir_node, shape=shape, dtype=dtype) - - -def import_placeholder( - ir: IRGraph, - node: Any, - ctx: ImportContext, - *, - shape: tuple[int, ...] | None, - dtype: str | None, -) -> None: - add_value_if_missing(ir, IRValue(id=value_id(node, ctx), shape=shape, dtype=dtype, producer=None)) - ir.inputs.append(value_id(node, ctx)) - - -def import_get_attr( - ir: IRGraph, - node: Any, - ctx: ImportContext, - value: Any, - *, - shape: tuple[int, ...] | None, - dtype: str | None, - source_name: str | None = None, -) -> None: - add_value_if_missing(ir, IRValue(id=value_id(node, ctx), shape=shape, dtype=dtype, producer=None)) - ir.constants[value_id(node, ctx)] = value.detach().cpu() if isinstance(value, torch.Tensor) else value - ir.values[value_id(node, ctx)].meta["source_name"] = source_name - if source_name is not None: - bindings = ir.meta.setdefault("weight_bindings", {}) - if isinstance(bindings, dict): - bindings.setdefault(value_id(node, ctx), {})["source_name"] = source_name - - -def import_output(ir: IRGraph, node: Any, ctx: ImportContext) -> None: - ir.outputs.extend(extract_input_ids(node.args[0], ctx)) - - -def import_call_function( - ir: IRGraph, - node: Any, - ctx: ImportContext, - *, - shape: tuple[int, ...] | None, - dtype: str | None, - torch_op: str, -) -> None: - op = normalize_target(torch_op) - importer = OP_IMPORTERS.get(op) - if importer is None: - import_opaque_call_function(ir, node, ctx, shape=shape, dtype=dtype, torch_op=torch_op, op_name=op) - else: - importer(ir, node, ctx, shape=shape, dtype=dtype, torch_op=torch_op) - - -def import_opaque_call_function( - ir: IRGraph, - node: Any, - ctx: ImportContext, - *, - shape: tuple[int, ...] | None, - dtype: str | None, - torch_op: str, - op_name: str, -) -> None: - ir_node = IRNode( - id=node_id(node), - op=op_name, - inputs=extract_input_ids(getattr(node, "args", ()), ctx) + extract_input_ids(getattr(node, "kwargs", {}), ctx), - outputs=[value_id(node, ctx)], - attrs={ - "opaque": True, - "torch_op": torch_op, - "args": extract_literals(getattr(node, "args", ())), - "kwargs": extract_literals(getattr(node, "kwargs", {})), - }, - meta=_base_meta(shape, dtype, torch_op, node), - kind="opaque", - ) - register_node(ir, ir_node, shape=shape, dtype=dtype) - - -def import_add(ir: IRGraph, node: Any, ctx: ImportContext, *, shape: tuple[int, ...] | None, dtype: str | None, torch_op: str) -> None: - _import_scalar_binary(ir, node, ctx, shape=shape, dtype=dtype, torch_op=torch_op, op_name="add") - - -def import_subtract(ir: IRGraph, node: Any, ctx: ImportContext, *, shape: tuple[int, ...] | None, dtype: str | None, torch_op: str) -> None: - _import_scalar_binary(ir, node, ctx, shape=shape, dtype=dtype, torch_op=torch_op, op_name="subtract") - - -def import_multiply(ir: IRGraph, node: Any, ctx: ImportContext, *, shape: tuple[int, ...] | None, dtype: str | None, torch_op: str) -> None: - _import_scalar_binary(ir, node, ctx, shape=shape, dtype=dtype, torch_op=torch_op, op_name="multiply") - - -def import_multiply_inplace(ir: IRGraph, node: Any, ctx: ImportContext, *, shape: tuple[int, ...] | None, dtype: str | None, torch_op: str) -> None: - _import_scalar_binary(ir, node, ctx, shape=shape, dtype=dtype, torch_op=torch_op, op_name="multiply") - - -def import_divide(ir: IRGraph, node: Any, ctx: ImportContext, *, shape: tuple[int, ...] | None, dtype: str | None, torch_op: str) -> None: - _import_scalar_binary(ir, node, ctx, shape=shape, dtype=dtype, torch_op=torch_op, op_name="divide") - - -def import_not_equal(ir: IRGraph, node: Any, ctx: ImportContext, *, shape: tuple[int, ...] | None, dtype: str | None, torch_op: str) -> None: - lhs, rhs = node.args[0], node.args[1] - lhs_is_node = is_fx_node(lhs) - rhs_is_node = is_fx_node(rhs) - lhs_literal = _extract_numeric_literal(lhs) - rhs_literal = _extract_numeric_literal(rhs) - - if lhs_is_node and rhs_is_node: - ir_node = IRNode( - id=node_id(node), - op="not_equal", - inputs=[value_id(lhs, ctx), value_id(rhs, ctx)], - outputs=[value_id(node, ctx)], - meta=_base_meta(shape, dtype, torch_op, node), - ) - register_node(ir, ir_node, shape=shape, dtype=dtype) - return - - if lhs_is_node and rhs_literal is not None: - inputs = [value_id(lhs, ctx)] - attrs = {"value": rhs_literal} - elif rhs_is_node and lhs_literal is not None: - inputs = [value_id(rhs, ctx)] - attrs = {"value": lhs_literal} - else: - ctx.fail(f"unsupported scalar form for {torch_op}: {node.args!r}") - - ir_node = IRNode( - id=node_id(node), - op="scalar_not_equal", - inputs=inputs, - outputs=[value_id(node, ctx)], - attrs=attrs, - meta=_base_meta(shape, dtype, torch_op, node), - ) - register_node(ir, ir_node, shape=shape, dtype=dtype) - - -def import_equal(ir: IRGraph, node: Any, ctx: ImportContext, *, shape: tuple[int, ...] | None, dtype: str | None, torch_op: str) -> None: - lhs, rhs = node.args[0], node.args[1] - lhs_is_node = is_fx_node(lhs) - rhs_is_node = is_fx_node(rhs) - lhs_literal = _extract_numeric_literal(lhs) - rhs_literal = _extract_numeric_literal(rhs) - - if lhs_is_node and rhs_is_node: - ir_node = IRNode( - id=node_id(node), - op="equal", - inputs=[value_id(lhs, ctx), value_id(rhs, ctx)], - outputs=[value_id(node, ctx)], - meta=_base_meta(shape, dtype, torch_op, node), - ) - register_node(ir, ir_node, shape=shape, dtype=dtype) - return - - if lhs_is_node and rhs_literal is not None: - inputs = [value_id(lhs, ctx)] - attrs = {"value": rhs_literal} - elif rhs_is_node and lhs_literal is not None: - inputs = [value_id(rhs, ctx)] - attrs = {"value": lhs_literal} - else: - ctx.fail(f"unsupported scalar form for {torch_op}: {node.args!r}") - - ir_node = IRNode( - id=node_id(node), - op="scalar_equal", - inputs=inputs, - outputs=[value_id(node, ctx)], - attrs=attrs, - meta=_base_meta(shape, dtype, torch_op, node), - ) - register_node(ir, ir_node, shape=shape, dtype=dtype) - - -def import_compare(ir: IRGraph, node: Any, ctx: ImportContext, *, shape: tuple[int, ...] | None, dtype: str | None, torch_op: str, op_name: str) -> None: - lhs, rhs = node.args[0], node.args[1] - lhs_is_node = is_fx_node(lhs) - rhs_is_node = is_fx_node(rhs) - lhs_literal = _extract_numeric_literal(lhs) - rhs_literal = _extract_numeric_literal(rhs) - - if lhs_is_node and rhs_is_node: - ir_node = IRNode( - id=node_id(node), - op=op_name, - inputs=[value_id(lhs, ctx), value_id(rhs, ctx)], - outputs=[value_id(node, ctx)], - meta=_base_meta(shape, dtype, torch_op, node), - ) - register_node(ir, ir_node, shape=shape, dtype=dtype) - return - - if lhs_is_node and rhs_literal is not None: - inputs = [value_id(lhs, ctx)] - scalar_op = f"scalar_{op_name}" - scalar_value = rhs_literal - elif rhs_is_node and lhs_literal is not None: - reverse_ops = { - "greater": "less", - "greater_equal": "less_equal", - "less": "greater", - "less_equal": "greater_equal", - } - inputs = [value_id(rhs, ctx)] - scalar_op = f"scalar_{reverse_ops[op_name]}" - scalar_value = lhs_literal - else: - ctx.fail(f"unsupported scalar form for {torch_op}: {node.args!r}") - - ir_node = IRNode( - id=node_id(node), - op=scalar_op, - inputs=inputs, - outputs=[value_id(node, ctx)], - attrs={"value": scalar_value}, - meta=_base_meta(shape, dtype, torch_op, node), - ) - register_node(ir, ir_node, shape=shape, dtype=dtype) - - -def import_binary(op_name: str): - def _importer( - ir: IRGraph, - node: Any, - ctx: ImportContext, - *, - shape: tuple[int, ...] | None, - dtype: str | None, - torch_op: str, - ) -> None: - ir_node = IRNode( - id=node_id(node), - op=op_name, - inputs=[value_id(node.args[0], ctx), value_id(node.args[1], ctx)], - outputs=[value_id(node, ctx)], - meta=_base_meta(shape, dtype, torch_op, node), - ) - register_node(ir, ir_node, shape=shape, dtype=dtype) - - return _importer - - -def import_where( - ir: IRGraph, - node: Any, - ctx: ImportContext, - *, - shape: tuple[int, ...] | None, - dtype: str | None, - torch_op: str, -) -> None: - condition = node.args[0] - true_value = node.args[1] - false_value = node.args[2] - - inputs = [value_id(condition, ctx)] - attrs: dict[str, object] = {} - - if is_fx_node(true_value): - inputs.append(value_id(true_value, ctx)) - attrs["true_is_scalar"] = False - else: - literal = _extract_numeric_literal(true_value) - if literal is None: - ctx.fail(f"unsupported true branch for {torch_op}: {node.args!r}") - attrs["true_is_scalar"] = True - attrs["true_value"] = float(literal) - - if is_fx_node(false_value): - inputs.append(value_id(false_value, ctx)) - attrs["false_is_scalar"] = False - else: - literal = _extract_numeric_literal(false_value) - if literal is None: - ctx.fail(f"unsupported false branch for {torch_op}: {node.args!r}") - attrs["false_is_scalar"] = True - attrs["false_value"] = float(literal) - - ir_node = IRNode( - id=node_id(node), - op="where", - inputs=inputs, - outputs=[value_id(node, ctx)], - attrs=attrs, - meta=_base_meta(shape, dtype, torch_op, node), - ) - register_node(ir, ir_node, shape=shape, dtype=dtype) - - -def import_negate(ir: IRGraph, node: Any, ctx: ImportContext, *, shape: tuple[int, ...] | None, dtype: str | None, torch_op: str) -> None: - ir_node = IRNode( - id=node_id(node), - op="scalar_multiply", - inputs=[value_id(node.args[0], ctx)], - outputs=[value_id(node, ctx)], - attrs={"value": -1.0}, - meta=_base_meta(shape, dtype, torch_op, node), - ) - register_node(ir, ir_node, shape=shape, dtype=dtype) - - -def import_precision_cast(ir: IRGraph, node: Any, ctx: ImportContext, *, shape: tuple[int, ...] | None, dtype: str | None, torch_op: str) -> None: - target_dtype = extract_literals(node.args[1]) if len(node.args) > 1 else None - ir_node = IRNode( - id=node_id(node), - op="precision_cast", - inputs=[value_id(node.args[0], ctx)], - outputs=[value_id(node, ctx)], - attrs={"dtype": target_dtype}, - meta=_base_meta(shape, dtype, torch_op, node), - ) - register_node(ir, ir_node, shape=shape, dtype=dtype) - - -def import_arange(ir: IRGraph, node: Any, ctx: ImportContext, *, shape: tuple[int, ...] | None, dtype: str | None, torch_op: str) -> None: - attrs: dict[str, object] = {} - if len(node.args) == 1: - attrs["start"] = 0 - attrs["end"] = int(node.args[0]) - elif len(node.args) >= 2: - attrs["start"] = int(node.args[0]) - attrs["end"] = int(node.args[1]) - else: - ctx.fail(f"unsupported arange signature for {torch_op}: {node.args!r}") - - if len(node.args) > 2 and extract_literals(node.args[2]) is not None: - attrs["step"] = int(node.args[2]) - - target_dtype = None - if "dtype" in getattr(node, "kwargs", {}) and extract_literals(node.kwargs["dtype"]) is not None: - target_dtype = extract_literals(node.kwargs["dtype"]) - elif dtype is not None: - target_dtype = dtype - if target_dtype is not None: - attrs["dtype"] = target_dtype - - ir_node = IRNode( - id=node_id(node), - op="arange", - inputs=[], - outputs=[value_id(node, ctx)], - attrs=attrs, - meta=_base_meta(shape, dtype, torch_op, node), - ) - register_node(ir, ir_node, shape=shape, dtype=dtype) - - -def import_type_as(ir: IRGraph, node: Any, ctx: ImportContext, *, shape: tuple[int, ...] | None, dtype: str | None, torch_op: str) -> None: - ir_node = IRNode( - id=node_id(node), - op="type_as", - inputs=[value_id(node.args[0], ctx), value_id(node.args[1], ctx)], - outputs=[value_id(node, ctx)], - meta=_base_meta(shape, dtype, torch_op, node), - ) - register_node(ir, ir_node, shape=shape, dtype=dtype) - - -def import_identity(ir: IRGraph, node: Any, ctx: ImportContext, *, shape: tuple[int, ...] | None, dtype: str | None, torch_op: str) -> None: - ir_node = IRNode( - id=node_id(node), - op="identity", - inputs=[value_id(node.args[0], ctx)], - outputs=[value_id(node, ctx)], - meta=_base_meta(shape, dtype, torch_op, node), - ) - register_node(ir, ir_node, shape=shape, dtype=dtype) - - -def import_reshape(ir: IRGraph, node: Any, ctx: ImportContext, *, shape: tuple[int, ...] | None, dtype: str | None, torch_op: str) -> None: - target_shape = _extract_int_sequence_literal(node.args[1]) - if target_shape is None: - ctx.fail(f"unsupported reshape shape for {torch_op}: {node.args!r}") - ir_node = IRNode( - id=node_id(node), - op="reshape", - inputs=[value_id(node.args[0], ctx)], - outputs=[value_id(node, ctx)], - attrs={"shape": target_shape}, - meta=_base_meta(shape, dtype, torch_op, node), - ) - register_node(ir, ir_node, shape=shape, dtype=dtype) - - -def import_flatten(ir: IRGraph, node: Any, ctx: ImportContext, *, shape: tuple[int, ...] | None, dtype: str | None, torch_op: str) -> None: - ir_node = IRNode( - id=node_id(node), - op="flatten", - inputs=[value_id(node.args[0], ctx)], - outputs=[value_id(node, ctx)], - attrs={"start_dim": int(node.args[1]), "end_dim": int(node.args[2])}, - meta=_base_meta(shape, dtype, torch_op, node), - ) - register_node(ir, ir_node, shape=shape, dtype=dtype) - - -def import_unsqueeze(ir: IRGraph, node: Any, ctx: ImportContext, *, shape: tuple[int, ...] | None, dtype: str | None, torch_op: str) -> None: - dim = int(node.args[1]) - attrs: dict[str, object] = {"dim": dim} - materialized_shape = _try_materialize_shape(shape) - if materialized_shape is not None: - attrs["shape"] = materialized_shape - ir_node = IRNode( - id=node_id(node), - op="unsqueeze", - inputs=[value_id(node.args[0], ctx)], - outputs=[value_id(node, ctx)], - attrs=attrs, - meta=_base_meta(shape, dtype, torch_op, node), - ) - register_node(ir, ir_node, shape=shape, dtype=dtype) - - -def import_squeeze(ir: IRGraph, node: Any, ctx: ImportContext, *, shape: tuple[int, ...] | None, dtype: str | None, torch_op: str) -> None: - dim = int(node.args[1]) if len(node.args) > 1 and node.args[1] is not None else 0 - attrs: dict[str, object] = {"dim": dim} - materialized_shape = _try_materialize_shape(shape) - if materialized_shape is not None: - attrs["shape"] = materialized_shape - ir_node = IRNode( - id=node_id(node), - op="squeeze", - inputs=[value_id(node.args[0], ctx)], - outputs=[value_id(node, ctx)], - attrs=attrs, - meta=_base_meta(shape, dtype, torch_op, node), - ) - register_node(ir, ir_node, shape=shape, dtype=dtype) - - -def import_masked_scatter( - ir: IRGraph, - node: Any, - ctx: ImportContext, - *, - shape: tuple[int, ...] | None, - dtype: str | None, - torch_op: str, -) -> None: - ir_node = IRNode( - id=node_id(node), - op="masked_scatter", - inputs=[value_id(node.args[0], ctx), value_id(node.args[1], ctx), value_id(node.args[2], ctx)], - outputs=[value_id(node, ctx)], - meta=_base_meta(shape, dtype, torch_op, node), - ) - register_node(ir, ir_node, shape=shape, dtype=dtype) - - -def import_masked_fill( - ir: IRGraph, - node: Any, - ctx: ImportContext, - *, - shape: tuple[int, ...] | None, - dtype: str | None, - torch_op: str, -) -> None: - value = _extract_numeric_literal(node.args[2]) - if value is None: - ctx.fail(f"unsupported masked_fill value for {torch_op}: {node.args!r}") - ir_node = IRNode( - id=node_id(node), - op="masked_fill", - inputs=[value_id(node.args[0], ctx), value_id(node.args[1], ctx)], - outputs=[value_id(node, ctx)], - attrs={"value": float(value)}, - meta=_base_meta(shape, dtype, torch_op, node), - ) - register_node(ir, ir_node, shape=shape, dtype=dtype) - - -def import_expand(ir: IRGraph, node: Any, ctx: ImportContext, *, shape: tuple[int, ...] | None, dtype: str | None, torch_op: str) -> None: - target_shape = _extract_int_sequence_literal(node.args[1]) - if target_shape is None and shape is not None: - target_shape = _extract_int_sequence_literal(shape) - if target_shape is None: - ctx.fail(f"unsupported expand shape for {torch_op}: {node.args!r}") - ir_node = IRNode( - id=node_id(node), - op="expand", - inputs=[value_id(node.args[0], ctx)], - outputs=[value_id(node, ctx)], - attrs={"shape": target_shape}, - meta=_base_meta(shape, dtype, torch_op, node), - ) - register_node(ir, ir_node, shape=shape, dtype=dtype) - - -def import_repeat(ir: IRGraph, node: Any, ctx: ImportContext, *, shape: tuple[int, ...] | None, dtype: str | None, torch_op: str) -> None: - repeats = _extract_int_sequence_literal(node.args[1]) if len(node.args) > 1 else None - if repeats is None: - ctx.fail(f"unsupported repeat factors for {torch_op}: {node.args!r}") - ir_node = IRNode( - id=node_id(node), - op="repeat", - inputs=[value_id(node.args[0], ctx)], - outputs=[value_id(node, ctx)], - attrs={"repeats": repeats}, - meta=_base_meta(shape, dtype, torch_op, node), - ) - register_node(ir, ir_node, shape=shape, dtype=dtype) - - -def import_one_hot(ir: IRGraph, node: Any, ctx: ImportContext, *, shape: tuple[int, ...] | None, dtype: str | None, torch_op: str) -> None: - num_classes = _extract_static_int(node.args[1]) if len(node.args) > 1 else None - if (num_classes is None or num_classes <= 0) and shape is not None and shape: - num_classes = int(shape[-1]) - if num_classes is None or num_classes <= 0: - ctx.fail(f"unsupported one_hot num_classes for {torch_op}: {node.args!r}") - ir_node = IRNode( - id=node_id(node), - op="one_hot", - inputs=[value_id(node.args[0], ctx)], - outputs=[value_id(node, ctx)], - attrs={"num_classes": int(num_classes)}, - meta=_base_meta(shape, dtype, torch_op, node), - ) - register_node(ir, ir_node, shape=shape, dtype=dtype) - - -def import_tril(ir: IRGraph, node: Any, ctx: ImportContext, *, shape: tuple[int, ...] | None, dtype: str | None, torch_op: str) -> None: - diagonal = int(extract_literals(node.args[1])) if len(node.args) > 1 and extract_literals(node.args[1]) is not None else 0 - ir_node = IRNode( - id=node_id(node), - op="tril", - inputs=[value_id(node.args[0], ctx)], - outputs=[value_id(node, ctx)], - attrs={"diagonal": diagonal}, - meta=_base_meta(shape, dtype, torch_op, node), - ) - register_node(ir, ir_node, shape=shape, dtype=dtype) - - -def import_unfold(ir: IRGraph, node: Any, ctx: ImportContext, *, shape: tuple[int, ...] | None, dtype: str | None, torch_op: str) -> None: - if len(node.args) < 4: - ctx.fail(f"unsupported unfold args for {torch_op}: {node.args!r}") - dimension = extract_literals(node.args[1]) - size = extract_literals(node.args[2]) - step = extract_literals(node.args[3]) - if not isinstance(dimension, int) or not isinstance(size, int) or not isinstance(step, int): - ctx.fail(f"unsupported unfold args for {torch_op}: {node.args!r}") - ir_node = IRNode( - id=node_id(node), - op="unfold", - inputs=[value_id(node.args[0], ctx)], - outputs=[value_id(node, ctx)], - attrs={"dimension": dimension, "size": size, "step": step}, - meta=_base_meta(shape, dtype, torch_op, node), - ) - register_node(ir, ir_node, shape=shape, dtype=dtype) - - -def import_transpose(ir: IRGraph, node: Any, ctx: ImportContext, *, shape: tuple[int, ...] | None, dtype: str | None, torch_op: str) -> None: - attrs = {"dim0": 0, "dim1": 1} if torch_op == "aten.t.default" else {"dim0": int(node.args[1]), "dim1": int(node.args[2])} - ir_node = IRNode( - id=node_id(node), - op="transpose", - inputs=[value_id(node.args[0], ctx)], - outputs=[value_id(node, ctx)], - attrs=attrs, - meta=_base_meta(shape, dtype, torch_op, node), - ) - register_node(ir, ir_node, shape=shape, dtype=dtype) - - -def import_permute(ir: IRGraph, node: Any, ctx: ImportContext, *, shape: tuple[int, ...] | None, dtype: str | None, torch_op: str) -> None: - permutation = _extract_int_sequence_literal(node.args[1]) - if permutation is None: - ctx.fail(f"unsupported permute dims for {torch_op}: {node.args!r}") - ir_node = IRNode( - id=node_id(node), - op="permute", - inputs=[value_id(node.args[0], ctx)], - outputs=[value_id(node, ctx)], - attrs={"permutation": permutation}, - meta=_base_meta(shape, dtype, torch_op, node), - ) - register_node(ir, ir_node, shape=shape, dtype=dtype) - - -def import_numpy_t(ir: IRGraph, node: Any, ctx: ImportContext, *, shape: tuple[int, ...] | None, dtype: str | None, torch_op: str) -> None: - input_shape = _extract_fx_tensor_shape(node.args[0]) - rank_source = input_shape if input_shape is not None else shape - if rank_source is None: - ctx.fail(f"unsupported numpy_T rank inference for {torch_op}: {node.args!r}") - permutation = tuple(range(len(rank_source) - 1, -1, -1)) - ir_node = IRNode( - id=node_id(node), - op="permute", - inputs=[value_id(node.args[0], ctx)], - outputs=[value_id(node, ctx)], - attrs={"permutation": permutation}, - meta=_base_meta(shape, dtype, torch_op, node), - ) - register_node(ir, ir_node, shape=shape, dtype=dtype) - - -def import_movedim(ir: IRGraph, node: Any, ctx: ImportContext, *, shape: tuple[int, ...] | None, dtype: str | None, torch_op: str) -> None: - source = _extract_int_sequence_literal(node.args[1]) - destination = _extract_int_sequence_literal(node.args[2]) - if source is None or destination is None: - ctx.fail(f"unsupported movedim dims for {torch_op}: {node.args!r}") - ir_node = IRNode( - id=node_id(node), - op="movedim", - inputs=[value_id(node.args[0], ctx)], - outputs=[value_id(node, ctx)], - attrs={"source": source, "destination": destination}, - meta=_base_meta(shape, dtype, torch_op, node), - ) - register_node(ir, ir_node, shape=shape, dtype=dtype) - - -def import_matmul(ir: IRGraph, node: Any, ctx: ImportContext, *, shape: tuple[int, ...] | None, dtype: str | None, torch_op: str) -> None: - ir_node = IRNode( - id=node_id(node), - op="matmul", - inputs=[value_id(node.args[0], ctx), value_id(node.args[1], ctx)], - outputs=[value_id(node, ctx)], - meta=_base_meta(shape, dtype, torch_op, node), - ) - register_node(ir, ir_node, shape=shape, dtype=dtype) - - -def import_linear(ir: IRGraph, node: Any, ctx: ImportContext, *, shape: tuple[int, ...] | None, dtype: str | None, torch_op: str) -> None: - inputs = [value_id(node.args[0], ctx), value_id(node.args[1], ctx)] - attrs = {"has_bias": False} - if len(node.args) > 2 and node.args[2] is not None: - inputs.append(value_id(node.args[2], ctx)) - attrs["has_bias"] = True - ir_node = IRNode( - id=node_id(node), - op="linear", - inputs=inputs, - outputs=[value_id(node, ctx)], - attrs=attrs, - meta=_base_meta(shape, dtype, torch_op, node), - ) - register_node(ir, ir_node, shape=shape, dtype=dtype) - - -def import_addmm(ir: IRGraph, node: Any, ctx: ImportContext, *, shape: tuple[int, ...] | None, dtype: str | None, torch_op: str) -> None: - kwargs = getattr(node, "kwargs", {}) - beta = _extract_numeric_literal(node.args[3]) if len(node.args) > 3 else _extract_numeric_literal(kwargs.get("beta")) - alpha = _extract_numeric_literal(node.args[4]) if len(node.args) > 4 else _extract_numeric_literal(kwargs.get("alpha")) - if (beta is not None and float(beta) != 1.0) or (alpha is not None and float(alpha) != 1.0): - ctx.fail(f"unsupported scaled addmm (beta/alpha != 1) for {torch_op}: {node.args!r} {kwargs!r}") - ir_node = IRNode( - id=node_id(node), - op="addmm", - inputs=[value_id(node.args[0], ctx), value_id(node.args[1], ctx), value_id(node.args[2], ctx)], - outputs=[value_id(node, ctx)], - meta=_base_meta(shape, dtype, torch_op, node), - ) - register_node(ir, ir_node, shape=shape, dtype=dtype) - - -def import_unary(op_name: str): - def _importer(ir: IRGraph, node: Any, ctx: ImportContext, *, shape: tuple[int, ...] | None, dtype: str | None, torch_op: str) -> None: - ir_node = IRNode( - id=node_id(node), - op=op_name, - inputs=[value_id(node.args[0], ctx)], - outputs=[value_id(node, ctx)], - meta=_base_meta(shape, dtype, torch_op, node), - ) - register_node(ir, ir_node, shape=shape, dtype=dtype) - return _importer - - -def import_clamp(ir: IRGraph, node: Any, ctx: ImportContext, *, shape: tuple[int, ...] | None, dtype: str | None, torch_op: str) -> None: - min_value = extract_literals(node.args[1]) if len(node.args) > 1 else None - max_value = extract_literals(node.args[2]) if len(node.args) > 2 else None - attrs: dict[str, object] = {} - inputs = [value_id(node.args[0], ctx)] - if isinstance(min_value, (int, float)): - attrs["min"] = float(min_value) - elif len(node.args) > 1 and is_fx_node(node.args[1]): - attrs["has_min_tensor"] = True - inputs.append(value_id(node.args[1], ctx)) - if isinstance(max_value, (int, float)): - attrs["max"] = float(max_value) - elif len(node.args) > 2 and is_fx_node(node.args[2]): - attrs["has_max_tensor"] = True - inputs.append(value_id(node.args[2], ctx)) - ir_node = IRNode( - id=node_id(node), - op="clamp", - inputs=inputs, - outputs=[value_id(node, ctx)], - attrs=attrs, - meta=_base_meta(shape, dtype, torch_op, node), - ) - register_node(ir, ir_node, shape=shape, dtype=dtype) - - -def import_pow_rsqrt(ir: IRGraph, node: Any, ctx: ImportContext, *, shape: tuple[int, ...] | None, dtype: str | None, torch_op: str) -> None: - ir_node = IRNode( - id=node_id(node), - op="pow", - inputs=[value_id(node.args[0], ctx)], - outputs=[value_id(node, ctx)], - attrs={"exponent": -0.5}, - meta=_base_meta(shape, dtype, torch_op, node), - ) - register_node(ir, ir_node, shape=shape, dtype=dtype) - - -def import_reciprocal(ir: IRGraph, node: Any, ctx: ImportContext, *, shape: tuple[int, ...] | None, dtype: str | None, torch_op: str) -> None: - ir_node = IRNode( - id=node_id(node), - op="pow", - inputs=[value_id(node.args[0], ctx)], - outputs=[value_id(node, ctx)], - attrs={"exponent": -1.0}, - meta=_base_meta(shape, dtype, torch_op, node), - ) - register_node(ir, ir_node, shape=shape, dtype=dtype) - - -def import_floor_divide( - ir: IRGraph, - node: Any, - ctx: ImportContext, - *, - shape: tuple[int, ...] | None, - dtype: str | None, - torch_op: str, -) -> None: - if len(node.args) < 2: - ctx.fail(f"unsupported floor_divide signature for {torch_op}: {node.args!r}") - divisor = _extract_numeric_literal(node.args[1]) - if divisor is None: - ctx.fail(f"unsupported floor_divide divisor for {torch_op}: {node.args!r}") - ir_node = IRNode( - id=node_id(node), - op="scalar_floor_divide", - inputs=[value_id(node.args[0], ctx)], - outputs=[value_id(node, ctx)], - attrs={"value": float(divisor)}, - meta=_base_meta(shape, dtype, torch_op, node), - ) - register_node(ir, ir_node, shape=shape, dtype=dtype) - - -def import_softmax(ir: IRGraph, node: Any, ctx: ImportContext, *, shape: tuple[int, ...] | None, dtype: str | None, torch_op: str) -> None: - ir_node = IRNode( - id=node_id(node), - op="softmax", - inputs=[value_id(node.args[0], ctx)], - outputs=[value_id(node, ctx)], - attrs={"axis": int(node.args[1])}, - meta=_base_meta(shape, dtype, torch_op, node), - ) - register_node(ir, ir_node, shape=shape, dtype=dtype) - - -def import_glu(ir: IRGraph, node: Any, ctx: ImportContext, *, shape: tuple[int, ...] | None, dtype: str | None, torch_op: str) -> None: - axis = -1 - if len(node.args) > 1 and extract_literals(node.args[1]) is not None: - axis = int(node.args[1]) - ir_node = IRNode( - id=node_id(node), - op="glu", - inputs=[value_id(node.args[0], ctx)], - outputs=[value_id(node, ctx)], - attrs={"axis": axis}, - meta=_base_meta(shape, dtype, torch_op, node), - ) - register_node(ir, ir_node, shape=shape, dtype=dtype) - - -def import_scaled_dot_product_attention( - ir: IRGraph, - node: Any, - ctx: ImportContext, - *, - shape: tuple[int, ...] | None, - dtype: str | None, - torch_op: str, -) -> None: - attrs: dict[str, object] = {} - inputs = [value_id(node.args[0], ctx), value_id(node.args[1], ctx), value_id(node.args[2], ctx)] - if len(node.args) > 3: - mask_literal = extract_literals(node.args[3]) - if mask_literal is not None: - attrs["mask"] = mask_literal - elif is_fx_node(node.args[3]): - inputs.append(value_id(node.args[3], ctx)) - mask_dtype = _extract_fx_torch_dtype(node.args[3]) - if mask_dtype is not None and mask_dtype != torch.bool: - attrs["additive_mask"] = True - if len(node.args) > 4 and extract_literals(node.args[4]) is not None: - attrs["dropout_p"] = float(node.args[4]) - if len(node.args) > 5 and extract_literals(node.args[5]) is not None: - attrs["is_causal"] = bool(node.args[5]) - if "scale" in getattr(node, "kwargs", {}) and extract_literals(node.kwargs["scale"]) is not None: - attrs["scale"] = float(node.kwargs["scale"]) - if "enable_gqa" in getattr(node, "kwargs", {}) and extract_literals(node.kwargs["enable_gqa"]) is not None: - attrs["enable_gqa"] = bool(node.kwargs["enable_gqa"]) - ir_node = IRNode( - id=node_id(node), - op="scaled_dot_product_attention", - inputs=inputs, - outputs=[value_id(node, ctx)], - attrs=attrs, - meta=_base_meta(shape, dtype, torch_op, node), - ) - register_node(ir, ir_node, shape=shape, dtype=dtype) - - -def import_reduce(op_name: str): - def _importer(ir: IRGraph, node: Any, ctx: ImportContext, *, shape: tuple[int, ...] | None, dtype: str | None, torch_op: str) -> None: - attrs: dict[str, object] = {} - if len(node.args) > 1 and extract_literals(node.args[1]) is not None: - attrs["axis"] = extract_literals(node.args[1]) - keepdim = None - if len(node.args) > 2 and extract_literals(node.args[2]) is not None: - keepdim = bool(extract_literals(node.args[2])) - elif "keepdim" in getattr(node, "kwargs", {}) and extract_literals(node.kwargs["keepdim"]) is not None: - keepdim = bool(extract_literals(node.kwargs["keepdim"])) - if keepdim is not None: - attrs["keepdim"] = keepdim - ir_node = IRNode( - id=node_id(node), - op=op_name, - inputs=[value_id(node.args[0], ctx)], - outputs=[value_id(node, ctx)], - attrs=attrs, - meta=_base_meta(shape, dtype, torch_op, node), - ) - register_node(ir, ir_node, shape=shape, dtype=dtype) - return _importer - - -def import_cat(ir: IRGraph, node: Any, ctx: ImportContext, *, shape: tuple[int, ...] | None, dtype: str | None, torch_op: str) -> None: - tensors_arg = node.args[0] - axis = int(node.args[1]) if len(node.args) > 1 else 0 - ir_node = IRNode( - id=node_id(node), - op="cat", - inputs=extract_input_ids(tensors_arg, ctx), - outputs=[value_id(node, ctx)], - attrs={"axis": axis}, - meta=_base_meta(shape, dtype, torch_op, node), - ) - register_node(ir, ir_node, shape=shape, dtype=dtype) - - -def import_split_with_sizes(ir: IRGraph, node: Any, ctx: ImportContext, *, shape: tuple[int, ...] | None, dtype: str | None, torch_op: str) -> None: - if len(node.args) < 2: - ctx.fail(f"unsupported split_with_sizes signature for {torch_op}: {node.args!r}") - sizes = _extract_int_sequence_literal(node.args[1]) - if sizes is None: - ctx.fail(f"unsupported split sizes for {torch_op}: {node.args!r}") - axis = -1 - if len(node.args) > 2 and extract_literals(node.args[2]) is not None: - axis = int(extract_literals(node.args[2])) - ir_node = IRNode( - id=node_id(node), - op="split_with_sizes", - inputs=[value_id(node.args[0], ctx)], - outputs=[value_id(node, ctx)], - attrs={"sizes": sizes, "axis": axis}, - meta=_base_meta(shape, dtype, torch_op, node), - ) - register_node(ir, ir_node, shape=shape, dtype=dtype) - - -def import_chunk(ir: IRGraph, node: Any, ctx: ImportContext, *, shape: tuple[int, ...] | None, dtype: str | None, torch_op: str) -> None: - if len(node.args) < 2: - ctx.fail(f"unsupported chunk signature for {torch_op}: {node.args!r}") - chunks = extract_literals(node.args[1]) - if not isinstance(chunks, int): - ctx.fail(f"unsupported chunk count for {torch_op}: {node.args!r}") - axis = 0 - if len(node.args) > 2 and extract_literals(node.args[2]) is not None: - axis = int(extract_literals(node.args[2])) - ir_node = IRNode( - id=node_id(node), - op="chunk", - inputs=[value_id(node.args[0], ctx)], - outputs=[value_id(node, ctx)], - attrs={"chunks": int(chunks), "axis": axis}, - meta=_base_meta(shape, dtype, torch_op, node), - ) - register_node(ir, ir_node, shape=shape, dtype=dtype) - - -def import_unbind(ir: IRGraph, node: Any, ctx: ImportContext, *, shape: tuple[int, ...] | None, dtype: str | None, torch_op: str) -> None: - if len(node.args) < 1: - ctx.fail(f"unsupported unbind signature for {torch_op}: {node.args!r}") - axis = 0 - if len(node.args) > 1 and extract_literals(node.args[1]) is not None: - axis = int(extract_literals(node.args[1])) - ir_node = IRNode( - id=node_id(node), - op="unbind", - inputs=[value_id(node.args[0], ctx)], - outputs=[value_id(node, ctx)], - attrs={"axis": axis}, - meta=_base_meta(shape, dtype, torch_op, node), - ) - register_node(ir, ir_node, shape=shape, dtype=dtype) - - -def import_pow(ir: IRGraph, node: Any, ctx: ImportContext, *, shape: tuple[int, ...] | None, dtype: str | None, torch_op: str) -> None: - base_literal = _extract_numeric_literal(node.args[0]) if len(node.args) > 0 else None - attrs: dict[str, object] = {} - exponent_literal = extract_literals(node.args[1]) if len(node.args) > 1 else None - - if base_literal is not None and isinstance(exponent_literal, (int, float)): - add_value_if_missing(ir, IRValue(id=value_id(node, ctx), shape=shape, dtype=dtype, producer=None)) - ir.constants[value_id(node, ctx)] = base_literal**exponent_literal - return - - if base_literal is not None: - if not (len(node.args) > 1 and is_fx_node(node.args[1])): - ctx.fail(f"unsupported pow signature for {torch_op}: scalar base with non-literal exponent {node.args!r}") - if float(base_literal) <= 0.0: - ctx.fail(f"unsupported pow signature for {torch_op}: non-positive scalar base {base_literal!r}") - - exponent_value_id = value_id(node.args[1], ctx) - mul_output_id = f"{value_id(node, ctx)}__pow_scalar_base_log_mul" - mul_node = IRNode( - id=f"{node_id(node)}__pow_scalar_base_log_mul", - op="scalar_multiply", - inputs=[exponent_value_id], - outputs=[mul_output_id], - attrs={"value": float(math.log(float(base_literal)))}, - meta=_base_meta(shape, dtype, torch_op, node), - ) - register_node(ir, mul_node, shape=shape, dtype=dtype) - - exp_node = IRNode( - id=node_id(node), - op="scalar_exp", - inputs=[mul_output_id], - outputs=[value_id(node, ctx)], - meta=_base_meta(shape, dtype, torch_op, node), - ) - register_node(ir, exp_node, shape=shape, dtype=dtype) - return - - inputs = [value_id(node.args[0], ctx)] - if exponent_literal is None and len(node.args) > 1 and is_fx_node(node.args[1]): - inputs.append(value_id(node.args[1], ctx)) - else: - attrs["exponent"] = exponent_literal - ir_node = IRNode( - id=node_id(node), - op="pow", - inputs=inputs, - outputs=[value_id(node, ctx)], - attrs=attrs, - meta=_base_meta(shape, dtype, torch_op, node), - ) - register_node(ir, ir_node, shape=shape, dtype=dtype) - - -def import_slice(ir: IRGraph, node: Any, ctx: ImportContext, *, shape: tuple[int, ...] | None, dtype: str | None, torch_op: str) -> None: - attrs = {"axis": int(node.args[1]), "start": int(node.args[2]), "end": int(node.args[3])} - if len(node.args) > 4 and extract_literals(node.args[4]) is not None: - attrs["step"] = int(node.args[4]) - ir_node = IRNode( - id=node_id(node), - op="slice", - inputs=[value_id(node.args[0], ctx)], - outputs=[value_id(node, ctx)], - attrs=attrs, - meta=_base_meta(shape, dtype, torch_op, node), - ) - register_node(ir, ir_node, shape=shape, dtype=dtype) - - -def import_diff(ir: IRGraph, node: Any, ctx: ImportContext, *, shape: tuple[int, ...] | None, dtype: str | None, torch_op: str) -> None: - inputs = list(getattr(node, "args", ())) - if not inputs: - ctx.fail(f"unsupported diff signature for {torch_op}: missing input tensor") - - source = inputs[0] - if not is_fx_node(source): - ctx.fail(f"unsupported diff signature for {torch_op}: source is not FX node") - - n_value = 1 - if len(inputs) > 1 and extract_literals(inputs[1]) is not None: - n_value = int(extract_literals(inputs[1])) - elif "n" in getattr(node, "kwargs", {}) and extract_literals(node.kwargs["n"]) is not None: - n_value = int(extract_literals(node.kwargs["n"])) - if n_value != 1: - ctx.fail(f"unsupported diff order for {torch_op}: n={n_value}") - - dim_value = None - if len(inputs) > 2 and extract_literals(inputs[2]) is not None: - dim_value = int(extract_literals(inputs[2])) - elif "dim" in getattr(node, "kwargs", {}) and extract_literals(node.kwargs["dim"]) is not None: - dim_value = int(extract_literals(node.kwargs["dim"])) - if dim_value is None: - ctx.fail(f"unsupported diff signature for {torch_op}: missing integer dim") - - prepend_value = None - append_value = None - prepend_node = None - append_node = None - if len(inputs) > 3: - if is_fx_node(inputs[3]): - prepend_node = inputs[3] - else: - prepend_value = extract_literals(inputs[3]) - if len(inputs) > 4: - if is_fx_node(inputs[4]): - append_node = inputs[4] - else: - append_value = extract_literals(inputs[4]) - if "prepend" in getattr(node, "kwargs", {}): - prepend_arg = node.kwargs["prepend"] - if is_fx_node(prepend_arg): - prepend_node = prepend_arg - prepend_value = None - else: - prepend_value = extract_literals(prepend_arg) - if "append" in getattr(node, "kwargs", {}): - append_arg = node.kwargs["append"] - if is_fx_node(append_arg): - append_node = append_arg - append_value = None - else: - append_value = extract_literals(append_arg) - if prepend_value is not None or append_value is not None: - ctx.fail(f"unsupported diff signature for {torch_op}: literal prepend/append are not supported") - - source_value_id = value_id(source, ctx) - source_shape = None - if source_value_id in ir.values: - source_shape = ir.values[source_value_id].shape - if source_shape is None: - source_shape = shape - if source_shape is None: - ctx.fail(f"unsupported diff import for {torch_op}: missing source shape") - - rank = len(source_shape) - normalized_dim = dim_value if dim_value >= 0 else dim_value + rank - if normalized_dim < 0 or normalized_dim >= rank: - ctx.fail(f"unsupported diff dimension for {torch_op}: dim={dim_value} rank={rank}") - - dim_extent = int(source_shape[normalized_dim]) - if dim_extent < 1: - ctx.fail(f"unsupported diff input extent for {torch_op}: dim size {dim_extent}") - - concat_inputs = [source_value_id] - concat_shapes = [tuple(int(v) for v in source_shape)] - - def _append_concat_input(arg: Any, label: str) -> None: - if arg is None: - return - arg_value_id = value_id(arg, ctx) - arg_value = ir.values.get(arg_value_id) - arg_shape = arg_value.shape if arg_value is not None else None - if arg_shape is None: - ctx.fail(f"unsupported diff import for {torch_op}: missing {label} shape") - arg_shape = tuple(int(v) for v in arg_shape) - if len(arg_shape) != rank: - ctx.fail( - f"unsupported diff import for {torch_op}: {label} rank {len(arg_shape)} does not match source rank {rank}" - ) - for axis, (arg_dim, src_dim) in enumerate(zip(arg_shape, source_shape, strict=True)): - if axis == normalized_dim: - continue - if int(arg_dim) != int(src_dim): - ctx.fail( - f"unsupported diff import for {torch_op}: {label} shape {arg_shape} " - f"is incompatible with source shape {source_shape}" - ) - concat_inputs.append(arg_value_id) - concat_shapes.append(arg_shape) - - if prepend_node is not None: - concat_inputs = [] - concat_shapes = [] - _append_concat_input(prepend_node, "prepend") - concat_inputs.append(source_value_id) - concat_shapes.append(tuple(int(v) for v in source_shape)) - if append_node is not None: - _append_concat_input(append_node, "append") - elif append_node is not None: - _append_concat_input(append_node, "append") - - augmented_shape = list(source_shape) - augmented_shape[normalized_dim] = sum(int(shape_item[normalized_dim]) for shape_item in concat_shapes) - augmented_shape_tuple = tuple(int(v) for v in augmented_shape) - diff_shape = list(augmented_shape) - diff_shape[normalized_dim] = max(0, int(augmented_shape[normalized_dim]) - 1) - diff_shape_tuple = tuple(int(v) for v in diff_shape) - - diff_source_value_id = source_value_id - if len(concat_inputs) > 1: - concat_output_id = f"{value_id(node, ctx)}__diff_source" - concat_node = IRNode( - id=f"{node_id(node)}__diff_source", - op="cat", - inputs=concat_inputs, - outputs=[concat_output_id], - attrs={"axis": normalized_dim}, - meta=_base_meta(augmented_shape_tuple, dtype, torch_op, node), - ) - register_node(ir, concat_node, shape=augmented_shape_tuple, dtype=dtype) - diff_source_value_id = concat_output_id - - left_shape = list(diff_shape_tuple) - left_shape_tuple = tuple(int(v) for v in left_shape) - - left_node = IRNode( - id=f"{node_id(node)}__diff_left", - op="slice", - inputs=[diff_source_value_id], - outputs=[f"{value_id(node, ctx)}__diff_left"], - attrs={"axis": normalized_dim, "start": 0, "end": int(augmented_shape[normalized_dim]) - 1, "step": 1}, - meta=_base_meta(left_shape_tuple, dtype, torch_op, node), - ) - register_node(ir, left_node, shape=left_shape_tuple, dtype=dtype) - - right_node = IRNode( - id=f"{node_id(node)}__diff_right", - op="slice", - inputs=[diff_source_value_id], - outputs=[f"{value_id(node, ctx)}__diff_right"], - attrs={"axis": normalized_dim, "start": 1, "end": int(augmented_shape[normalized_dim]), "step": 1}, - meta=_base_meta(left_shape_tuple, dtype, torch_op, node), - ) - register_node(ir, right_node, shape=left_shape_tuple, dtype=dtype) - - diff_node = IRNode( - id=node_id(node), - op="subtract", - inputs=[right_node.outputs[0], left_node.outputs[0]], - outputs=[value_id(node, ctx)], - meta=_base_meta(diff_shape_tuple, dtype, torch_op, node), - ) - register_node(ir, diff_node, shape=diff_shape_tuple, dtype=dtype) - - -def import_index(ir: IRGraph, node: Any, ctx: ImportContext, *, shape: tuple[int, ...] | None, dtype: str | None, torch_op: str) -> None: - if torch_op.startswith("aten.select"): - axis = int(node.args[1]) - index_value = int(node.args[2]) - else: - index_value = int(node.args[1]) - axis = int(node.args[2]) if len(node.args) > 2 else 0 - ir_node = IRNode( - id=node_id(node), - op="index", - inputs=[value_id(node.args[0], ctx)], - outputs=[value_id(node, ctx)], - attrs={"index_value": index_value, "axis": axis}, - meta=_base_meta(shape, dtype, torch_op, node), - ) - register_node(ir, ir_node, shape=shape, dtype=dtype) - - -def import_gather(ir: IRGraph, node: Any, ctx: ImportContext, *, shape: tuple[int, ...] | None, dtype: str | None, torch_op: str) -> None: - ir_node = IRNode( - id=node_id(node), - op="gather", - inputs=[value_id(node.args[0], ctx), value_id(node.args[2], ctx)], - outputs=[value_id(node, ctx)], - attrs={"axis": int(node.args[1])}, - meta=_base_meta(shape, dtype, torch_op, node), - ) - register_node(ir, ir_node, shape=shape, dtype=dtype) - - -def import_embedding(ir: IRGraph, node: Any, ctx: ImportContext, *, shape: tuple[int, ...] | None, dtype: str | None, torch_op: str) -> None: - inputs = [value_id(node.args[0], ctx), value_id(node.args[1], ctx)] - attrs: dict[str, object] = {} - if len(node.args) > 2 and extract_literals(node.args[2]) is not None: - attrs["padding_idx"] = int(node.args[2]) - if len(node.args) > 3 and extract_literals(node.args[3]) is not None: - attrs["scale_grad_by_freq"] = bool(node.args[3]) - if len(node.args) > 4 and extract_literals(node.args[4]) is not None: - attrs["sparse"] = bool(node.args[4]) - ir_node = IRNode( - id=node_id(node), - op="embedding", - inputs=inputs, - outputs=[value_id(node, ctx)], - attrs=attrs, - meta=_base_meta(shape, dtype, torch_op, node), - ) - register_node(ir, ir_node, shape=shape, dtype=dtype) - - -def _extract_int_list_or_scalar(value: Any, *, default: int) -> int: - literal = extract_literals(value) - if literal is None: - return default - if isinstance(literal, int): - return int(literal) - if isinstance(literal, (tuple, list)) and literal: - first = literal[0] - if isinstance(first, int): - return int(first) - raise UnsupportedImportError(f"unsupported integer/list literal: {value!r}") - - -def import_conv1d(ir: IRGraph, node: Any, ctx: ImportContext, *, shape: tuple[int, ...] | None, dtype: str | None, torch_op: str) -> None: - if len(node.args) < 2: - ctx.fail(f"unsupported conv1d signature for {torch_op}: {node.args!r}") - - inputs = [value_id(node.args[0], ctx), value_id(node.args[1], ctx)] - if len(node.args) > 2 and is_fx_node(node.args[2]): - inputs.append(value_id(node.args[2], ctx)) - - stride = _extract_int_list_or_scalar(node.args[3], default=1) if len(node.args) > 3 else 1 - padding = _extract_int_list_or_scalar(node.args[4], default=0) if len(node.args) > 4 else 0 - dilation = _extract_int_list_or_scalar(node.args[5], default=1) if len(node.args) > 5 else 1 - groups = _extract_int_list_or_scalar(node.args[6], default=1) if len(node.args) > 6 else 1 - - ir_node = IRNode( - id=node_id(node), - op="conv1d", - inputs=inputs, - outputs=[value_id(node, ctx)], - attrs={"stride": stride, "padding": padding, "dilation": dilation, "groups": groups}, - meta=_base_meta(shape, dtype, torch_op, node), - ) - register_node(ir, ir_node, shape=shape, dtype=dtype) - - -def import_conv2d(ir: IRGraph, node: Any, ctx: ImportContext, *, shape: tuple[int, ...] | None, dtype: str | None, torch_op: str) -> None: - if len(node.args) < 2: - ctx.fail(f"unsupported conv2d signature for {torch_op}: {node.args!r}") - - inputs = [value_id(node.args[0], ctx), value_id(node.args[1], ctx)] - if len(node.args) > 2 and is_fx_node(node.args[2]): - inputs.append(value_id(node.args[2], ctx)) - - stride = _extract_int_list_or_scalar(node.args[3], default=1) if len(node.args) > 3 else 1 - padding = _extract_int_list_or_scalar(node.args[4], default=0) if len(node.args) > 4 else 0 - dilation = _extract_int_list_or_scalar(node.args[5], default=1) if len(node.args) > 5 else 1 - groups = _extract_int_list_or_scalar(node.args[6], default=1) if len(node.args) > 6 else 1 - - ir_node = IRNode( - id=node_id(node), - op="conv2d", - inputs=inputs, - outputs=[value_id(node, ctx)], - attrs={"stride": stride, "padding": padding, "dilation": dilation, "groups": groups}, - meta=_base_meta(shape, dtype, torch_op, node), - ) - register_node(ir, ir_node, shape=shape, dtype=dtype) - - -def import_pad(ir: IRGraph, node: Any, ctx: ImportContext, *, shape: tuple[int, ...] | None, dtype: str | None, torch_op: str) -> None: - if len(node.args) < 2: - ctx.fail(f"unsupported pad signature for {torch_op}: {node.args!r}") - - pads = extract_literals(node.args[1]) - if not isinstance(pads, (tuple, list)) or not all(isinstance(v, int) for v in pads): - ctx.fail(f"unsupported pads for {torch_op}: {node.args!r}") - - mode = "constant" - if len(node.args) > 2 and extract_literals(node.args[2]) is not None: - mode = str(extract_literals(node.args[2])) - if "mode" in getattr(node, "kwargs", {}) and extract_literals(node.kwargs["mode"]) is not None: - mode = str(extract_literals(node.kwargs["mode"])) - - value = 0.0 - if len(node.args) > 3 and extract_literals(node.args[3]) is not None: - value = float(extract_literals(node.args[3])) - if "value" in getattr(node, "kwargs", {}) and extract_literals(node.kwargs["value"]) is not None: - value = float(extract_literals(node.kwargs["value"])) - - ir_node = IRNode( - id=node_id(node), - op="pad", - inputs=[value_id(node.args[0], ctx)], - outputs=[value_id(node, ctx)], - attrs={"pads": tuple(int(v) for v in pads), "mode": mode, "value": value}, - meta=_base_meta(shape, dtype, torch_op, node), - ) - register_node(ir, ir_node, shape=shape, dtype=dtype) - - -def import_ones(ir: IRGraph, node: Any, ctx: ImportContext, *, shape: tuple[int, ...] | None, dtype: str | None, torch_op: str) -> None: - if not node.args: - ctx.fail(f"unsupported ones signature for {torch_op}: {node.args!r}") - - output_shape = _extract_int_sequence_literal(node.args[0]) - if output_shape is None: - ctx.fail(f"unsupported ones shape for {torch_op}: {node.args!r}") - - node_dtype = dtype - if "dtype" in getattr(node, "kwargs", {}) and extract_literals(node.kwargs["dtype"]) is not None: - node_dtype = dtype_to_ir(extract_literals(node.kwargs["dtype"])) - elif len(node.args) > 1 and extract_literals(node.args[1]) is not None: - node_dtype = dtype_to_ir(extract_literals(node.args[1])) - if node_dtype is None: - node_dtype = "fp32" - - ir_node = IRNode( - id=node_id(node), - op="ones", - inputs=[], - outputs=[value_id(node, ctx)], - attrs={"shape": output_shape, "dtype": node_dtype}, - meta=_base_meta(shape or output_shape, node_dtype, torch_op, node), - ) - register_node(ir, ir_node, shape=shape or output_shape, dtype=node_dtype) - - -def import_layer_norm(ir: IRGraph, node: Any, ctx: ImportContext, *, shape: tuple[int, ...] | None, dtype: str | None, torch_op: str) -> None: - if not node.args: - ctx.fail(f"unsupported layer_norm signature for {torch_op}: {node.args!r}") - - normalized_shape_value = None - if len(node.args) > 1: - normalized_shape_value = _extract_int_sequence_literal(node.args[1]) - if normalized_shape_value is None: - normalized_shape_value = _extract_int_sequence_literal(getattr(node, "kwargs", {}).get("normalized_shape")) - if normalized_shape_value is None: - ctx.fail(f"unsupported normalized_shape for {torch_op}: {node.args!r} {getattr(node, 'kwargs', {})!r}") - - eps_value = None - if len(node.args) > 4: - eps_value = _extract_numeric_literal(node.args[4]) - if eps_value is None: - eps_value = _extract_numeric_literal(getattr(node, "kwargs", {}).get("eps")) - if eps_value is None: - eps_value = 1e-5 - - inputs = [value_id(node.args[0], ctx)] - weight = node.args[2] if len(node.args) > 2 else getattr(node, "kwargs", {}).get("weight") - bias = node.args[3] if len(node.args) > 3 else getattr(node, "kwargs", {}).get("bias") - if not is_fx_node(weight): - ctx.fail(f"unsupported layer_norm weight form for {torch_op}: {node.args!r} {getattr(node, 'kwargs', {})!r}") - inputs.append(value_id(weight, ctx)) - if bias is None: - pass - elif is_fx_node(bias): - inputs.append(value_id(bias, ctx)) - else: - ctx.fail(f"unsupported layer_norm bias form for {torch_op}: {node.args!r} {getattr(node, 'kwargs', {})!r}") - - attrs = {"normalized_shape": tuple(int(v) for v in normalized_shape_value), "eps": float(eps_value)} - ir_node = IRNode( - id=node_id(node), - op="layer_norm", - inputs=inputs, - outputs=[value_id(node, ctx)], - attrs=attrs, - meta=_base_meta(shape, dtype, torch_op, node), - ) - register_node(ir, ir_node, shape=shape, dtype=dtype) - - -def import_rms_norm(ir: IRGraph, node: Any, ctx: ImportContext, *, shape: tuple[int, ...] | None, dtype: str | None, torch_op: str) -> None: - inputs = [value_id(node.args[0], ctx), value_id(node.args[2], ctx)] - attrs = {"normalized_shape": tuple(int(v) for v in node.args[1]), "eps": float(node.args[3])} - ir_node = IRNode( - id=node_id(node), - op="rms_norm", - inputs=inputs, - outputs=[value_id(node, ctx)], - attrs=attrs, - meta=_base_meta(shape, dtype, torch_op, node), - ) - register_node(ir, ir_node, shape=shape, dtype=dtype) - - -def import_group_norm(ir: IRGraph, node: Any, ctx: ImportContext, *, shape: tuple[int, ...] | None, dtype: str | None, torch_op: str) -> None: - inputs = [value_id(node.args[0], ctx), value_id(node.args[2], ctx), value_id(node.args[3], ctx)] - attrs = {"num_groups": int(node.args[1]), "eps": float(node.args[4])} - ir_node = IRNode( - id=node_id(node), - op="group_norm", - inputs=inputs, - outputs=[value_id(node, ctx)], - attrs=attrs, - meta=_base_meta(shape, dtype, torch_op, node), - ) - register_node(ir, ir_node, shape=shape, dtype=dtype) - - -def import_batch_norm(ir: IRGraph, node: Any, ctx: ImportContext, *, shape: tuple[int, ...] | None, dtype: str | None, torch_op: str) -> None: - inputs = [ - value_id(node.args[0], ctx), - value_id(node.args[1], ctx), - value_id(node.args[2], ctx), - value_id(node.args[3], ctx), - value_id(node.args[4], ctx), - ] - attrs = {"training": bool(node.args[5]), "momentum": float(node.args[6]), "eps": float(node.args[7])} - ir_node = IRNode( - id=node_id(node), - op="batch_norm", - inputs=inputs, - outputs=[value_id(node, ctx)], - attrs=attrs, - meta=_base_meta(shape, dtype, torch_op, node), - ) - register_node(ir, ir_node, shape=shape, dtype=dtype) - - -def import_contiguous(ir: IRGraph, node: Any, ctx: ImportContext, *, shape: tuple[int, ...] | None, dtype: str | None, torch_op: str) -> None: - ir_node = IRNode( - id=node_id(node), - op="contiguous", - inputs=[value_id(node.args[0], ctx)], - outputs=[value_id(node, ctx)], - meta=_base_meta(shape, dtype, torch_op, node), - ) - register_node(ir, ir_node, shape=shape, dtype=dtype) - - -def import_getitem(ir: IRGraph, node: Any, ctx: ImportContext, *, shape: tuple[int, ...] | None, dtype: str | None, torch_op: str) -> None: - index_value = extract_literals(node.args[1]) - if not isinstance(index_value, int): - ctx.fail(f"unsupported getitem index for {torch_op}: {node.args!r}") - ir_node = IRNode( - id=node_id(node), - op="getitem", - inputs=[value_id(node.args[0], ctx)], - outputs=[value_id(node, ctx)], - attrs={"index": index_value}, - meta=_base_meta(shape, dtype, torch_op, node), - ) - register_node(ir, ir_node, shape=shape, dtype=dtype) - - -OP_IMPORTERS = { - "arange": import_arange, - "identity": import_identity, - "add": import_add, - "subtract": import_subtract, - "multiply": import_multiply, - "multiply_inplace": import_multiply_inplace, - "divide": import_divide, - "precision_cast": import_precision_cast, - "type_as": import_type_as, - "negate": import_negate, - "abs": import_unary("abs"), - "clamp": import_clamp, - "not_equal": import_not_equal, - "equal": import_equal, - "greater": lambda ir, node, ctx, *, shape, dtype, torch_op: import_compare( - ir, node, ctx, shape=shape, dtype=dtype, torch_op=torch_op, op_name="greater" - ), - "greater_equal": lambda ir, node, ctx, *, shape, dtype, torch_op: import_compare( - ir, node, ctx, shape=shape, dtype=dtype, torch_op=torch_op, op_name="greater_equal" - ), - "less": lambda ir, node, ctx, *, shape, dtype, torch_op: import_compare( - ir, node, ctx, shape=shape, dtype=dtype, torch_op=torch_op, op_name="less" - ), - "less_equal": lambda ir, node, ctx, *, shape, dtype, torch_op: import_compare( - ir, node, ctx, shape=shape, dtype=dtype, torch_op=torch_op, op_name="less_equal" - ), - "logical_and": import_binary("logical_and"), - "logical_or": import_binary("logical_or"), - "logical_not": import_unary("logical_not"), - "where": import_where, - "masked_scatter": import_masked_scatter, - "masked_fill": import_masked_fill, - "cos": import_unary("cos"), - "sin": import_unary("sin"), - "floor_divide": import_floor_divide, - "scalar_exp": import_unary("scalar_exp"), - "scalar_sqrt": import_unary("scalar_sqrt"), - "reciprocal": import_reciprocal, - "scalar_log": import_unary("scalar_log"), - "rsqrt": import_pow_rsqrt, - "reshape": import_reshape, - "flatten": import_flatten, - "unsqueeze": import_unsqueeze, - "squeeze": import_squeeze, - "expand": import_expand, - "repeat": import_repeat, - "one_hot": import_one_hot, - "tril": import_tril, - "unfold": import_unfold, - "numpy_T": import_numpy_t, - "transpose": import_transpose, - "permute": import_permute, - "movedim": import_movedim, - "matmul": import_matmul, - "linear": import_linear, - "addmm": import_addmm, - "relu": import_unary("relu"), - "silu": import_unary("silu"), - "gelu": import_unary("gelu"), - "gelu_erf": import_unary("gelu_erf"), - "sigmoid": import_unary("sigmoid"), - "glu": import_glu, - "softplus": import_unary("softplus"), - "tanh": import_unary("tanh"), - "softmax": import_softmax, - "scaled_dot_product_attention": import_scaled_dot_product_attention, - "sum": import_reduce("sum"), - "mean": import_reduce("mean"), - "variance": import_reduce("variance"), - "min": import_reduce("min"), - "max": import_reduce("max"), - "cat": import_cat, - "split_with_sizes": import_split_with_sizes, - "chunk": import_chunk, - "unbind": import_unbind, - "aten.unbind.int": import_unbind, - "ones": import_ones, - "pad": import_pad, - "pow": import_pow, - "aten.diff.default": import_diff, - "slice": import_slice, - "index": import_index, - "gather": import_gather, - "embedding": import_embedding, - "conv1d": import_conv1d, - "conv2d": import_conv2d, - "layer_norm": import_layer_norm, - "rms_norm": import_rms_norm, - "group_norm": import_group_norm, - "batch_norm": import_batch_norm, - "contiguous": import_contiguous, - "getitem": import_getitem, -} diff --git a/python/cactus/transpile/jax_semantic_rewrites.py b/python/cactus/transpile/jax_semantic_rewrites.py deleted file mode 100644 index d0dd63f15..000000000 --- a/python/cactus/transpile/jax_semantic_rewrites.py +++ /dev/null @@ -1,875 +0,0 @@ -from __future__ import annotations - -from collections.abc import Callable -from dataclasses import dataclass - -import numpy as np - -from cactus.transpile.graph_ir import IRGraph -from cactus.transpile.graph_ir import IRNode -from cactus.transpile.graph_ir import IRValue -from cactus.transpile.import_semantics import apply_import_semantics - - -@dataclass(frozen=True) -class JaxPattern: - name: str - apply: Callable[[IRGraph], None] - - -def _constant_scalar_or_singleton(graph: IRGraph, value_id: str) -> float | None: - value = graph.constants.get(value_id) - if value is None: - return None - array = np.asarray(value) - if graph.values[value_id].meta.get("jax_closed_constant") and array.size != 1: - return None - if array.size == 0: - return None - first = float(array.reshape(-1)[0].item()) - if np.allclose(array, first, rtol=0.0, atol=0.0): - return first - return None - - -def _rebuild_users(graph: IRGraph) -> None: - for value in graph.values.values(): - value.users.clear() - for node_id in graph.order: - node = graph.nodes[node_id] - for input_id in node.inputs: - graph.values[input_id].users.append(node_id) - - -def _producer(graph: IRGraph, value_id: str) -> IRNode | None: - value = graph.values.get(value_id) - if value is None or value.producer is None: - return None - return graph.nodes.get(value.producer) - - -def _strip_simple_wrappers(graph: IRGraph, value_id: str) -> str: - current = value_id - for _ in range(8): - producer = _producer(graph, current) - if producer is None or producer.op not in {"precision_cast", "reshape", "view", "expand"} or not producer.inputs: - return current - current = producer.inputs[0] - return current - - -def _trace_reduction_source(graph: IRGraph, value_id: str) -> str | None: - current = _strip_simple_wrappers(graph, value_id) - node = _producer(graph, current) - if node is not None and node.op == "divide" and len(node.inputs) == 2: - divisor = _constant_scalar_or_singleton(graph, node.inputs[1]) - if divisor is None: - return None - current = _strip_simple_wrappers(graph, node.inputs[0]) - node = _producer(graph, current) - if node is None or node.op not in {"mean", "sum"} or not node.inputs: - return None - return _trace_precision_cast_source(graph, _strip_simple_wrappers(graph, node.inputs[0])) - - -def _trace_square_reduction_source(graph: IRGraph, value_id: str, term_trace) -> str | None: - current = _strip_simple_wrappers(graph, value_id) - node = _producer(graph, current) - if node is not None and node.op == "divide" and len(node.inputs) == 2: - divisor = _constant_scalar_or_singleton(graph, node.inputs[1]) - if divisor is None: - return None - current = _strip_simple_wrappers(graph, node.inputs[0]) - node = _producer(graph, current) - if node is None or node.op not in {"mean", "sum"} or not node.inputs: - return None - square_node = _producer(graph, _strip_simple_wrappers(graph, node.inputs[0])) - if square_node is None or square_node.op != "multiply" or len(square_node.inputs) != 2: - return None - lhs = term_trace(graph, square_node.inputs[0]) - rhs = term_trace(graph, square_node.inputs[1]) - if lhs is None or rhs is None: - return None - if _strip_simple_wrappers(graph, lhs) != _strip_simple_wrappers(graph, rhs): - return None - return _strip_simple_wrappers(graph, lhs) - - -def _trace_raw_square_source(graph: IRGraph, value_id: str) -> str | None: - return _strip_simple_wrappers(graph, _trace_precision_cast_source(graph, _strip_simple_wrappers(graph, value_id))) - - -def _trace_mean_square_source(graph: IRGraph, value_id: str) -> str | None: - return _trace_square_reduction_source(graph, value_id, _trace_raw_square_source) - - -def _trace_layer_norm_variance_source(graph: IRGraph, value_id: str) -> str | None: - current = _strip_simple_wrappers(graph, value_id) - node = _producer(graph, current) - if node is not None and node.op == "where" and node.inputs: - current = _strip_simple_wrappers(graph, node.inputs[-1]) - node = _producer(graph, current) - if node is None or node.op != "subtract" or len(node.inputs) != 2: - return None - mean_square_source = _trace_mean_square_source(graph, node.inputs[0]) - square_mean_node = _producer(graph, _strip_simple_wrappers(graph, node.inputs[1])) - if mean_square_source is None or square_mean_node is None or square_mean_node.op != "multiply": - return None - if len(square_mean_node.inputs) != 2: - return None - lhs_mean_source = _trace_reduction_source(graph, square_mean_node.inputs[0]) - rhs_mean_source = _trace_reduction_source(graph, square_mean_node.inputs[1]) - if lhs_mean_source is None or rhs_mean_source is None: - return None - if _strip_simple_wrappers(graph, lhs_mean_source) != _strip_simple_wrappers(graph, rhs_mean_source): - return None - if _strip_simple_wrappers(graph, mean_square_source) != _strip_simple_wrappers(graph, lhs_mean_source): - return None - return mean_square_source - - -def _trace_layer_norm_centered_source(graph: IRGraph, value_id: str) -> str | None: - centered = _trace_layer_norm_centered(graph, value_id) - return None if centered is None else centered[0] - - -def _trace_layer_norm_centered(graph: IRGraph, value_id: str) -> tuple[str, str] | None: - centered_id = _strip_simple_wrappers(graph, value_id) - node = _producer(graph, centered_id) - if node is None or node.op != "subtract" or len(node.inputs) != 2: - return None - source_id = _trace_precision_cast_source(graph, _strip_simple_wrappers(graph, node.inputs[0])) - mean_source_id = _trace_reduction_source(graph, node.inputs[1]) - if mean_source_id is None: - return None - if _strip_simple_wrappers(graph, source_id) != _strip_simple_wrappers(graph, mean_source_id): - return None - return _strip_simple_wrappers(graph, source_id), centered_id - - -def _trace_inv_std_addend(graph: IRGraph, value_id: str) -> str | None: - current = _strip_simple_wrappers(graph, value_id) - node = _producer(graph, current) - if node is None: - return None - if node.op == "pow" and node.inputs: - exponent = float(node.attrs.get("exponent", 0.0)) - if exponent == -0.5: - return node.inputs[0] - if exponent == -1.0: - return _trace_inv_std_addend(graph, node.inputs[0]) - if node.op == "scalar_sqrt" and node.inputs: - return node.inputs[0] - if node.op == "divide" and len(node.inputs) == 2: - numerator = _constant_scalar_or_singleton(graph, node.inputs[0]) - if numerator == 1.0: - return _trace_inv_std_addend(graph, node.inputs[1]) or node.inputs[1] - return _trace_inv_std_addend(graph, node.inputs[1]) - return None - - -def _trace_norm_inv_std( - graph: IRGraph, - value_id: str, - default_eps: float, - variance_trace, -) -> tuple[str, float] | None: - add_id = _trace_inv_std_addend(graph, value_id) - add_node = _producer(graph, _strip_simple_wrappers(graph, add_id)) if add_id is not None else None - if add_node is None or add_node.op != "add" or len(add_node.inputs) != 2: - return None - - eps = default_eps - source_id: str | None = None - for add_input in add_node.inputs: - scalar = _constant_scalar_or_singleton(graph, add_input) - if scalar is not None: - eps = float(scalar) - continue - variance_source = variance_trace(graph, add_input) - if variance_source is not None: - source_id = variance_source - if source_id is None: - return None - return source_id, eps - - -def _trace_rms_inv_std(graph: IRGraph, value_id: str) -> tuple[str, float] | None: - return _trace_norm_inv_std(graph, value_id, 1.0e-6, _trace_mean_square_source) - - -def _trace_layer_norm_inv_std(graph: IRGraph, value_id: str) -> tuple[str, float] | None: - def _variance_source(g: IRGraph, candidate_id: str) -> str | None: - return _trace_layer_norm_variance_source(g, candidate_id) or _trace_square_reduction_source( - g, - candidate_id, - _trace_layer_norm_centered_source, - ) - - return _trace_norm_inv_std(graph, value_id, 1.0e-5, _variance_source) - - -def _valid_last_dim_weight(graph: IRGraph, weight_id: str, source_id: str) -> str | None: - weight_id = _strip_simple_wrappers(graph, weight_id) - weight_value = graph.values.get(weight_id) - source_value = graph.values.get(source_id) - if weight_value is None or source_value is None: - return None - if weight_value.shape is None or source_value.shape is None: - return None - if len(weight_value.shape) != 1 or int(weight_value.shape[0]) != int(source_value.shape[-1]): - return None - return weight_id - - -def _rewrite_norm_node(node: IRNode, op: str, inputs: list[str], eps: float, rewritten_from: str) -> None: - node.op = op - node.inputs = inputs - node.attrs = {"eps": eps} - node.kind = "semantic" - node.meta = {**node.meta, "rewritten_from": rewritten_from} - - -def _trace_norm_branch(graph: IRGraph, value_id: str, source_trace, inv_trace) -> tuple[str, float] | None: - node = _producer(graph, _strip_simple_wrappers(graph, value_id)) - if node is None or len(node.inputs) != 2: - return None - pairs = ((node.inputs[0], node.inputs[1]),) if node.op == "divide" else () - if node.op == "multiply": - pairs = ((node.inputs[0], node.inputs[1]), (node.inputs[1], node.inputs[0])) - for source_candidate, inv_candidate in pairs: - source_id = source_trace(graph, source_candidate) - inv = inv_trace(graph, inv_candidate) - if source_id is None or inv is None: - continue - inv_source_id, eps = inv - if _strip_simple_wrappers(graph, source_id) == _strip_simple_wrappers(graph, inv_source_id): - return source_id, eps - return None - - -def _trace_weighted_norm(graph: IRGraph, value_id: str, branch_trace) -> tuple[str, str, float] | None: - node = _producer(graph, _strip_simple_wrappers(graph, value_id)) - if node is None or node.op != "multiply" or len(node.inputs) != 2: - return None - for norm_id, weight_id in ((node.inputs[0], node.inputs[1]), (node.inputs[1], node.inputs[0])): - traced = branch_trace(graph, norm_id) - if traced is None: - continue - source_id, eps = traced - weight_id = _valid_last_dim_weight(graph, weight_id, source_id) - if weight_id is None: - continue - return source_id, weight_id, eps - return None - - -def _trace_factored_weighted_norm( - graph: IRGraph, - value_id: str, - source_trace, - inv_trace, -) -> tuple[str, str, float] | None: - node = _producer(graph, _strip_simple_wrappers(graph, value_id)) - if node is None or node.op != "multiply" or len(node.inputs) != 2: - return None - for source_candidate, scaled_id in ((node.inputs[0], node.inputs[1]), (node.inputs[1], node.inputs[0])): - source_id = source_trace(graph, source_candidate) - scaled_node = _producer(graph, _strip_simple_wrappers(graph, scaled_id)) - if source_id is None or scaled_node is None or scaled_node.op != "multiply" or len(scaled_node.inputs) != 2: - continue - for inv_id, weight_id in ((scaled_node.inputs[0], scaled_node.inputs[1]), (scaled_node.inputs[1], scaled_node.inputs[0])): - inv = inv_trace(graph, inv_id) - if inv is None: - continue - inv_source_id, eps = inv - if _strip_simple_wrappers(graph, source_id) != _strip_simple_wrappers(graph, inv_source_id): - continue - weight_id = _valid_last_dim_weight(graph, weight_id, source_id) - if weight_id is not None: - return source_id, weight_id, eps - return None - - -def _trace_weighted_rms_norm(graph: IRGraph, value_id: str) -> tuple[str, str, float] | None: - traced = _trace_weighted_norm( - graph, - value_id, - lambda g, candidate: _trace_norm_branch(g, candidate, _trace_raw_square_source, _trace_rms_inv_std), - ) - if traced is not None: - source_id, _, _ = traced - source_producer = _producer(graph, source_id) - if source_producer is None or source_producer.op != "subtract": - return traced - return _trace_factored_weighted_norm(graph, value_id, _trace_raw_square_source, _trace_rms_inv_std) - - -def _trace_weighted_layer_norm(graph: IRGraph, value_id: str) -> tuple[str, str, float] | None: - traced = _trace_weighted_norm( - graph, - value_id, - lambda g, candidate: _trace_norm_branch(g, candidate, _trace_layer_norm_centered_source, _trace_layer_norm_inv_std), - ) - if traced is not None: - return traced - return _trace_factored_weighted_norm( - graph, - value_id, - _trace_layer_norm_centered_source, - _trace_layer_norm_inv_std, - ) - - -def _rewrite_jax_rms_norms(graph: IRGraph) -> None: - for node_id in list(graph.order): - node = graph.nodes[node_id] - if node.op == "multiply" and len(node.inputs) == 2: - traced = _trace_weighted_rms_norm(graph, node.outputs[0] if node.outputs else "") - if traced is not None: - source_id, weight_id, eps = traced - _rewrite_norm_node(node, "rms_norm", [source_id, weight_id], eps, "jax_rms_norm_multiply") - continue - if node.op == "divide" and len(node.inputs) == 2: - denominator = _trace_rms_inv_std(graph, node.inputs[1]) - if denominator is None: - continue - source_id, eps = denominator - numerator = _producer(graph, _strip_simple_wrappers(graph, node.inputs[0])) - if numerator is None or numerator.op != "multiply" or len(numerator.inputs) != 2: - continue - for x_id, weight_id in ((numerator.inputs[0], numerator.inputs[1]), (numerator.inputs[1], numerator.inputs[0])): - if _strip_simple_wrappers(graph, x_id) != _strip_simple_wrappers(graph, source_id): - continue - weight_id = _valid_last_dim_weight(graph, weight_id, x_id) - if weight_id is not None: - _rewrite_norm_node(node, "rms_norm", [x_id, weight_id], eps, "jax_rms_norm") - break - - -def _rewrite_jax_layer_norms(graph: IRGraph) -> None: - for node_id in list(graph.order): - node = graph.nodes[node_id] - if node.op == "add" and len(node.inputs) == 2: - for affine_id, bias_id in ((node.inputs[0], node.inputs[1]), (node.inputs[1], node.inputs[0])): - bias_id = _strip_simple_wrappers(graph, bias_id) - bias_value = graph.values.get(bias_id) - if bias_value is None or bias_value.shape is None or len(bias_value.shape) != 1: - continue - traced_affine = _trace_weighted_layer_norm(graph, affine_id) - if traced_affine is None: - continue - source_id, weight_id, eps = traced_affine - source_value = graph.values.get(source_id) - if source_value is None or source_value.shape is None: - continue - if int(bias_value.shape[0]) != int(source_value.shape[-1]): - continue - _rewrite_norm_node(node, "layer_norm", [source_id, weight_id, bias_id], eps, "jax_layer_norm") - break - continue - if node.op == "multiply" and len(node.inputs) == 2: - traced = _trace_weighted_layer_norm(graph, node.outputs[0] if node.outputs else "") - if traced is not None: - source_id, weight_id, eps = traced - _rewrite_norm_node(node, "layer_norm", [source_id, weight_id], eps, "jax_layer_norm_multiply") - - -def _rewrite_jax_silus(graph: IRGraph) -> None: - for node_id in list(graph.order): - node = graph.nodes[node_id] - if node.op != "multiply" or len(node.inputs) != 2: - continue - for source_id, sigmoid_id in ((node.inputs[0], node.inputs[1]), (node.inputs[1], node.inputs[0])): - sigmoid_node = _producer(graph, _strip_simple_wrappers(graph, sigmoid_id)) - if sigmoid_node is None or sigmoid_node.op not in {"sigmoid", "logistic"} or len(sigmoid_node.inputs) != 1: - continue - if _strip_simple_wrappers(graph, source_id) != _strip_simple_wrappers(graph, sigmoid_node.inputs[0]): - continue - node.op = "silu" - node.inputs = [source_id] - node.attrs = {} - node.kind = "semantic" - node.meta = {**node.meta, "rewritten_from": "jax_silu"} - break - - -def _trace_batch_norm_centered(graph: IRGraph, value_id: str) -> tuple[str, str] | None: - node = _producer(graph, _strip_simple_wrappers(graph, value_id)) - if node is None or node.op != "subtract" or len(node.inputs) != 2: - return None - source_id = _trace_precision_cast_source(graph, _strip_simple_wrappers(graph, node.inputs[0])) - mean_id = _strip_simple_wrappers(graph, node.inputs[1]) - source_value = graph.values.get(source_id) - mean_value = graph.values.get(mean_id) - if source_value is None or source_value.shape is None: - return None - if mean_value is None or mean_value.shape is None or len(mean_value.shape) != 1: - return None - if int(source_value.shape[-1]) != int(mean_value.shape[0]): - return None - return source_id, mean_id - - -def _trace_batch_norm_inv_scale(graph: IRGraph, value_id: str) -> tuple[str, str, float] | None: - node = _producer(graph, _strip_simple_wrappers(graph, value_id)) - if node is None or node.op != "multiply" or len(node.inputs) != 2: - return None - for inv_id, scale_id in ((node.inputs[0], node.inputs[1]), (node.inputs[1], node.inputs[0])): - scale_id = _strip_simple_wrappers(graph, scale_id) - scale_value = graph.values.get(scale_id) - if scale_value is None or scale_value.shape is None or len(scale_value.shape) != 1: - continue - inv_node = _producer(graph, _strip_simple_wrappers(graph, inv_id)) - if inv_node is None or inv_node.op not in {"pow", "divide"} or not inv_node.inputs: - continue - if inv_node.op == "pow" and float(inv_node.attrs.get("exponent", 0.0)) != -0.5: - continue - add_id = inv_node.inputs[0] if inv_node.op == "pow" else inv_node.inputs[1] - add_node = _producer(graph, _strip_simple_wrappers(graph, add_id)) - if add_node is None or add_node.op != "add" or len(add_node.inputs) != 2: - continue - eps = 1.0e-5 - var_id: str | None = None - for add_input in add_node.inputs: - scalar = _constant_scalar_or_singleton(graph, add_input) - if scalar is not None: - eps = float(scalar) - continue - candidate = _strip_simple_wrappers(graph, add_input) - candidate_value = graph.values.get(candidate) - if candidate_value is not None and candidate_value.shape is not None and len(candidate_value.shape) == 1: - var_id = candidate - if var_id is None: - continue - var_value = graph.values.get(var_id) - if var_value is None or var_value.shape is None or int(var_value.shape[0]) != int(scale_value.shape[0]): - continue - return var_id, scale_id, eps - return None - - -def _trace_batch_norm_affine(graph: IRGraph, value_id: str) -> tuple[str, str, str, str, float] | None: - node = _producer(graph, _strip_simple_wrappers(graph, value_id)) - if node is None or node.op != "multiply" or len(node.inputs) != 2: - return None - for centered_id, inv_scale_id in ((node.inputs[0], node.inputs[1]), (node.inputs[1], node.inputs[0])): - centered = _trace_batch_norm_centered(graph, centered_id) - inv_scale = _trace_batch_norm_inv_scale(graph, inv_scale_id) - if centered is None or inv_scale is None: - continue - source_id, mean_id = centered - var_id, scale_id, eps = inv_scale - channel_count = int(graph.values[scale_id].shape[0]) # type: ignore[index] - if int(graph.values[mean_id].shape[0]) != channel_count: # type: ignore[index] - continue - if int(graph.values[var_id].shape[0]) != channel_count: # type: ignore[index] - continue - return source_id, scale_id, mean_id, var_id, eps - return None - - -def _rewrite_jax_batch_norms(graph: IRGraph) -> None: - for node_id in list(graph.order): - node = graph.nodes[node_id] - if node.op != "add" or len(node.inputs) != 2: - continue - for affine_id, bias_id in ((node.inputs[0], node.inputs[1]), (node.inputs[1], node.inputs[0])): - bias_id = _strip_simple_wrappers(graph, bias_id) - bias_value = graph.values.get(bias_id) - if bias_value is None or bias_value.shape is None or len(bias_value.shape) != 1: - continue - traced = _trace_batch_norm_affine(graph, affine_id) - if traced is None: - continue - source_id, scale_id, mean_id, var_id, eps = traced - source_value = graph.values.get(source_id) - if source_value is None or source_value.shape is None or not source_value.shape: - continue - channel_count = int(bias_value.shape[0]) - if int(source_value.shape[-1]) != channel_count: - continue - node.op = "batch_norm" - node.inputs = [source_id, scale_id, bias_id, mean_id, var_id] - node.attrs = {"axis": len(source_value.shape) - 1, "eps": eps} - node.kind = "semantic" - node.meta = {**node.meta, "rewritten_from": "jax_batch_norm"} - break - - -def _branch_uses_rms_norm_of(graph: IRGraph, branch_value_id: str, residual_value_id: str) -> bool: - residual_base = _strip_simple_wrappers(graph, residual_value_id) - stack = [branch_value_id] - visited: set[str] = set() - while stack and len(visited) < 512: - current = stack.pop() - if current in visited: - continue - visited.add(current) - node = _producer(graph, _strip_simple_wrappers(graph, current)) - if node is None: - continue - if node.op == "rms_norm" and node.inputs: - if _strip_simple_wrappers(graph, node.inputs[0]) == residual_base: - return True - stack.extend(node.inputs) - return False - - -def _rewrite_prenorm_residual_adds_to_clipped(graph: IRGraph) -> None: - for node_id in list(graph.order): - node = graph.nodes[node_id] - if node.op != "add" or len(node.inputs) != 2: - continue - lhs, rhs = node.inputs - if _branch_uses_rms_norm_of(graph, rhs, lhs) or _branch_uses_rms_norm_of(graph, lhs, rhs): - node.op = "add_clipped" - node.kind = "semantic" - node.meta = {**node.meta, "rewritten_from": "prenorm_residual_add"} - - -def _shape(graph: IRGraph, value_id: str) -> tuple[int, ...] | None: - value = graph.values.get(value_id) - if value is None or value.shape is None: - return None - return tuple(int(dim) for dim in value.shape) - - -def _trace_view_input(graph: IRGraph, value_id: str) -> str: - current = value_id - for _ in range(8): - node = _producer(graph, current) - if node is None or node.op not in {"view", "reshape", "precision_cast"} or not node.inputs: - return current - current = node.inputs[0] - return current - - -def _trace_gqa_repeat_base(graph: IRGraph, value_id: str) -> str: - current = _trace_view_input(graph, value_id) - shape = _shape(graph, current) - if shape is None or len(shape) != 5: - return current - expand = _producer(graph, current) - if expand is None or expand.op != "expand" or not expand.inputs: - return current - base = _trace_view_input(graph, expand.inputs[0]) - base_shape = _shape(graph, base) - if base_shape is None or len(base_shape) != 4: - return current - if base_shape[0] == shape[0] and base_shape[1] == shape[1] and base_shape[2] == shape[3] and base_shape[3] == shape[4]: - return base - return current - - -def _trace_jax_attention_logits(graph: IRGraph, value_id: str) -> tuple[str, str, float, str | None, bool] | None: - current = _trace_view_input(graph, value_id) - node = _producer(graph, current) - mask_id: str | None = None - additive_mask = False - if node is not None and node.op == "where" and len(node.inputs) >= 2: - mask_id = node.inputs[0] - current = _trace_view_input(graph, node.inputs[1]) - node = _producer(graph, current) - if node is not None and node.op == "add" and len(node.inputs) == 2: - for score_id, additive_mask_id in ((node.inputs[0], node.inputs[1]), (node.inputs[1], node.inputs[0])): - score_node = _producer(graph, _trace_view_input(graph, score_id)) - if score_node is not None and score_node.op == "divide": - mask_id = additive_mask_id - additive_mask = True - current = _trace_view_input(graph, score_id) - node = score_node - break - if node is None or node.op != "divide" or len(node.inputs) != 2: - return None - divisor = _constant_scalar_or_singleton(graph, node.inputs[1]) - if divisor is None or divisor == 0.0: - return None - matmul = _producer(graph, _trace_view_input(graph, node.inputs[0])) - if matmul is None or matmul.op != "matmul" or len(matmul.inputs) != 2: - return None - query_id = _trace_view_input(graph, matmul.inputs[0]) - key_t_id = _trace_view_input(graph, matmul.inputs[1]) - key_t_node = _producer(graph, key_t_id) - if key_t_node is None or key_t_node.op != "permute" or len(key_t_node.inputs) != 1: - return None - if tuple(int(dim) for dim in key_t_node.attrs.get("permutation", ())) != (0, 1, 3, 2): - return None - key_id = _trace_gqa_repeat_base(graph, key_t_node.inputs[0]) - return query_id, key_id, 1.0 / float(divisor), mask_id, additive_mask - - -def _trace_jax_attention_probs(graph: IRGraph, value_id: str) -> tuple[str, str, float, str | None, bool] | None: - current = _trace_view_input(graph, value_id) - cast = _producer(graph, current) - if cast is not None and cast.op == "precision_cast" and cast.inputs: - current = _trace_view_input(graph, cast.inputs[0]) - divide = _producer(graph, current) - if divide is None or divide.op != "divide" or len(divide.inputs) != 2: - return None - exp_node = _producer(graph, _trace_view_input(graph, divide.inputs[0])) - if exp_node is None or exp_node.op != "scalar_exp" or not exp_node.inputs: - return None - subtract = _producer(graph, _trace_view_input(graph, exp_node.inputs[0])) - if subtract is None or subtract.op != "subtract" or not subtract.inputs: - return None - return _trace_jax_attention_logits(graph, subtract.inputs[0]) - - -def _rewrite_attention_node( - node: IRNode, - query_id: str, - key_id: str, - value_id: str, - mask_id: str | None, - *, - additive_mask: bool, - scale: float, - output_layout: str, - is_causal: bool, - rewritten_from: str, -) -> None: - node.op = "attention" - node.inputs = [query_id, key_id, value_id] + ([] if mask_id is None else [mask_id]) - node.attrs = { - "scale": float(scale), - "is_causal": bool(is_causal), - "window_size": 0, - "q_layout": "bhsd", - "k_layout": "bhsd", - "v_layout": "bhsd", - "output_layout": output_layout, - **({"additive_mask": True} if additive_mask else {}), - } - node.kind = "semantic" - node.meta = {**node.meta, "rewritten_from": rewritten_from} - - -def _trace_attention_output(graph: IRGraph, node: IRNode) -> tuple[str, str, str, float, str | None, bool] | None: - if len(node.inputs) != 1 or len(node.outputs) != 1: - return None - matmul = _producer(graph, _trace_view_input(graph, node.inputs[0])) - if matmul is None or matmul.op != "matmul" or len(matmul.inputs) != 2: - return None - traced = _trace_jax_attention_probs(graph, matmul.inputs[0]) - if traced is None: - return None - query_id, key_id, scale, mask_id, additive_mask = traced - return query_id, key_id, _trace_gqa_repeat_base(graph, matmul.inputs[1]), scale, mask_id, additive_mask - - -def _valid_attention_qkv_shapes( - query_shape: tuple[int, ...] | None, - key_shape: tuple[int, ...] | None, - value_shape: tuple[int, ...] | None, -) -> bool: - if query_shape is None or key_shape is None or value_shape is None: - return False - if len(query_shape) != 4 or len(key_shape) != 4 or len(value_shape) != 4: - return False - if query_shape[0] != key_shape[0] or key_shape[0] != value_shape[0]: - return False - return key_shape[1:] == value_shape[1:] - - -def _valid_attention_mask_shape( - mask_shape: tuple[int, ...] | None, - query_shape: tuple[int, ...], - key_shape: tuple[int, ...], -) -> bool: - return mask_shape == (query_shape[0], query_shape[1], query_shape[2], key_shape[2]) - - -def _rewrite_jax_prefill_attentions(graph: IRGraph) -> None: - for node_id in list(graph.order): - node = graph.nodes[node_id] - if node.op not in {"permute", "view", "reshape"}: - continue - component = str(graph.meta.get("component", "") or "").strip().lower() - is_decoder_step = component == "decoder_step" - is_causal_prefill = component in {"decoder_prefill_chunk", "decoder_media_step"} - if node.op == "permute" and tuple(int(dim) for dim in node.attrs.get("permutation", ())) != (0, 2, 1, 3): - continue - if node.op in {"view", "reshape"} and not is_decoder_step: - continue - output_shape = _shape(graph, node.outputs[0]) - if output_shape is None or len(output_shape) != 4: - continue - traced = _trace_attention_output(graph, node) - if traced is None: - continue - query_id, key_id, value_id, scale, mask_id, additive_mask = traced - query_shape = _shape(graph, query_id) - key_shape = _shape(graph, key_id) - value_shape = _shape(graph, value_id) - if not _valid_attention_qkv_shapes(query_shape, key_shape, value_shape): - continue - assert query_shape is not None and key_shape is not None - if mask_id is not None and not _valid_attention_mask_shape(_shape(graph, mask_id), query_shape, key_shape): - continue - is_self_attention = query_shape[2] == key_shape[2] - if is_decoder_step: - if query_shape[2] != 1 or (not is_self_attention and mask_id is None): - continue - _rewrite_attention_node( - node, - query_id, - key_id, - value_id, - mask_id, - additive_mask=additive_mask, - scale=scale, - output_layout="bhsd", - is_causal=is_self_attention, - rewritten_from="jax_decoder_step_self_attention" if is_self_attention else "jax_decoder_step_cross_attention", - ) - continue - if output_shape != (query_shape[0], query_shape[2], query_shape[1], query_shape[3]): - continue - _rewrite_attention_node( - node, - query_id, - key_id, - value_id, - mask_id, - additive_mask=additive_mask, - scale=scale, - output_layout="bthd", - is_causal=is_causal_prefill and mask_id is None and is_self_attention, - rewritten_from="jax_prefill_attention", - ) - - -def _trace_precision_cast_source(graph: IRGraph, value_id: str) -> str: - producer = _producer(graph, value_id) - if producer is not None and producer.op == "precision_cast" and producer.inputs: - return producer.inputs[0] - return value_id - - -def _prune_dead_nodes(graph: IRGraph) -> None: - live_values = set(graph.outputs) - live_nodes: set[str] = set() - stack = list(graph.outputs) - while stack: - value_id = stack.pop() - value = graph.values.get(value_id) - if value is None or value.producer is None: - continue - node = graph.nodes.get(value.producer) - if node is None or node.id in live_nodes: - continue - live_nodes.add(node.id) - for input_id in node.inputs: - if input_id not in live_values: - live_values.add(input_id) - stack.append(input_id) - - removed_nodes = set(graph.nodes) - live_nodes - if not removed_nodes: - return - for node_id in removed_nodes: - node = graph.nodes.pop(node_id) - for output_id in node.outputs: - graph.values.pop(output_id, None) - graph.constants.pop(output_id, None) - graph.order = [node_id for node_id in graph.order if node_id in live_nodes] - - -def _unique_graph_id(existing: set[str], base: str) -> str: - candidate = base - suffix = 0 - while candidate in existing: - suffix += 1 - candidate = f"{base}_{suffix}" - return candidate - - -def _legalize_remaining_jax_reductions(graph: IRGraph) -> None: - new_order: list[str] = [] - for node_id in list(graph.order): - node = graph.nodes[node_id] - if node.op not in {"sum", "mean"} or len(node.inputs) != 1 or len(node.outputs) != 1: - new_order.append(node_id) - continue - input_value = graph.values.get(node.inputs[0]) - output_value = graph.values.get(node.outputs[0]) - if input_value is None or output_value is None or input_value.dtype == "fp16": - new_order.append(node_id) - continue - - cast_input_id = _unique_graph_id(set(graph.values), f"{node.outputs[0]}__reduce_fp16_input") - cast_input_node_id = _unique_graph_id(set(graph.nodes), f"{node.id}_reduce_input_cast") - graph.add_node( - IRNode( - cast_input_node_id, - "precision_cast", - [node.inputs[0]], - [cast_input_id], - attrs={"dtype": "fp16"}, - meta={ - "jax_reduce_legalization": "input_fp16", - "source_dtype": input_value.dtype, - }, - ) - ) - graph.values[cast_input_id].shape = input_value.shape - graph.values[cast_input_id].dtype = "fp16" - node.inputs = [cast_input_id] - new_order.append(cast_input_node_id) - new_order.append(node_id) - - if output_value.dtype == "fp16": - continue - - original_output_id = node.outputs[0] - reduce_output_id = _unique_graph_id(set(graph.values), f"{original_output_id}__reduce_fp16") - node.outputs = [reduce_output_id] - graph.values[reduce_output_id] = IRValue( - id=reduce_output_id, - shape=output_value.shape, - dtype="fp16", - producer=node.id, - meta={ - "jax_reduce_legalization": "fp16_output", - "source_output": original_output_id, - }, - ) - cast_output_node_id = _unique_graph_id(set(graph.nodes), f"{node.id}_reduce_output_cast") - graph.nodes[cast_output_node_id] = IRNode( - cast_output_node_id, - "precision_cast", - [reduce_output_id], - [original_output_id], - attrs={"dtype": output_value.dtype}, - meta={ - "jax_reduce_legalization": "restore_output_dtype", - "source_dtype": "fp16", - }, - ) - output_value.producer = cast_output_node_id - new_order.append(cast_output_node_id) - graph.order = new_order - - -JAX_PATTERNS = ( - JaxPattern("batch_norm", _rewrite_jax_batch_norms), - JaxPattern("layer_norm", _rewrite_jax_layer_norms), - JaxPattern("rms_norm", _rewrite_jax_rms_norms), - JaxPattern("silu", _rewrite_jax_silus), - JaxPattern("attention", _rewrite_jax_prefill_attentions), - JaxPattern("prenorm_add_clipped", _rewrite_prenorm_residual_adds_to_clipped), -) - - -def _apply_pattern(graph: IRGraph, pattern: JaxPattern) -> None: - pattern.apply(graph) - - -def apply_jax_semantic_rewrites(graph: IRGraph) -> None: - for pattern in JAX_PATTERNS: - _apply_pattern(graph, pattern) - apply_import_semantics(graph) - _prune_dead_nodes(graph) - _legalize_remaining_jax_reductions(graph) - _rebuild_users(graph) diff --git a/python/cactus/transpile/jax_user_graph_bundle.py b/python/cactus/transpile/jax_user_graph_bundle.py deleted file mode 100644 index b06dc3657..000000000 --- a/python/cactus/transpile/jax_user_graph_bundle.py +++ /dev/null @@ -1,575 +0,0 @@ -from __future__ import annotations - -from collections.abc import Sequence -from dataclasses import dataclass -from dataclasses import fields -import json -import os -import tempfile -from pathlib import Path -from typing import Any - -import numpy as np -import torch - -from cactus.convert.cactus_adapters.tensor_io import save_tensor_with_header -from cactus.transpile.capture_jax import capture_jax_graphs -from cactus.transpile.capture_jax import JaxGraphSpec -from cactus.transpile.graph_ir import IRGraph -from cactus.transpile.lower import TranspiledGraph -from cactus.transpile.runtime_compat import Graph -from cactus.transpile.runtime_compat import Tensor - - -@dataclass(frozen=True) -class JaxUserGraphBundleResult: - bundle: Any - output_dir: Path - components_manifest_path: Path - weights_dir: Path - - -@dataclass -class LoadedJaxUserGraph: - name: str - graph: Graph - runtime_inputs: list[Tensor] - outputs: list[Tensor] - logical_inputs: list[str] - logical_outputs: list[str] - - def set_inputs(self, inputs: Sequence[Any]) -> None: - if len(inputs) != len(self.runtime_inputs): - raise ValueError( - f"graph {self.name!r} expected {len(self.runtime_inputs)} inputs, got {len(inputs)}" - ) - for tensor, value in zip(self.runtime_inputs, inputs, strict=True): - if isinstance(value, Tensor): - if tuple(int(dim) for dim in value.shape) != tuple(int(dim) for dim in tensor.shape): - raise ValueError( - f"graph {self.name!r} input shape mismatch for node {tensor.id}: " - f"expected {tensor.shape}, got {value.shape}" - ) - if int(value.dtype) != int(tensor.dtype): - raise ValueError( - f"graph {self.name!r} input dtype mismatch for node {tensor.id}: " - f"expected {tensor.dtype}, got {value.dtype}" - ) - self.graph.set_input(tensor, value.numpy()) - else: - self.graph.set_input(tensor, np.asarray(value)) - - def execute(self, *inputs: Any) -> list[Tensor]: - self.set_inputs(inputs) - self.graph.execute() - return self.outputs - - def reset(self) -> None: - self.graph.hard_reset() - - -@dataclass -class LoadedJaxUserGraphBundle: - root: Path - manifest_path: Path - manifest: dict[str, object] - graphs: dict[str, LoadedJaxUserGraph] - - def execute(self, graph_name: str, *inputs: Any) -> list[Tensor]: - if graph_name not in self.graphs: - available = ", ".join(sorted(self.graphs)) or "" - raise ValueError(f"unknown JAX user graph {graph_name!r}; available graphs: {available}") - return self.graphs[graph_name].execute(*inputs) - - def reset(self, graph_name: str | None = None) -> None: - if graph_name is not None: - if graph_name not in self.graphs: - available = ", ".join(sorted(self.graphs)) or "" - raise ValueError(f"unknown JAX user graph {graph_name!r}; available graphs: {available}") - self.graphs[graph_name].reset() - return - for graph in self.graphs.values(): - graph.reset() - - -def flatten_jax_params(params: object) -> dict[str, np.ndarray]: - import jax - - return { - _tree_path_name(path): np.asarray(value) - for path, value in jax.tree_util.tree_flatten_with_path(params)[0] - } - - -def _tree_path_name(path: Sequence[Any]) -> str: - parts: list[str] = [] - for entry in path: - if hasattr(entry, "key"): - part = str(entry.key) - elif hasattr(entry, "name"): - part = str(entry.name) - elif hasattr(entry, "idx"): - part = str(entry.idx) - else: - part = str(entry) - parts.append(part) - return ".".join(parts) or "param" - - -def write_fp16_weights_manifest( - weights_dir: str | Path, - params: dict[str, object], - *, - exclude: set[str] | None = None, -) -> Path: - weights_root = Path(weights_dir) - weights_root.mkdir(parents=True, exist_ok=True) - excluded = exclude or set() - manifest: dict[str, dict[str, str]] = {} - for name, value in params.items(): - if name in excluded: - continue - filename = f"{name}.weights" - save_tensor_with_header(np.asarray(value), weights_root / filename, precision="FP16") - manifest[name] = {"filename": filename, "kind": "weight"} - manifest_path = weights_root / "weights_manifest.json" - _atomic_write_text(manifest_path, json.dumps(manifest, indent=2) + "\n") - return manifest_path - - -def load_jax_user_graph_bundle(bundle_dir_or_manifest: str | Path) -> LoadedJaxUserGraphBundle: - manifest_path = _resolve_components_manifest(bundle_dir_or_manifest) - root = manifest_path.parent.parent - manifest = json.loads(manifest_path.read_text()) - weights_dir = _resolve_manifest_path(root, manifest.get("weights_dir")) - graphs: dict[str, LoadedJaxUserGraph] = {} - for component in manifest.get("components", []): - if not isinstance(component, dict): - continue - name = str(component.get("component", "") or "") - graph_relpath = component.get("graph") - if not name or not isinstance(graph_relpath, str) or not graph_relpath: - continue - graph = Graph.load(_resolve_manifest_path(root, graph_relpath)) - for binding in component.get("bound_constant_bindings", []) or []: - if not isinstance(binding, dict): - continue - node_id = binding.get("node_id") - path = binding.get("path") - if not isinstance(node_id, int) or not isinstance(path, str): - continue - graph.bind_mmap_weights( - graph._tensor_from_node(int(node_id)), - _resolve_weight_path(root, weights_dir, path), - ) - graphs[name] = LoadedJaxUserGraph( - name=name, - graph=graph, - runtime_inputs=[ - graph._tensor_from_node(int(node_id)) - for node_id in component.get("runtime_input_node_ids", []) or [] - ], - outputs=[ - graph._tensor_from_node(int(node_id)) - for node_id in component.get("output_node_ids", []) or [] - ], - logical_inputs=[str(value) for value in component.get("logical_inputs", []) or []], - logical_outputs=[str(value) for value in component.get("logical_outputs", []) or []], - ) - return LoadedJaxUserGraphBundle( - root=root, - manifest_path=manifest_path, - manifest=manifest, - graphs=graphs, - ) - - -def _resolve_components_manifest(bundle_dir_or_manifest: str | Path) -> Path: - path = Path(bundle_dir_or_manifest).expanduser() - if path.is_dir(): - components_manifest = path / "components" / "manifest.json" - if components_manifest.exists(): - return components_manifest - if path.name == "components" and (path / "manifest.json").exists(): - return path / "manifest.json" - if path.exists(): - return path - raise FileNotFoundError(f"JAX user graph bundle manifest not found: {path}") - - -def _reject_relative_escape(value: str, *, field: str) -> None: - """Reject relative manifest entries that try to escape via `..`. Absolute - paths are a legitimate workflow (shared weights cache) and stay allowed.""" - if not value: - return - p = Path(value) - if p.is_absolute(): - return - if any(part == ".." for part in p.parts): - raise RuntimeError( - f"refusing to follow bundle manifest relative path with '..' segments: " - f"{field}={value}" - ) - - -def _resolve_manifest_path(root: Path, value: object) -> str: - if not isinstance(value, str) or not value: - return str(root) - _reject_relative_escape(value, field="manifest path") - path = Path(value).expanduser() - return str(path if path.is_absolute() else root / path) - - -def _resolve_weight_path(root: Path, weights_dir: str, value: str) -> str: - _reject_relative_escape(value, field="weight path") - path = Path(value).expanduser() - if path.is_absolute(): - return str(path) - candidate = Path(weights_dir) / path - if candidate.exists(): - return str(candidate) - return str(root / path) - - -def _embed_materialized_bound_constants(transpiled_graph: TranspiledGraph) -> None: - external_node_ids = { - int(binding["node_id"]) - for binding in transpiled_graph.bound_constant_bindings - } - for constant in transpiled_graph.bound_constants: - if int(constant.id) in external_node_ids: - continue - transpiled_graph.graph.mark_embedded_input(constant) - - -def _serialize_json_compatible(value: Any) -> Any: - if isinstance(value, dict): - return {str(key): _serialize_json_compatible(inner) for key, inner in value.items()} - if isinstance(value, (list, tuple)): - return [_serialize_json_compatible(inner) for inner in value] - if isinstance(value, Path): - return str(value) - if isinstance(value, (str, int, float, bool)) or value is None: - return value - if isinstance(value, torch.dtype): - return str(value) - if isinstance(value, torch.Tensor): - payload: dict[str, object] = { - "type": "torch.Tensor", - "dtype": str(value.dtype), - "shape": list(value.shape), - } - if value.numel() == 1: - payload["data"] = value.detach().cpu().reshape(-1)[0].item() - return payload - if isinstance(value, np.ndarray): - payload = { - "type": "numpy.ndarray", - "dtype": str(value.dtype), - "shape": list(value.shape), - } - if value.size == 1: - payload["data"] = np.asarray(value).reshape(-1)[0].item() - return payload - if hasattr(value, "__dataclass_fields__"): - return { - field.name: _serialize_json_compatible(getattr(value, field.name)) - for field in fields(value) - } - try: - return repr(value) - except Exception: - return f"<{type(value).__module__}.{type(value).__name__}>" - - -def _ir_graph_to_dict(graph: IRGraph) -> dict[str, object]: - return { - "meta": _serialize_json_compatible(graph.meta), - "inputs": list(graph.inputs), - "outputs": list(graph.outputs), - "constants": { - value_id: _serialize_json_compatible(constant) - for value_id, constant in graph.constants.items() - }, - "values": { - value_id: _serialize_json_compatible(value) - for value_id, value in graph.values.items() - }, - "nodes": [ - _serialize_json_compatible(graph.nodes[node_id]) - for node_id in graph.order - ], - } - - -def _atomic_write_text(path: Path, text: str) -> None: - """Write text to path via tempfile + os.replace so a crash mid-write - cannot leave a half-written file the next loader would choke on.""" - path.parent.mkdir(parents=True, exist_ok=True) - fd, tmp = tempfile.mkstemp(prefix=path.name + ".", suffix=".tmp", dir=str(path.parent)) - try: - with os.fdopen(fd, "w") as f: - f.write(text) - os.replace(tmp, path) - except Exception: - try: - os.unlink(tmp) - except OSError: - pass - raise - - -def _write_json(path: Path, payload: dict[str, object]) -> None: - _atomic_write_text(path, json.dumps(payload, indent=2, sort_keys=True) + "\n") - - -def build_jax_user_graph_bundle( - *, - params: Any, - specs: Sequence[JaxGraphSpec], - output_dir: str | Path, - model_id: str, - task: str = "generic", - family: str = "jax_user_graph", - inputs_metadata: dict[str, object] | None = None, - graph_meta: dict[str, object] | None = None, - weight_arrays: dict[str, object] | None = None, - exclude_weights: set[str] | None = None, - weights_dir: str | Path | None = None, -) -> JaxUserGraphBundleResult: - arrays = weight_arrays if weight_arrays is not None else flatten_jax_params(params) - root = Path(output_dir) - weights_root = Path(weights_dir) if weights_dir is not None else root - write_fp16_weights_manifest(weights_root, arrays, exclude=exclude_weights) - - bundle = capture_jax_graphs( - params, - specs, - weights_dir=str(weights_root), - graph_meta={ - "frontend": "jax", - "adapter_family": "generic", - "graph_family": "jax_user_graph_bundle", - **dict(graph_meta or {}), - }, - ) - return _write_jax_user_graph_bundle( - bundle=bundle, - specs=specs, - output_dir=root, - weights_dir=weights_root, - model_id=model_id, - task=task, - family=family, - inputs_metadata=inputs_metadata, - ) - - -def build_jax_generation_graph_bundle( - *, - params: Any, - output_dir: str | Path, - model_id: str, - encoder: JaxGraphSpec | None = None, - decoder_prefill: JaxGraphSpec | None = None, - decoder_step: JaxGraphSpec | None = None, - task: str = "text-generation", - family: str = "jax_user_graph", - inputs_metadata: dict[str, object] | None = None, - graph_meta: dict[str, object] | None = None, - weight_arrays: dict[str, object] | None = None, - exclude_weights: set[str] | None = None, - weights_dir: str | Path | None = None, -) -> JaxUserGraphBundleResult: - root = Path(output_dir) - weights_root = Path(weights_dir) if weights_dir is not None else root - arrays = weight_arrays if weight_arrays is not None else flatten_jax_params(params) - write_fp16_weights_manifest(weights_root, arrays, exclude=exclude_weights) - - input_specs = tuple(spec for spec in (encoder, decoder_prefill, decoder_step) if spec is not None) - bundle = capture_jax_generation_graphs( - params, - encoder=encoder, - decoder_prefill=decoder_prefill, - decoder_step=decoder_step, - weights_dir=str(weights_root), - graph_meta={ - "frontend": "jax", - "adapter_family": "generic", - "graph_family": "jax_user_graph_bundle", - **dict(graph_meta or {}), - }, - ) - return _write_jax_user_graph_bundle( - bundle=bundle, - specs=input_specs, - output_dir=root, - weights_dir=weights_root, - model_id=model_id, - task=task, - family=family, - inputs_metadata=inputs_metadata, - ) - - -def capture_jax_generation_graphs( - params: Any, - *, - encoder: JaxGraphSpec | None = None, - decoder_prefill: JaxGraphSpec | None = None, - decoder_step: JaxGraphSpec | None = None, - weights_dir: str | None = None, - weight_bindings: dict[str, dict[str, str]] | None = None, - graph_meta: dict[str, object] | None = None, -) -> Any: - specs: list[JaxGraphSpec] = [] - - def _with_meta( - spec: JaxGraphSpec | None, - *, - name: str, - role: str, - component: str, - ) -> None: - if spec is None: - return - specs.append( - JaxGraphSpec( - name=spec.name or name, - role=spec.role if spec.role != "generic" else role, - fn=spec.fn, - example_args=spec.example_args, - input_names=spec.input_names, - output_names=spec.output_names, - graph_meta={ - "component": component, - **dict(spec.graph_meta or {}), - }, - ) - ) - - _with_meta(encoder, name="encoder", role="encoder", component="encoder") - _with_meta( - decoder_prefill, - name="decoder_prefill", - role="decoder_prefill", - component="decoder_prefill_chunk", - ) - _with_meta( - decoder_step, - name="decoder_step", - role="decoder_step", - component="decoder_step", - ) - if not specs: - raise ValueError("capture_jax_generation_graphs requires at least one graph spec") - return capture_jax_graphs( - params, - specs, - weights_dir=weights_dir, - weight_bindings=weight_bindings, - graph_meta={ - "frontend": "jax", - "adapter_family": "generic", - "graph_family": "generation", - **dict(graph_meta or {}), - }, - ) - - -def _write_jax_user_graph_bundle( - *, - bundle: Any, - specs: Sequence[JaxGraphSpec], - output_dir: Path, - weights_dir: Path, - model_id: str, - task: str, - family: str, - inputs_metadata: dict[str, object] | None, -) -> JaxUserGraphBundleResult: - root = output_dir - weights_root = weights_dir - component_root = root / "components" - component_root.mkdir(parents=True, exist_ok=True) - - components = [] - component_order = [spec.name for spec in specs] - base_payload = { - "model_id": model_id, - "model_source": "jax_user_graph", - "task": task, - "family": family, - "inputs": _serialize_json_compatible(dict(inputs_metadata or {})), - } - for name, captured in bundle.graphs.items(): - component_dir = component_root / name - component_dir.mkdir(parents=True, exist_ok=True) - graph_path = component_dir / "graph.cactus" - raw_ir_path = component_dir / "raw_ir.json" - optimized_ir_path = component_dir / "optimized_ir.json" - _write_json( - raw_ir_path, - { - **base_payload, - "component": name, - "graph": _ir_graph_to_dict(captured.raw_ir_graph), - }, - ) - _write_json( - optimized_ir_path, - { - **base_payload, - "component": name, - "graph": _ir_graph_to_dict(captured.ir_graph), - }, - ) - _embed_materialized_bound_constants(captured.graph) - captured.graph.graph.save(str(graph_path)) - components.append( - { - "component": name, - "directory": str(component_dir.relative_to(root)), - "raw_ir": str(raw_ir_path.relative_to(root)), - "optimized_ir": str(optimized_ir_path.relative_to(root)), - "graph": str(graph_path.relative_to(root)), - "inputs": list(captured.ir_graph.inputs), - "outputs": list(captured.ir_graph.outputs), - "logical_inputs": list(captured.spec.input_names or ()), - "logical_outputs": list(captured.spec.output_names or ()), - "node_count": len(captured.ir_graph.order), - "weight_binding_count": len(captured.graph.bound_constant_bindings), - "runtime_input_node_ids": [int(tensor.id) for tensor in captured.graph.runtime_inputs], - "output_node_ids": [int(tensor.id) for tensor in captured.graph.outputs], - "cache_state_node_ids": [ - { - "layer_key": str(layer_key), - "key": int(key_tensor.id), - "value": int(value_tensor.id), - } - for layer_key, key_tensor, value_tensor in getattr(captured.graph, "cache_state_tensors", []) - ], - "bound_constant_bindings": captured.graph.bound_constant_bindings, - } - ) - - manifest = { - "model_id": model_id, - "model_source": "jax_user_graph", - "task": task, - "family": family, - "component_order": component_order, - "inputs": _serialize_json_compatible(dict(inputs_metadata or {})), - "components": components, - "weights_dir": str(weights_root), - "weights_manifest": str(weights_root / "weights_manifest.json"), - } - manifest_path = component_root / "manifest.json" - _atomic_write_text(manifest_path, json.dumps(manifest, indent=2) + "\n") - - return JaxUserGraphBundleResult( - bundle=bundle, - output_dir=root, - components_manifest_path=manifest_path, - weights_dir=weights_root, - ) diff --git a/python/cactus/transpile/lower.py b/python/cactus/transpile/lower.py deleted file mode 100644 index 4d195310f..000000000 --- a/python/cactus/transpile/lower.py +++ /dev/null @@ -1,3437 +0,0 @@ -from __future__ import annotations - -from dataclasses import dataclass -from dataclasses import field -import math -import os -from typing import Any - -import numpy as np -import torch - -from cactus.transpile.runtime_compat import Graph -from cactus.transpile.runtime_compat import Tensor -from cactus.transpile.capture_pytorch import CapturedModel -from cactus.transpile.canonicalize.cleanup import canonicalize_exported_graph -from cactus.transpile.graph_ir import IRGraph -from cactus.transpile.graph_ir import IRNode -from cactus.transpile.graph_ir import IRValue -from cactus.transpile.graph_ir import verify_ir -from cactus.transpile.optimize_graph import optimize_graph -from cactus.transpile.weight_compat import ensure_binding_compatible -from cactus.transpile.weight_binding import WeightBinding - - -_DYNAMIC_KV_POSITION_OFFSET = (1 << 64) - 1 -_DYNAMIC_KV_QUERY_POSITION_OFFSET = (1 << 64) - 2 -_KV_CACHE_TRANSPARENT_OPS = { - "expand", - "precision_cast", - "reshape", - "slice", - "transpose", - "unsqueeze", - "view", -} - - -@dataclass -class TranspiledGraph: - graph: Graph - runtime_inputs: list[Tensor] - bound_constants: list[Tensor] - bound_constant_bindings: list[dict[str, object]] - bound_constant_value_ids: dict[int, str] - outputs: list[Tensor] - cache_state_tensors: list[tuple[str, Tensor, Tensor]] = field(default_factory=list) - - def set_input(self, index: int, data: Any, *, dtype: int | None = None) -> None: - if index < 0 or index >= len(self.runtime_inputs): - raise IndexError( - f"runtime input index out of range: {index} (have {len(self.runtime_inputs)})" - ) - self.graph.set_input(self.runtime_inputs[index], data, dtype=dtype) - - def set_inputs(self, inputs: list[Any] | tuple[Any, ...]) -> None: - if len(inputs) != len(self.runtime_inputs): - raise ValueError( - f"expected {len(self.runtime_inputs)} runtime inputs, got {len(inputs)}" - ) - for index, value in enumerate(inputs): - self.set_input(index, value) - - def execute(self) -> list[Tensor]: - self.graph.execute() - return self.outputs - - -@dataclass -class BroadcastAlias: - tensor: Tensor - logical_shape: tuple[int, ...] - kind: str - - -def _is_quantized_runtime_dtype(dtype: int) -> bool: - return int(dtype) in { - int(Graph.INT8), - int(getattr(Graph, "CQ1", 3)), - int(getattr(Graph, "CQ2", 4)), - int(getattr(Graph, "CQ3", 5)), - int(getattr(Graph, "CQ4", 6)), - } - - -def transpile_captured(captured: CapturedModel) -> TranspiledGraph: - return transpile_ir(captured.ir_graph) - - -def transpile_ir(ir: IRGraph) -> TranspiledGraph: - verify_ir(ir) - canonicalize_exported_graph(ir) - optimize_graph(ir) - return transpile_preoptimized_ir(ir) - - -def transpile_preoptimized_ir(ir: IRGraph) -> TranspiledGraph: - verify_ir(ir) - g = Graph() - g._transpile_materialized_constants = [] # type: ignore[attr-defined] - trace_lower = os.environ.get("CACTUS_LOWER_TRACE_NODES", "0") != "0" - env: dict[str, Any] = {} - runtime_inputs: list[Tensor] = [] - bound_constants: list[Tensor] = [] - bound_constant_bindings: list[dict[str, object]] = [] - bound_constant_value_ids: dict[int, str] = {} - - for input_index, value_id in enumerate(ir.inputs): - value = ir.values[value_id] - tensor = _lower_input_value(g, value, ir=ir, input_index=input_index) - env[value_id] = tensor - runtime_inputs.append(tensor) - - for value_id, const in ir.constants.items(): - value = ir.values[value_id] - binding = _lookup_weight_binding(value) - if binding is not None: - binding = ensure_binding_compatible(binding, source_tensor=const) - lowered_const = _lower_constant_value(g, value, const, binding=binding) - env[value_id] = lowered_const - if isinstance(lowered_const, Tensor): - bound_constants.append(lowered_const) - bound_constant_value_ids[int(lowered_const.id)] = str(value_id) - if binding is not None: - bound_constant_bindings.append( - { - "node_id": int(lowered_const.id), - "value_id": str(value_id), - "path": binding.path, - "kind": binding.kind, - "source_name": binding.source_name, - } - ) - - for node_id in ir.order: - node = ir.nodes[node_id] - outputs = _lower_ir_node(g, node, env, ir) - if len(outputs) != len(node.outputs): - raise ValueError( - f"node {node.id} produced {len(outputs)} outputs, expected {len(node.outputs)}" - ) - for output_id, tensor in zip(node.outputs, outputs): - env[output_id] = tensor - if trace_lower and isinstance(tensor, Tensor): - print( - "lower_trace" - f" cactus_id={int(tensor.id)}" - f" ir_node={node.id}" - f" op={node.op}" - f" output={output_id}" - f" shape={tuple(int(dim) for dim in tensor.shape)}", - flush=True, - ) - - outputs = [env[value_id] for value_id in ir.outputs] - seen_bound_constant_ids = {int(tensor.id) for tensor in bound_constants} - for tensor in getattr(g, "_transpile_materialized_constants", []): - tensor_id = int(tensor.id) - if tensor_id in seen_bound_constant_ids: - continue - bound_constants.append(tensor) - seen_bound_constant_ids.add(tensor_id) - cache_state_tensors = list(env.get("__internal_kv_cache_state_entries", [])) - return TranspiledGraph( - graph=g, - runtime_inputs=runtime_inputs, - bound_constants=bound_constants, - bound_constant_bindings=bound_constant_bindings, - bound_constant_value_ids=bound_constant_value_ids, - outputs=outputs, - cache_state_tensors=cache_state_tensors, - ) - - -def _cross_kv_cache_buffer_elements(max_seq_len: int, num_kv_heads: int, head_dim: int) -> int: - kv_quant_group_size = 32 - num_groups = (int(head_dim) + kv_quant_group_size - 1) // kv_quant_group_size - kv_values = int(max_seq_len) * int(num_kv_heads) * int(head_dim) - kv_scales = int(max_seq_len) * int(num_kv_heads) * num_groups * 4 - return int(64 + kv_values + kv_scales) - - -def _external_cross_kv_cache_input_indices(ir: IRGraph) -> set[int]: - if not bool(ir.meta.get("external_cross_kv_cache_inputs", False)): - return set() - start = int(ir.meta.get("cross_kv_input_start_index", -1)) - count = int(ir.meta.get("cross_kv_input_count", 0)) - if start < 0 or count <= 0: - return set() - return set(range(start, start + count)) - - -def _is_external_cross_kv_cache_value(ir: IRGraph, value_id: str) -> bool: - return any( - input_value_id == value_id and input_index in _external_cross_kv_cache_input_indices(ir) - for input_index, input_value_id in enumerate(ir.inputs) - ) - - -def _lower_input_value(g: Graph, value: IRValue, *, ir: IRGraph | None = None, input_index: int | None = None) -> Tensor: - if value.shape is None or value.dtype is None: - raise ValueError(f"IR input missing shape or dtype: {value.id}") - if ir is not None and input_index is not None and input_index in _external_cross_kv_cache_input_indices(ir): - shape = tuple(int(dim) for dim in value.shape) - if len(shape) != 4: - raise ValueError(f"external cross-KV cache input expects original rank-4 K/V tensor, got {shape}") - layout = str(ir.meta.get("cross_kv_input_layout", "bthd") or "bthd").lower() - if layout == "bhsd": - max_seq_len, num_kv_heads, head_dim = shape[2], shape[1], shape[3] - else: - max_seq_len, num_kv_heads, head_dim = shape[1], shape[2], shape[3] - cache_elements = _cross_kv_cache_buffer_elements(max_seq_len, num_kv_heads, head_dim) - return g.input(shape=(cache_elements,), dtype=Graph.INT8) - dynamic_dims = value.meta.get("dynamic_dims") if value.meta else None - return g.input(shape=value.shape, dtype=_map_ir_dtype(value.dtype), dynamic_dims=dynamic_dims) - - -def _should_lower_attention_with_internal_kv_cache(ir: IRGraph, node: IRNode) -> bool: - if not bool(ir.meta.get("use_internal_kv_cache", False)): - return False - if node.op not in {"attention", "scaled_dot_product_attention"}: - return False - component = str(ir.meta.get("component", "") or "").strip().lower() - if component not in {"decoder_step", "decoder_prefill_chunk", "decoder_media_step", - "decoder_embed_chunk"}: - return False - if len(node.inputs) < 3: - return False - - def _sequence_dim(layout: str) -> int: - return 1 if layout == "bthd" else 2 - - q_layout = str(node.attrs.get("q_layout", node.attrs.get("qkv_layout", "bhsd")) or "bhsd").lower() - k_layout = str(node.attrs.get("k_layout", node.attrs.get("qkv_layout", "bhsd")) or "bhsd").lower() - query_value = ir.values.get(node.inputs[0]) - key_value = ir.values.get(node.inputs[1]) - query_shape = tuple(int(dim) for dim in (query_value.shape or ())) if query_value is not None else () - key_shape = tuple(int(dim) for dim in (key_value.shape or ())) if key_value is not None else () - if len(query_shape) >= 4 and len(key_shape) >= 4: - query_seq = int(query_shape[_sequence_dim(q_layout)]) - key_seq = int(key_shape[_sequence_dim(k_layout)]) - # Seq2seq decoders interleave self-attention (Q/K over the current - # decoder tokens) with cross-attention (Q over decoder tokens, K/V over - # encoder frames). Only the self-attention stream belongs in the - # autoregressive KV cache; cross-attention must continue to read the - # full encoder output every step. - if query_seq != key_seq: - return False - return True - - -def _kv_cache_layer_key(node: IRNode) -> str: - for meta_key in ("attention_layer_index", "layer_index"): - value = node.meta.get(meta_key) - if value is not None: - return str(value) - return str(node.id) - - -def _lower_attention_with_internal_kv_cache( - g: Graph, - ir: IRGraph, - env: dict[str, Any], - node: IRNode, - *, - query: Tensor, - key: Tensor, - value: Tensor, - output_layout: str, -) -> list[Tensor]: - key_shape = tuple(int(dim) for dim in key.shape) - value_shape = tuple(int(dim) for dim in value.shape) - if len(key_shape) != 4 or len(value_shape) != 4: - raise NotImplementedError( - f"cached attention expects key/value as [batch, seq, heads, dim], got {key_shape} and {value_shape}" - ) - if int(key_shape[0]) != 1: - raise NotImplementedError("cached attention currently expects batch size 1") - - cache_states: dict[str, tuple[Tensor, Tensor]] = env.setdefault("__internal_kv_cache_states", {}) # type: ignore[assignment] - layer_key = _kv_cache_layer_key(node) - if layer_key not in cache_states: - default_cache_len = max(512, int(key_shape[1])) - requested_cache_len = node.meta.get("max_cache_seq_len", ir.meta.get("max_cache_seq_len", default_cache_len)) - max_cache_seq_len = max(default_cache_len, int(requested_cache_len or default_cache_len)) - sink_size = int(ir.meta.get("cache_sink_size", 4) or 4) - window_size = int(node.attrs.get("window_size", 0) or 0) - cache_num_slots = int(ir.meta.get("cache_num_slots", 1) or 1) - k_cache = g.kv_cache_state( - max_cache_seq_len, - int(key_shape[2]), - int(key_shape[3]), - window_size=window_size, - sink_size=sink_size, - num_slots=cache_num_slots, - ) - v_cache = g.kv_cache_state( - max_cache_seq_len, - int(value_shape[2]), - int(value_shape[3]), - window_size=window_size, - sink_size=sink_size, - num_slots=cache_num_slots, - ) - cache_states[layer_key] = (k_cache, v_cache) - cache_entries: list[tuple[str, Tensor, Tensor]] = env.setdefault("__internal_kv_cache_state_entries", []) # type: ignore[assignment] - cache_entries.append((layer_key, k_cache, v_cache)) - - k_cache, v_cache = cache_states[layer_key] - window_size = int(node.attrs.get("window_size", 0) or 0) - sink_size = int(ir.meta.get("cache_sink_size", 4) or 4) - g.kv_cache_append(key, k_cache, window_size=window_size, sink_size=sink_size) - g.kv_cache_append(value, v_cache, window_size=window_size, sink_size=sink_size) - kv_input_cache_states: dict[tuple[str, str], tuple[Tensor, Tensor, int, int]] = env.setdefault( - "__internal_kv_cache_by_kv_inputs", - {}, - ) # type: ignore[assignment] - cache_entry = (k_cache, v_cache, window_size, int(value_shape[3])) - kv_input_cache_states[(node.inputs[1], node.inputs[2])] = cache_entry - kv_input_cache_states[_canonical_kv_cache_input_key(ir, node)] = cache_entry - out = g.attention_cached( - query, - key, - value, - k_cache, - v_cache, - scale=_resolve_attention_scale(node, query), - position_offset=_DYNAMIC_KV_POSITION_OFFSET, - window_size=window_size, - v_head_dim=int(value_shape[3]), - ) - if output_layout == "bthd": - return [out] - return [g.permute(out, (0, 2, 1, 3))] - - -def _lower_attention_with_shared_internal_kv_cache( - g: Graph, - ir: IRGraph, - env: dict[str, Any], - node: IRNode, - *, - query: Tensor, - key: Tensor, - value: Tensor, - output_layout: str, -) -> list[Tensor] | None: - kv_input_cache_states: dict[tuple[str, str], tuple[Tensor, Tensor, int, int]] = env.get( - "__internal_kv_cache_by_kv_inputs", - {}, - ) - cache_state = kv_input_cache_states.get((node.inputs[1], node.inputs[2])) - if cache_state is None: - cache_state = kv_input_cache_states.get(_canonical_kv_cache_input_key(ir, node)) - if cache_state is None: - return None - k_cache, v_cache, window_size, v_head_dim = cache_state - out = g.attention_cached( - query, - key, - value, - k_cache, - v_cache, - scale=_resolve_attention_scale(node, query), - position_offset=_DYNAMIC_KV_QUERY_POSITION_OFFSET, - window_size=window_size, - v_head_dim=v_head_dim, - ) - if output_layout == "bthd": - return [out] - return [g.permute(out, (0, 2, 1, 3))] - - -def _canonical_kv_cache_input_key(ir: IRGraph, node: IRNode) -> tuple[str, str]: - return ( - _canonical_kv_cache_value_id(ir, node.inputs[1]), - _canonical_kv_cache_value_id(ir, node.inputs[2]), - ) - - -def _canonical_kv_cache_value_id(ir: IRGraph, value_id: str) -> str: - current = value_id - seen: set[str] = set() - while current not in seen: - seen.add(current) - value = ir.values.get(current) - producer_id = None if value is None else value.producer - if producer_id is None: - break - producer = ir.nodes.get(producer_id) - if producer is None or producer.op not in _KV_CACHE_TRANSPARENT_OPS or len(producer.inputs) != 1: - break - current = producer.inputs[0] - return current - - -def _should_lower_conv1d_with_internal_conv_cache(ir: IRGraph, node: IRNode, x: Tensor, weight: Tensor) -> bool: - if not bool(ir.meta.get("use_internal_conv_cache", False)): - return False - if node.op != "conv1d": - return False - component = str(ir.meta.get("component", "") or "").strip().lower() - if component not in {"decoder_step", "decoder_prefill_chunk", "decoder_media_step", - "decoder_embed_chunk"}: - return False - stride = int(node.attrs.get("stride", 1)) - padding = int(node.attrs.get("padding", 0)) - dilation = int(node.attrs.get("dilation", 1)) - groups = int(node.attrs.get("groups", 1)) - return ( - len(x.shape) == 3 - and len(weight.shape) == 3 - and int(x.shape[0]) == 1 - and groups == int(x.shape[1]) == int(weight.shape[0]) - and int(weight.shape[1]) == 1 - and stride == 1 - and dilation == 1 - and padding == max(int(weight.shape[2]) - 1, 0) - ) - - -def _conv_cache_layer_key(ir: IRGraph, node: IRNode) -> str: - if len(node.inputs) > 1: - value = ir.values.get(node.inputs[1]) - if value is not None and isinstance(value.meta, dict): - source_name = value.meta.get("source_name") - if isinstance(source_name, str) and source_name: - return source_name - return str(node.id) - - -def _lower_conv1d_with_internal_conv_cache( - g: Graph, - ir: IRGraph, - env: dict[str, Any], - node: IRNode, - *, - x: Tensor, - weight: Tensor, - bias: Tensor | None, -) -> list[Tensor]: - channels = int(x.shape[1]) - seq_len = int(x.shape[2]) - kernel_size = int(weight.shape[2]) - x_nlc = g.permute(x, (0, 2, 1)) - - layer_key = f"conv:{_conv_cache_layer_key(ir, node)}" - conv_states: dict[str, Tensor] = env.setdefault("__internal_conv_cache_states", {}) # type: ignore[assignment] - if layer_key not in conv_states: - cache_state = g.conv_cache_state(kernel_size, channels) - conv_states[layer_key] = cache_state - cache_entries: list[tuple[str, Tensor, Tensor]] = env.setdefault("__internal_kv_cache_state_entries", []) # type: ignore[assignment] - # Reuse the existing manifest/cache-transfer shape. Conv state is a - # single mutable tensor, so key/value intentionally point to the same node. - cache_entries.append((layer_key, cache_state, cache_state)) - cache_state = conv_states[layer_key] - - if seq_len > 1: - # Prefill chunks still compute the chunk output with the regular causal - # convolution; the cache append is the side-effect needed by later step - # graphs. The cache only needs the last K rows for the next token. - x_rows = g.reshape(x_nlc, (seq_len, channels)) - g.conv_cache_append(x_rows, cache_state) - out_nlc = g.conv1d_causal(x_nlc, weight, kernel_size=kernel_size, dilation=1) - else: - x_rows = g.reshape(x_nlc, (1, channels)) - window = g.conv_cache_append(x_rows, cache_state) - window_nlc = g.reshape(window, (1, kernel_size, channels)) - out_nlc = g.conv1d_causal(window_nlc, weight, kernel_size=kernel_size, dilation=1) - out_nlc = g.slice(out_nlc, axis=1, start=kernel_size - 1, length=1) - - if bias is not None: - bias_reshaped = g.reshape(bias, (1, 1, int(weight.shape[0]))) - out_nlc, bias_reshaped = _legalize_elementwise_binary_inputs(g, out_nlc, bias_reshaped) - out_nlc = g.add(out_nlc, bias_reshaped) - return [g.permute(out_nlc, (0, 2, 1))] - - -def _should_lower_gated_deltanet_with_internal_cache( - ir: IRGraph, - op: str, - *, - seq_len: int, -) -> bool: - if not bool(ir.meta.get("use_internal_gated_deltanet_cache", False)): - return False - if op not in {"gated_deltanet_prefill", "gated_deltanet_decode"}: - return False - component = str(ir.meta.get("component", "") or "").strip().lower() - return component in {"decoder_prefill_chunk", "decoder_step", "decoder_media_step"} - - -def _gated_deltanet_layer_key(ir: IRGraph, node: IRNode) -> str: - if len(node.inputs) > 1: - value = ir.values.get(node.inputs[1]) - if value is not None and isinstance(value.meta, dict): - source_name = value.meta.get("source_name") - if isinstance(source_name, str) and source_name: - return source_name - return str(node.id) - - -def _lower_gated_deltanet_initial_state( - g: Graph, - ir: IRGraph, - env: dict[str, Any], - node: IRNode, - *, - batch_size: int, - key_dim: int, - num_v_heads: int, - value_dim: int, - seq_len: int, - op: str, -) -> tuple[str | None, Tensor]: - state_shape = (batch_size, key_dim, num_v_heads, value_dim) - if not _should_lower_gated_deltanet_with_internal_cache(ir, op, seq_len=seq_len): - return None, _materialize_constant_tensor(g, torch.zeros(state_shape, dtype=torch.float16)) - - layer_key = f"gdn:{_gated_deltanet_layer_key(ir, node)}" - gdn_states: dict[str, Tensor] = env.setdefault("__internal_gated_deltanet_cache_states", {}) # type: ignore[assignment] - if layer_key not in gdn_states: - gdn_states[layer_key] = g.recurrent_cache_state(state_shape, dtype=g.FP16) - return layer_key, gdn_states[layer_key] - - -def _lower_gated_deltanet_conv1d( - g: Graph, - ir: IRGraph, - env: dict[str, Any], - node: IRNode, - *, - mixed_qkv: Tensor, - conv_weight: Tensor, - batch_size: int, - seq_len: int, - mixed_qkv_dim: int, - op: str, -) -> Tensor: - kernel_size = int(conv_weight.shape[2]) - if not _should_lower_gated_deltanet_with_internal_cache(ir, op, seq_len=seq_len): - return g.conv1d_causal(mixed_qkv, conv_weight, kernel_size=kernel_size, dilation=1) - - layer_key = f"conv:gdn:{_gated_deltanet_layer_key(ir, node)}" - conv_states: dict[str, Tensor] = env.setdefault("__internal_conv_cache_states", {}) # type: ignore[assignment] - if layer_key not in conv_states: - cache_state = g.conv_cache_state(kernel_size, mixed_qkv_dim) - conv_states[layer_key] = cache_state - cache_entries: list[tuple[str, Tensor, Tensor]] = env.setdefault("__internal_kv_cache_state_entries", []) # type: ignore[assignment] - cache_entries.append((layer_key, cache_state, cache_state)) - - cache_state = conv_states[layer_key] - if batch_size != 1: - raise NotImplementedError( - f"gated_deltanet conv cache lowering currently supports batch_size=1 only, got {batch_size}" - ) - x_rows = g.reshape(mixed_qkv, (batch_size * seq_len, mixed_qkv_dim)) - - if seq_len == 1: - window = g.conv_cache_append(x_rows, cache_state) - window_nlc = g.reshape(window, (batch_size, kernel_size, mixed_qkv_dim)) - conv_out = g.conv1d_causal(window_nlc, conv_weight, kernel_size=kernel_size, dilation=1) - return g.slice(conv_out, axis=1, start=kernel_size - 1, length=1) - - conv_out = g.conv1d_causal(mixed_qkv, conv_weight, kernel_size=kernel_size, dilation=1) - g.conv_cache_initialize(x_rows, cache_state) - return conv_out - - -def _lower_constant_value( - g: Graph, - value: IRValue, - const: Any, - *, - binding: WeightBinding | None = None, -) -> Any: - if binding is None: - binding = _lookup_weight_binding(value) - if binding is not None: - _debug_mmap_binding(value.id, binding) - return g.mmap_weights(binding.path) - - _debug_constant_fallback(value, const) - - if isinstance(const, torch.nn.Parameter): - const = const.detach() - if isinstance(const, torch.Tensor): - tensor_value = const.detach().cpu() - else: - raise NotImplementedError( - f"unsupported IR constant type for {value.id}: {type(const).__name__}" - ) - - if tensor_value.numel() == 1 and value.shape in {None, ()}: - return tensor_value.item() - - return _materialize_constant_tensor(g, tensor_value) - - -def _lookup_weight_binding(value: IRValue) -> WeightBinding | None: - meta = getattr(value, "meta", None) - if isinstance(meta, dict): - path = meta.get("path") - kind = meta.get("kind") - source_name = meta.get("source_name") - if isinstance(path, str) and isinstance(kind, str) and isinstance(source_name, str): - return WeightBinding(path=path, kind=kind, source_name=source_name) - return None - - -def _debug_mmap_binding(value_id: str, binding: WeightBinding) -> None: - if os.environ.get("CACTUS_TRANSPILER_DEBUG_MMAP") != "1": - return - print( - "[transpile:mmap] " - f"value={value_id} " - f"kind={binding.kind} " - f"source={binding.source_name} " - f"path={binding.path}" - ) - - -def _debug_constant_fallback(value: IRValue, const: Any) -> None: - if os.environ.get("CACTUS_TRANSPILER_DEBUG_MMAP") != "1": - return - const_type = type(const).__name__ - shape = None - dtype = None - if isinstance(const, torch.Tensor): - shape = tuple(const.shape) - dtype = str(const.dtype) - source_name = None - if isinstance(value.meta, dict): - source_name = value.meta.get("source_name") - print( - "[transpile:fallback] " - f"value={value.id} " - f"source={source_name} " - f"const_type={const_type} " - f"shape={shape} " - f"dtype={dtype}" - ) - - -def _debug_embedding_lowering(node: IRNode, embedding_tensor: Tensor, indices_tensor: Tensor) -> None: - if os.environ.get("CACTUS_TRANSPILER_DEBUG_MMAP") != "1": - return - print( - "[transpile:embedding] " - f"node={node.id} " - f"embedding_id={embedding_tensor.id} " - f"embedding_shape={tuple(embedding_tensor.shape)} " - f"embedding_dtype={embedding_tensor.dtype} " - f"indices_id={indices_tensor.id} " - f"indices_shape={tuple(indices_tensor.shape)} " - f"indices_dtype={indices_tensor.dtype}" - ) - - -def _matmul_with_quantized_rhs_legalization( - g: Graph, - lhs: Tensor, - rhs: Tensor, - *, - pretransposed_rhs: bool = False, - output_dtype: int | None = None, -) -> Tensor: - if _is_quantized_runtime_dtype(rhs.dtype) and lhs.dtype == Graph.FP32: - lhs = g.precision_cast(lhs, Graph.FP16) - return g.matmul(lhs, rhs, pretransposed_rhs=pretransposed_rhs, output_dtype=output_dtype) - - -def _matmul_output_dtype( - lhs: Tensor, - rhs: Tensor, - *, - output_dtype: int | None, -) -> int | None: - if output_dtype != Graph.FP16: - return output_dtype - if _is_quantized_runtime_dtype(rhs.dtype): - return output_dtype - if lhs.dtype == Graph.FP32 or rhs.dtype == Graph.FP32: - return Graph.FP32 - return output_dtype - - -def _trim_padded_last_dim(g: Graph, tensor: Tensor, expected_shape: tuple[int, ...] | None) -> Tensor: - if expected_shape is None: - return tensor - actual_shape = tuple(int(dim) for dim in tensor.shape) - target_shape = tuple(int(dim) for dim in expected_shape) - if actual_shape == target_shape: - return tensor - if ( - len(actual_shape) == len(target_shape) - and len(actual_shape) >= 1 - and actual_shape[:-1] == target_shape[:-1] - and actual_shape[-1] > target_shape[-1] - ): - return g.slice(tensor, len(actual_shape) - 1, 0, target_shape[-1]) - return tensor - - -def _legalize_for_transpose(g: Graph, x: Tensor) -> Tensor: - return x - - -def _reshape_attention_output_for_linear( - g: Graph, - attn_out: Tensor, - target_shape: tuple[int, ...], -) -> Tensor: - actual_shape = tuple(int(dim) for dim in attn_out.shape) - target_shape = tuple(int(dim) for dim in target_shape) - if len(actual_shape) == 4 and len(target_shape) == 3 and actual_shape[0] == target_shape[0]: - batch, dim1, dim2, dim3 = actual_shape - _, target_seq, target_hidden = target_shape - if dim1 == target_seq and dim2 * dim3 == target_hidden: - return g.reshape(attn_out, target_shape) - if dim2 == target_seq and dim1 * dim3 == target_hidden: - return g.reshape(g.permute(attn_out, (0, 2, 1, 3)), target_shape) - return g.reshape(attn_out, target_shape) - - -def _lower_ir_node(g: Graph, node: IRNode, env: dict[str, Any], ir: IRGraph) -> list[Any]: - op = node.op - - if op == "arange": - start = int(node.attrs.get("start", 0)) - end = int(node.attrs["end"]) - step = int(node.attrs.get("step", 1)) - dtype = _materialize_constant_torch_dtype(node.attrs.get("dtype")) - tensor_value = torch.arange(start, end, step=step, dtype=dtype) - return [_materialize_constant_tensor(g, tensor_value)] - - if op == "add": - return [_lower_binary_op(g, env[node.inputs[0]], env[node.inputs[1]], "add")] - - if op == "add_clipped": - lhs, rhs = _legalize_elementwise_binary_inputs(g, _tensor(env, node.inputs[0]), _tensor(env, node.inputs[1])) - return [g.add_clipped(lhs, rhs)] - - if op == "subtract": - return [_lower_binary_op(g, env[node.inputs[0]], env[node.inputs[1]], "subtract")] - - if op == "multiply": - return [_lower_binary_op(g, env[node.inputs[0]], env[node.inputs[1]], "multiply")] - - if op == "multiply_inplace": - return [_lower_binary_op(g, env[node.inputs[0]], env[node.inputs[1]], "multiply")] - - if op == "divide": - return [_lower_binary_op(g, env[node.inputs[0]], env[node.inputs[1]], "divide")] - - if op == "not_equal": - return [_lower_compare_op(g, env[node.inputs[0]], env[node.inputs[1]], "not_equal")] - - if op == "equal": - return [_lower_compare_op(g, env[node.inputs[0]], env[node.inputs[1]], "equal")] - - if op == "greater": - return [_lower_compare_op(g, env[node.inputs[0]], env[node.inputs[1]], "greater")] - - if op == "greater_equal": - return [_lower_compare_op(g, env[node.inputs[0]], env[node.inputs[1]], "greater_equal")] - - if op == "less": - return [_lower_compare_op(g, env[node.inputs[0]], env[node.inputs[1]], "less")] - - if op == "less_equal": - return [_lower_compare_op(g, env[node.inputs[0]], env[node.inputs[1]], "less_equal")] - - if op == "logical_and": - lhs = _lower_compare_op(g, env[node.inputs[0]], 0.0, "not_equal") - rhs = _lower_compare_op(g, env[node.inputs[1]], 0.0, "not_equal") - return [_lower_binary_op(g, lhs, rhs, "multiply")] - - if op == "logical_or": - lhs = _lower_compare_op(g, env[node.inputs[0]], 0.0, "not_equal") - rhs = _lower_compare_op(g, env[node.inputs[1]], 0.0, "not_equal") - return [_lower_compare_op(g, _lower_binary_op(g, lhs, rhs, "add"), 0.0, "not_equal")] - - if op == "logical_not": - return [_lower_compare_op(g, env[node.inputs[0]], 0.0, "equal")] - - if op == "where": - return [_lower_where_op(g, node, env)] - - if op == "scalar_add": - x = _ensure_scalar_math_tensor(g, _tensor(env, node.inputs[0])) - return [g.scalar_add(x, float(node.attrs["value"]))] - - if op == "scalar_subtract": - x = _ensure_scalar_math_tensor(g, _tensor(env, node.inputs[0])) - return [g.scalar_subtract(x, float(node.attrs["value"]))] - - if op == "scalar_subtract_reverse": - x = _ensure_scalar_math_tensor(g, _tensor(env, node.inputs[0])) - return [g.scalar_add(g.scalar_multiply(x, -1.0), float(node.attrs["value"]))] - - if op == "scalar_multiply": - x = _ensure_scalar_math_tensor(g, _tensor(env, node.inputs[0])) - return [g.scalar_multiply(x, float(node.attrs["value"]))] - - if op == "scalar_divide": - x = _ensure_scalar_math_tensor(g, _tensor(env, node.inputs[0])) - return [g.scalar_divide(x, float(node.attrs["value"]))] - - if op == "scalar_floor_divide": - source = _tensor(env, node.inputs[0]) - x = _ensure_scalar_math_tensor(g, source) - y = g.scalar_floor_divide(x, float(node.attrs["value"])) - output_dtype = ir.values[node.outputs[0]].dtype - if output_dtype is None: - return [y] - target_dtype = _map_ir_dtype(output_dtype) - if y.dtype == target_dtype: - return [y] - return [g.precision_cast(y, target_dtype)] - - if op == "scalar_not_equal": - value = float(node.attrs["value"]) - source = _tensor(env, node.inputs[0]) - target_dtype = Graph.FP32 if source.dtype != Graph.FP16 or abs(value) > 65504.0 else Graph.FP16 - x = _ensure_tensor_dtype(g, source, target_dtype) - return [g.scalar_not_equal(x, value)] - - if op == "scalar_equal": - value = float(node.attrs["value"]) - source = _tensor(env, node.inputs[0]) - target_dtype = Graph.FP32 if source.dtype != Graph.FP16 or abs(value) > 65504.0 else Graph.FP16 - x = _ensure_tensor_dtype(g, source, target_dtype) - not_equal = g.scalar_not_equal(x, value) - return [g.scalar_not_equal(not_equal, 1.0)] - - if op == "scalar_greater": - x = _ensure_scalar_math_tensor(g, _tensor(env, node.inputs[0])) - return [_lower_compare_op(g, x, float(node.attrs["value"]), "greater")] - - if op == "scalar_greater_equal": - x = _ensure_scalar_math_tensor(g, _tensor(env, node.inputs[0])) - return [_lower_compare_op(g, x, float(node.attrs["value"]), "greater_equal")] - - if op == "scalar_less": - x = _ensure_scalar_math_tensor(g, _tensor(env, node.inputs[0])) - return [_lower_compare_op(g, x, float(node.attrs["value"]), "less")] - - if op == "scalar_less_equal": - x = _ensure_scalar_math_tensor(g, _tensor(env, node.inputs[0])) - return [_lower_compare_op(g, x, float(node.attrs["value"]), "less_equal")] - - if op == "scalar_divide_reverse": - raise NotImplementedError("scalar_divide_reverse is not directly supported by Cactus graph ops") - - if op == "precision_cast": - target_dtype = node.attrs.get("dtype") - if target_dtype is None: - return [_tensor(env, node.inputs[0])] - return [g.precision_cast(_tensor(env, node.inputs[0]), _map_ir_or_torch_dtype(target_dtype))] - - if op == "type_as": - source = _tensor(env, node.inputs[0]) - target = _tensor(env, node.inputs[1]) - return [g.precision_cast(source, target.dtype)] - - if op == "abs": - return [g.abs(_tensor(env, node.inputs[0]))] - - if op == "clamp": - x = _ensure_fp16_tensor(g, _tensor(env, node.inputs[0])) - min_value = node.attrs.get("min") - max_value = node.attrs.get("max") - input_index = 1 - if min_value is None and ( - bool(node.attrs.get("has_min_tensor", False)) or len(node.inputs) > input_index - ): - min_value = _constant_scalar_from_ir(ir, node.inputs[input_index]) - input_index += 1 - if max_value is None and ( - bool(node.attrs.get("has_max_tensor", False)) or len(node.inputs) > input_index - ): - max_value = _constant_scalar_from_ir(ir, node.inputs[input_index]) - lo = float(min_value) if min_value is not None and math.isfinite(float(min_value)) else -65504.0 - hi = float(max_value) if max_value is not None and math.isfinite(float(max_value)) else 65504.0 - return [g.clamp(x, lo, hi)] - - if op == "negate": - return [g.scalar_multiply(_tensor(env, node.inputs[0]), -1.0)] - - if op == "pow": - exponent = node.attrs.get("exponent") - if exponent is None: - raise NotImplementedError("tensor-tensor pow is not supported by Cactus graph ops") - return [g.pow(_tensor(env, node.inputs[0]), float(exponent))] - - if op == "scalar_exp": - return [g.scalar_exp(_tensor(env, node.inputs[0]))] - - if op == "scalar_sqrt": - return [g.scalar_sqrt(_tensor(env, node.inputs[0]))] - - if op == "scalar_log": - return [g.scalar_log(_tensor(env, node.inputs[0]))] - - if op == "scalar_cos": - return [g.scalar_cos(_tensor(env, node.inputs[0]))] - - if op == "scalar_sin": - return [g.scalar_sin(_tensor(env, node.inputs[0]))] - - if op in {"reshape", "view"}: - source = env[node.inputs[0]] - if isinstance(source, BroadcastAlias): - input_shape = tuple(source.logical_shape) - else: - input_shape = tuple(_tensor(env, node.inputs[0]).shape) - target_shape = _resolve_reshape_shape(input_shape, tuple(node.attrs["shape"])) - if isinstance(source, BroadcastAlias): - if source.kind != "gqa_repeat_kv": - raise NotImplementedError(f"unsupported broadcast alias kind in reshape: {source.kind}") - base_shape = tuple(source.tensor.shape) - if len(base_shape) != 4: - raise NotImplementedError(f"gqa_repeat_kv alias requires 4D base tensor, got {base_shape}") - if len(target_shape) != 4: - return [g.reshape(_materialize_broadcast_alias(g, source), target_shape)] - batch, kv_heads, seq_len, head_dim = base_shape - if target_shape[0] != batch or target_shape[2] != seq_len or target_shape[3] != head_dim: - raise NotImplementedError( - f"gqa_repeat_kv reshape mismatch: base {base_shape}, target {target_shape}" - ) - if target_shape[1] % max(kv_heads, 1) != 0: - raise NotImplementedError( - f"gqa_repeat_kv target head count must be a multiple of kv heads: {base_shape} -> {target_shape}" - ) - return [BroadcastAlias(source.tensor, target_shape, source.kind)] - return [g.reshape(_tensor(env, node.inputs[0]), target_shape)] - - if op == "flatten": - return [ - g.flatten( - _tensor(env, node.inputs[0]), - start_dim=int(node.attrs["start_dim"]), - end_dim=int(node.attrs["end_dim"]), - ) - ] - - if op == "unsqueeze": - x = _tensor(env, node.inputs[0]) - target_shape = node.attrs.get("shape") - if target_shape is None: - dim = _normalize_dim(int(node.attrs["dim"]), len(x.shape) + 1) - shape_list = list(x.shape) - shape_list.insert(dim, 1) - target_shape = tuple(shape_list) - return [g.reshape(x, tuple(target_shape))] - - if op == "expand": - matched_base = _match_gqa_expand_alias(ir, node) - if matched_base is not None: - return [BroadcastAlias(_tensor(env, matched_base), tuple(node.attrs["shape"]), "gqa_repeat_kv")] - x = _tensor(env, node.inputs[0]) - target_shape = _resolve_expand_shape(tuple(x.shape), tuple(node.attrs["shape"])) - if target_shape == tuple(x.shape): - return [x] - return [g.expand(x, target_shape)] - - if op == "repeat": - x = _tensor(env, node.inputs[0]) - repeats = tuple(int(value) for value in node.attrs["repeats"]) - return [_lower_repeat(g, x, repeats)] - - if op == "one_hot": - indices = _tensor(env, node.inputs[0]) - num_classes = int(node.attrs["num_classes"]) - eye = torch.eye(num_classes, dtype=torch.float16) - embedding = _materialize_constant_tensor(g, eye) - return [g.embedding_from_tensor(embedding, indices)] - - if op == "transpose": - x = _legalize_for_transpose(g, _tensor(env, node.inputs[0])) - dim0 = _normalize_dim(int(node.attrs["dim0"]), len(x.shape)) - dim1 = _normalize_dim(int(node.attrs["dim1"]), len(x.shape)) - rank = len(x.shape) - permutation = list(range(rank)) - permutation[dim0], permutation[dim1] = permutation[dim1], permutation[dim0] - if rank == 2 and permutation == [1, 0]: - return [g.transpose(x)] - return [g.permute(x, permutation)] - - if op == "permute": - source = env[node.inputs[0]] - if isinstance(source, BroadcastAlias): - x = _materialize_broadcast_alias(g, source) - else: - x = _tensor(env, node.inputs[0]) - x = _legalize_for_transpose(g, x) - permutation = tuple(_normalize_dim(int(dim), len(x.shape)) for dim in node.attrs["permutation"]) - if len(permutation) == 2 and permutation == (1, 0): - return [g.transpose(x)] - return [g.permute(x, permutation)] - - if op == "matmul": - lhs = _tensor_or_materialized_alias(g, env, node.inputs[0]) - rhs = _tensor_or_materialized_alias(g, env, node.inputs[1]) - output_value = ir.values.get(node.outputs[0]) - output_dtype = ( - _map_ir_dtype(output_value.dtype) - if output_value is not None and output_value.dtype is not None - else None - ) - matmul_output_dtype = _matmul_output_dtype(lhs, rhs, output_dtype=output_dtype) - if len(lhs.shape) == 2 and len(rhs.shape) == 2: - out = _matmul_with_quantized_rhs_legalization(g, lhs, rhs, output_dtype=matmul_output_dtype) - if output_dtype is not None and out.dtype != output_dtype: - out = g.precision_cast(out, output_dtype) - return [out] - legalized = _legalize_matmul_inputs(g, lhs, rhs, node) - if legalized is None: - static_batched = _lower_static_batched_matmul(g, lhs, rhs, output_dtype=matmul_output_dtype) - if static_batched is not None: - if output_dtype is not None and static_batched.dtype != output_dtype: - static_batched = g.precision_cast(static_batched, output_dtype) - return [static_batched] - out = _matmul_with_quantized_rhs_legalization(g, lhs, rhs, output_dtype=matmul_output_dtype) - if output_dtype is not None and out.dtype != output_dtype: - out = g.precision_cast(out, output_dtype) - return [out] - lhs_2d, rhs_2d, output_shape = legalized - out = _matmul_with_quantized_rhs_legalization(g, lhs_2d, rhs_2d, output_dtype=matmul_output_dtype) - out = g.reshape(out, output_shape) - if output_dtype is not None and out.dtype != output_dtype: - out = g.precision_cast(out, output_dtype) - return [out] - - if op == "linear": - x = _tensor(env, node.inputs[0]) - weight = _tensor(env, node.inputs[1]) - output_value = ir.values.get(node.outputs[0]) - output_dtype = _map_ir_dtype(output_value.dtype) if output_value is not None and output_value.dtype is not None else x.dtype - reshape_back: tuple[int, ...] | None = None - if len(x.shape) > 2: - x = _flatten_to_2d_for_linear(g, x) - output_shape = output_value.shape if output_value is not None else None - if output_shape is None: - output_shape = node.meta.get("shape") - if not isinstance(output_shape, tuple): - raise NotImplementedError(f"linear missing output shape metadata for node {node.id}") - reshape_back = tuple(int(v) for v in output_shape) - - matmul_output_dtype = _matmul_output_dtype(x, weight, output_dtype=output_dtype) - if node.attrs.get("has_bias"): - bias = _tensor(env, node.inputs[2]) - if ( - x.dtype == Graph.FP16 - and weight.dtype == Graph.FP16 - and output_dtype == Graph.FP16 - ): - matmul_output_dtype = Graph.FP32 - elif bias.dtype == Graph.FP32 and not _is_quantized_runtime_dtype(weight.dtype): - matmul_output_dtype = Graph.FP32 - out = _matmul_with_quantized_rhs_legalization( - g, - x, - weight, - pretransposed_rhs=True, - output_dtype=matmul_output_dtype, - ) - out = _trim_padded_last_dim(g, out, output_value.shape if output_value is not None else None) - if node.attrs.get("has_bias"): - bias = _ensure_tensor_dtype(g, bias, matmul_output_dtype) - if output_value is not None and output_value.shape is not None: - expected_last_dim = int(output_value.shape[-1]) - if len(bias.shape) >= 1 and int(bias.shape[-1]) > expected_last_dim: - bias = g.slice(bias, len(bias.shape) - 1, 0, expected_last_dim) - out, bias = _legalize_elementwise_binary_inputs(g, out, bias) - try: - out = g.add(out, bias) - except RuntimeError as exc: - raise RuntimeError( - "linear bias add failed while lowering " - f"{node.id}: out_shape={tuple(out.shape)} bias_shape={tuple(bias.shape)} " - f"expected_shape={None if output_value is None else output_value.shape}" - ) from exc - if reshape_back is not None: - out = g.reshape(out, reshape_back) - if out.dtype != output_dtype: - out = g.precision_cast(out, output_dtype) - return [out] - - if op == "dense_mlp_tq_fused": - hidden = _tensor(env, node.inputs[0]) - gate_weight = _tensor(env, node.inputs[1]) - up_weight = _tensor(env, node.inputs[2]) - down_weight = _tensor(env, node.inputs[3]) - product_scale = ( - node.attrs.get("product_scale") - if "product_scale" in node.attrs - else node.meta.get("product_scale_from_export", node.meta.get("product_scale_elided_by_post_norm", 1.0)) - ) - out = g.dense_mlp_tq_fused( - hidden, - gate_weight, - up_weight, - down_weight, - product_scale=float(product_scale), - ) - output_value = ir.values.get(node.outputs[0]) - output_dtype = _map_ir_dtype(output_value.dtype) if output_value is not None and output_value.dtype is not None else out.dtype - if out.dtype != output_dtype: - out = g.precision_cast(out, output_dtype) - return [out] - - if op == "addmm": - bias = _tensor(env, node.inputs[0]) - lhs = _tensor(env, node.inputs[1]) - rhs = _tensor(env, node.inputs[2]) - output_value = ir.values.get(node.outputs[0]) - output_dtype = _map_ir_dtype(output_value.dtype) if output_value is not None and output_value.dtype is not None else lhs.dtype - matmul_output_dtype = _matmul_output_dtype(lhs, rhs, output_dtype=output_dtype) - if lhs.dtype == Graph.FP16 and rhs.dtype == Graph.FP16 and bias.dtype == Graph.FP16 and output_dtype == Graph.FP16: - matmul_output_dtype = Graph.FP32 - elif bias.dtype == Graph.FP32 and not _is_quantized_runtime_dtype(rhs.dtype): - matmul_output_dtype = Graph.FP32 - out = _matmul_with_quantized_rhs_legalization(g, lhs, rhs, output_dtype=matmul_output_dtype) - out = _trim_padded_last_dim(g, out, output_value.shape if output_value is not None else None) - bias = _ensure_tensor_dtype(g, bias, matmul_output_dtype) - out = g.add(bias, out) - if out.dtype != output_dtype: - out = g.precision_cast(out, output_dtype) - return [out] - - if op == "relu": - return [g.relu(_tensor(env, node.inputs[0]))] - - if op == "silu": - return [g.silu(_tensor(env, node.inputs[0]))] - - if op == "gelu": - return [g.gelu(_tensor(env, node.inputs[0]))] - - if op == "gelu_erf": - return [g.gelu_erf(_tensor(env, node.inputs[0]))] - - if op == "sigmoid": - return [g.sigmoid(_tensor(env, node.inputs[0]))] - - if op == "glu": - return [g.glu(_tensor(env, node.inputs[0]), axis=int(node.attrs.get("axis", -1)))] - - if op == "softplus": - x = _tensor(env, node.inputs[0]) - return [_lower_softplus(g, x)] - - if op == "tanh": - return [g.tanh(_tensor(env, node.inputs[0]))] - - if op == "softmax": - return [g.softmax(_tensor(env, node.inputs[0]), axis=int(node.attrs.get("axis", -1)))] - - if op == "lstm_cell": - x = _tensor(env, node.inputs[0]) - h_prev = _tensor(env, node.inputs[1]) - c_prev = _tensor(env, node.inputs[2]) - weight_ih = _tensor(env, node.inputs[3]) - weight_hh = _tensor(env, node.inputs[4]) - bias_ih = _tensor(env, node.inputs[5]) - bias_hh = _tensor(env, node.inputs[6]) - - lstm_out = g.lstm_cell( - _ensure_fp16_tensor(g, x), - _ensure_fp16_tensor(g, h_prev), - _ensure_fp16_tensor(g, c_prev), - _ensure_fp16_tensor(g, weight_ih), - _ensure_fp16_tensor(g, weight_hh), - _ensure_fp16_tensor(g, bias_ih), - _ensure_fp16_tensor(g, bias_hh), - ) - h_new = g.slice(lstm_out, axis=2, start=0, length=1) - c_new = g.slice(lstm_out, axis=2, start=1, length=1) - h_shape = ir.values[node.outputs[0]].shape - c_shape = ir.values[node.outputs[1]].shape - if h_shape is not None: - h_new = g.reshape(h_new, tuple(int(dim) for dim in h_shape)) - if c_shape is not None: - c_new = g.reshape(c_new, tuple(int(dim) for dim in c_shape)) - return [h_new, c_new] - - if op in {"scaled_dot_product_attention", "attention"}: - if _should_lower_gemma4_decoder_attention_without_kernel(ir, node): - return [_lower_gemma4_decoder_attention_without_kernel(g, ir, env, node)] - mask = node.attrs.get("mask") - mask_tensor: Tensor | None = None - additive_mask = bool(node.attrs.get("additive_mask", False)) - literal_mask_is_causal = False - if len(node.inputs) > 3: - raw_mask = env.get(node.inputs[3]) - if isinstance(raw_mask, bool): - literal_mask_is_causal = bool(raw_mask) - mask_tensor = None - else: - mask_tensor = _tensor(env, node.inputs[3]) - elif mask is not None: - raise NotImplementedError(f"{op} with literal mask is not supported yet") - if op == "scaled_dot_product_attention": - dropout_p = float(node.attrs.get("dropout_p", 0.0)) - if dropout_p != 0.0: - raise NotImplementedError("scaled_dot_product_attention with dropout is not supported yet") - - # PyTorch SDPA exports tensors as [batch, heads, seq, dim], while Cactus attention - # expects [batch, seq, heads, dim]. The optimizer may strip surrounding layout - # permutes and mark the node as already-native to avoid large round-trip copies. - qkv_layout = str(node.attrs.get("qkv_layout", "bhsd") or "bhsd").lower() - q_layout = str(node.attrs.get("q_layout", qkv_layout) or qkv_layout).lower() - k_layout = str(node.attrs.get("k_layout", qkv_layout) or qkv_layout).lower() - v_layout = str(node.attrs.get("v_layout", qkv_layout) or qkv_layout).lower() - output_layout = str(node.attrs.get("output_layout", "bhsd") or "bhsd").lower() - query = _attention_tensor(env, node.inputs[0]) - key = _attention_tensor(env, node.inputs[1]) - value = _attention_tensor(env, node.inputs[2]) - if q_layout != "bthd": - query = g.permute(query, (0, 2, 1, 3)) - if _is_external_cross_kv_cache_value(ir, node.inputs[1]) and _is_external_cross_kv_cache_value(ir, node.inputs[2]): - key_value = ir.values.get(node.inputs[1]) - value_value = ir.values.get(node.inputs[2]) - if key_value is None or value_value is None or key_value.shape is None or value_value.shape is None: - raise ValueError("external cross-KV cache attention requires original key/value shapes") - key_shape = tuple(int(dim) for dim in key_value.shape) - value_shape = tuple(int(dim) for dim in value_value.shape) - if len(key_shape) != 4 or len(value_shape) != 4: - raise ValueError(f"external cross-KV cache attention expects rank-4 key/value tensors, got {key_shape} and {value_shape}") - if k_layout == "bhsd": - num_kv_heads, head_dim = key_shape[1], key_shape[3] - else: - num_kv_heads, head_dim = key_shape[2], key_shape[3] - v_head_dim = value_shape[3] - dummy_key = _materialize_constant_tensor( - g, - torch.zeros((1, 1, int(num_kv_heads), int(head_dim)), dtype=torch.float16), - ) - dummy_value = _materialize_constant_tensor( - g, - torch.zeros((1, 1, int(num_kv_heads), int(v_head_dim)), dtype=torch.float16), - ) - out = g.attention_cached( - query, - dummy_key, - dummy_value, - key, - value, - scale=_resolve_attention_scale(node, query), - position_offset=_DYNAMIC_KV_QUERY_POSITION_OFFSET, - window_size=int(node.attrs.get("window_size", 0)), - v_head_dim=int(v_head_dim), - ) - if output_layout == "bthd": - return [out] - return [g.permute(out, (0, 2, 1, 3))] - if k_layout != "bthd": - key = g.permute(key, (0, 2, 1, 3)) - if v_layout != "bthd": - value = g.permute(value, (0, 2, 1, 3)) - if mask_tensor is not None: - mask_tensor = _ensure_fp16_tensor(g, mask_tensor) - mask_tensor = _normalize_attention_mask_for_cactus(g, mask_tensor, query) - shared_cached = _lower_attention_with_shared_internal_kv_cache( - g, - ir, - env, - node, - query=query, - key=key, - value=value, - output_layout=output_layout, - ) - if shared_cached is not None: - return shared_cached - if _should_lower_attention_with_internal_kv_cache(ir, node): - return _lower_attention_with_internal_kv_cache( - g, - ir, - env, - node, - query=query, - key=key, - value=value, - output_layout=output_layout, - ) - out = g.attention( - query, - key, - value, - scale=_resolve_attention_scale(node, query), - is_causal=bool(node.attrs.get("is_causal", False)) or literal_mask_is_causal, - window_size=int(node.attrs.get("window_size", 0)), - mask=mask_tensor, - additive_mask=additive_mask, - ) - if output_layout == "bthd": - return [out] - return [g.permute(out, (0, 2, 1, 3))] - - if op == "attention_block": - has_mask = bool(node.attrs.get("has_mask", False)) - has_gate = bool(node.attrs.get("has_gate", False)) - has_bias = bool(node.attrs.get("has_bias", False)) - qkv_layout = str(node.attrs.get("qkv_layout", "bhsd") or "bhsd").lower() - input_index = 0 - query = _attention_tensor(env, node.inputs[input_index]) - if qkv_layout != "bthd": - query = g.permute(query, (0, 2, 1, 3)) - input_index += 1 - key = _attention_tensor(env, node.inputs[input_index]) - if qkv_layout != "bthd": - key = g.permute(key, (0, 2, 1, 3)) - input_index += 1 - value = _attention_tensor(env, node.inputs[input_index]) - if qkv_layout != "bthd": - value = g.permute(value, (0, 2, 1, 3)) - input_index += 1 - mask_tensor: Tensor | None = None - literal_mask_is_causal = False - if has_mask: - raw_mask = env.get(node.inputs[input_index]) - input_index += 1 - if isinstance(raw_mask, bool): - literal_mask_is_causal = bool(raw_mask) - mask_tensor = None - else: - mask_tensor = _tensor(env, node.inputs[input_index - 1]) - mask_tensor = _ensure_fp16_tensor(g, mask_tensor) - mask_tensor = _normalize_attention_mask_for_cactus(g, mask_tensor, query) - - attn_out = g.attention( - query, - key, - value, - scale=_resolve_attention_scale(node, query), - is_causal=bool(node.attrs.get("is_causal", False)) or literal_mask_is_causal, - window_size=int(node.attrs.get("window_size", 0)), - mask=mask_tensor, - additive_mask=bool(node.attrs.get("additive_mask", False)), - ) - flat_shape_attr = tuple(int(v) for v in node.attrs.get("attention_output_shape", ())) - flat_shape = _resolve_reshape_shape(tuple(attn_out.shape), flat_shape_attr) if flat_shape_attr else () - if flat_shape: - attn_out = _reshape_attention_output_for_linear(g, attn_out, flat_shape) - - if has_gate: - gate = _tensor(env, node.inputs[input_index]) - input_index += 1 - attn_out, gate = _legalize_elementwise_binary_inputs(g, attn_out, gate) - attn_out = g.multiply(attn_out, gate) - - weight = _tensor(env, node.inputs[input_index]) - input_index += 1 - reshape_back: tuple[int, ...] | None = None - linear_input = attn_out - if len(linear_input.shape) > 2: - linear_input = _flatten_to_2d_for_linear(g, linear_input) - output_value = ir.values.get(node.outputs[0]) - output_shape = output_value.shape if output_value is not None else None - if output_shape is None: - output_shape = node.meta.get("shape") - if not isinstance(output_shape, tuple): - raise NotImplementedError(f"attention_block missing output shape metadata for node {node.id}") - reshape_back = tuple(int(v) for v in output_shape) - - out = _matmul_with_quantized_rhs_legalization(g, linear_input, weight, pretransposed_rhs=True) - if has_bias: - out = g.add(out, _tensor(env, node.inputs[input_index])) - if reshape_back is not None: - out = g.reshape(out, reshape_back) - return [out] - - if op == "self_attention_block": - has_mask = bool(node.attrs.get("has_mask", False)) - has_gate = bool(node.attrs.get("has_gate", False)) - has_bias = bool(node.attrs.get("has_bias", False)) - has_query_projection_bias = bool(node.attrs.get("has_query_projection_bias", False)) - has_query_add = bool(node.attrs.get("has_query_add", False)) - has_rel_query_add = bool(node.attrs.get("has_rel_query_add", False)) - has_key_projection_bias = bool(node.attrs.get("has_key_projection_bias", False)) - has_value_projection_bias = bool(node.attrs.get("has_value_projection_bias", False)) - has_rel_pos_bias = bool(node.attrs.get("has_rel_pos_bias", False)) - has_relative_key_projection_bias = bool(node.attrs.get("has_relative_key_projection_bias", False)) - - input_index = 0 - hidden = _tensor(env, node.inputs[input_index]) - input_index += 1 - - query_weight = _tensor(env, node.inputs[input_index]) - input_index += 1 - query_projection_bias = None - if has_query_projection_bias: - query_projection_bias = _tensor(env, node.inputs[input_index]) - input_index += 1 - query_add = None - if has_query_add: - query_add = _tensor(env, node.inputs[input_index]) - input_index += 1 - rel_query_add = None - if has_rel_query_add: - rel_query_add = _tensor(env, node.inputs[input_index]) - input_index += 1 - - key_weight = _tensor(env, node.inputs[input_index]) - input_index += 1 - key_projection_bias = None - if has_key_projection_bias: - key_projection_bias = _tensor(env, node.inputs[input_index]) - input_index += 1 - - value_weight = _tensor(env, node.inputs[input_index]) - input_index += 1 - value_projection_bias = None - if has_value_projection_bias: - value_projection_bias = _tensor(env, node.inputs[input_index]) - input_index += 1 - - mask_tensor: Tensor | None = None - if has_mask: - mask_tensor = _tensor(env, node.inputs[input_index]) - input_index += 1 - - relative_key_input = None - relative_key_weight = None - relative_key_projection_bias = None - if has_rel_pos_bias: - relative_key_input = _tensor(env, node.inputs[input_index]) - input_index += 1 - relative_key_weight = _tensor(env, node.inputs[input_index]) - input_index += 1 - if has_relative_key_projection_bias: - relative_key_projection_bias = _tensor(env, node.inputs[input_index]) - input_index += 1 - - gate = None - if has_gate: - gate = _tensor(env, node.inputs[input_index]) - input_index += 1 - - output_weight = _tensor(env, node.inputs[input_index]) - input_index += 1 - output_bias = None - if has_bias: - output_bias = _tensor(env, node.inputs[input_index]) - - query_shape = tuple(int(v) for v in node.attrs.get("query_shape", ())) - key_shape = tuple(int(v) for v in node.attrs.get("key_shape", ())) - value_shape = tuple(int(v) for v in node.attrs.get("value_shape", ())) - relative_key_shape = tuple(int(v) for v in node.attrs.get("relative_key_shape", ())) - - query_base = _lower_projected_attention_tensor(g, hidden, query_weight, query_projection_bias, query_shape) - query = query_base - if query_add is not None: - query_add = _normalize_attention_add_tensor(g, query_add, query_shape) - query, query_add = _legalize_elementwise_binary_inputs(g, query, query_add) - query = g.add(query, query_add) - - key = _lower_projected_attention_tensor(g, hidden, key_weight, key_projection_bias, key_shape) - value = _lower_projected_attention_tensor(g, hidden, value_weight, value_projection_bias, value_shape) - - if has_rel_pos_bias: - if relative_key_input is None or relative_key_weight is None or not relative_key_shape: - raise NotImplementedError(f"self_attention_block missing relative position bias inputs for node {node.id}") - rel_query = query_base - if rel_query_add is not None: - rel_query_add = _normalize_attention_add_tensor(g, rel_query_add, query_shape) - rel_query, rel_query_add = _legalize_elementwise_binary_inputs(g, rel_query, rel_query_add) - rel_query = g.add(rel_query, rel_query_add) - relative_key = _lower_projected_attention_tensor( - g, - relative_key_input, - relative_key_weight, - relative_key_projection_bias, - relative_key_shape, - ) - rel_mask = g.rel_pos_bias(rel_query, relative_key, float(node.attrs.get("rel_pos_scale", 1.0))) - if mask_tensor is None: - mask_tensor = rel_mask - else: - mask_tensor, rel_mask = _legalize_elementwise_binary_inputs(g, mask_tensor, rel_mask) - mask_tensor = g.add(mask_tensor, rel_mask) - mask_tensor = _ensure_fp16_tensor(g, mask_tensor) - elif mask_tensor is not None: - mask_tensor = _ensure_fp16_tensor(g, mask_tensor) - if mask_tensor is not None: - mask_tensor = _normalize_attention_mask_for_cactus(g, mask_tensor, query) - - attn_out = g.attention( - query, - key, - value, - scale=_resolve_attention_scale(node, query), - is_causal=bool(node.attrs.get("is_causal", False)), - window_size=int(node.attrs.get("window_size", 0)), - mask=mask_tensor, - additive_mask=bool(node.attrs.get("additive_mask", False)), - ) - flat_shape_attr = tuple(int(v) for v in node.attrs.get("attention_output_shape", ())) - flat_shape = _resolve_reshape_shape(tuple(attn_out.shape), flat_shape_attr) if flat_shape_attr else () - if flat_shape: - attn_out = _reshape_attention_output_for_linear(g, attn_out, flat_shape) - - if gate is not None: - attn_out, gate = _legalize_elementwise_binary_inputs(g, attn_out, gate) - attn_out = g.multiply(attn_out, gate) - - linear_input = attn_out - reshape_back: tuple[int, ...] | None = None - if len(linear_input.shape) > 2: - linear_input = _flatten_to_2d_for_linear(g, linear_input) - output_value = ir.values.get(node.outputs[0]) - output_shape = output_value.shape if output_value is not None else None - if output_shape is None: - output_shape = node.meta.get("shape") - if not isinstance(output_shape, tuple): - raise NotImplementedError(f"self_attention_block missing output shape metadata for node {node.id}") - reshape_back = tuple(int(v) for v in output_shape) - - out = _matmul_with_quantized_rhs_legalization(g, linear_input, output_weight, pretransposed_rhs=True) - if output_bias is not None: - out = g.add(out, output_bias) - if reshape_back is not None: - out = g.reshape(out, reshape_back) - return [out] - - if op in ("sum", "mean", "variance", "min", "max"): - x = _ensure_scalar_math_tensor(g, _tensor(env, node.inputs[0])) - axes = _normalize_reduction_axes(node.attrs.get("axis"), len(x.shape)) - keepdim = bool(node.attrs.get("keepdim", False)) - out = _lower_reduction(g, op, x, axes=axes) - if keepdim: - new_shape = list(x.shape) - for reduced_axis in axes: - new_shape[reduced_axis] = 1 - out = g.reshape(out, tuple(new_shape)) - return [out] - - if op == "cumsum": - x = _ensure_scalar_math_tensor(g, _tensor(env, node.inputs[0])) - raw_axis = node.attrs.get("axis") - if not isinstance(raw_axis, int): - args = node.attrs.get("args", ()) - if isinstance(args, (list, tuple)) and len(args) > 1 and isinstance(args[1], int): - raw_axis = int(args[1]) - else: - raw_axis = -1 - axis = _normalize_dim(int(raw_axis), len(x.shape)) - return [g.cumsum(x, axis)] - - if op == "cat": - tensors = [_tensor(env, value_id) for value_id in node.inputs] - output_dtype: int | None = None - if node.outputs: - output_value = ir.values.get(node.outputs[0]) - if output_value is not None and output_value.dtype is not None: - output_dtype = _map_ir_or_torch_dtype(output_value.dtype) - return [_cat_with_legalized_dtype(g, tensors, axis=int(node.attrs.get("axis", 0)), dtype=output_dtype)] - - if op in {"stack", "aten.stack.default"}: - tensors = [_tensor(env, value_id) for value_id in node.inputs] - if not tensors: - raise NotImplementedError("stack requires at least one input tensor") - - raw_axis = node.attrs.get("axis") - if not isinstance(raw_axis, int): - args = node.attrs.get("args", ()) - if isinstance(args, (list, tuple)) and len(args) > 1 and isinstance(args[1], int): - raw_axis = int(args[1]) - else: - raw_axis = 0 - - base_shape = tuple(int(dim) for dim in tensors[0].shape) - axis = _normalize_dim(int(raw_axis), len(base_shape) + 1) - reshaped: list[Tensor] = [] - for tensor in tensors: - if tuple(int(dim) for dim in tensor.shape) != base_shape: - raise NotImplementedError( - f"stack requires equal input shapes, got {[tuple(int(dim) for dim in t.shape) for t in tensors]}" - ) - expanded_shape = list(base_shape) - expanded_shape.insert(axis, 1) - reshaped.append(g.reshape(tensor, tuple(expanded_shape))) - return [g.cat(reshaped, axis=axis)] - - if op == "slice": - source = env[node.inputs[0]] - if isinstance(source, BroadcastAlias): - x = _materialize_broadcast_alias(g, source) - else: - x = _tensor(env, node.inputs[0]) - axis = _normalize_dim(int(node.attrs["axis"]), len(x.shape)) - start = int(node.attrs["start"]) - end = int(node.attrs["end"]) - step = int(node.attrs.get("step", 1)) - dim_size = x.shape[axis] - if step == 0: - raise NotImplementedError("slice with step == 0 is invalid") - if step != 1: - indices = list(range(*slice(start, end, step).indices(dim_size))) - if axis == 0 or x.dtype == Graph.FP16: - return [_lower_static_strided_slice_via_gather(g, x, axis=axis, indices=indices)] - if not indices: - return [g.slice(x, axis=axis, start=0, length=0)] - expanded_shape = list(int(dim) for dim in x.shape) - expanded_shape[axis] = 1 - pieces: list[Tensor] = [] - for idx in indices: - piece = g.index(x, index_value=idx, axis=axis) - piece = g.reshape(piece, tuple(expanded_shape)) - pieces.append(piece) - if len(pieces) == 1: - return [pieces[0]] - return [g.cat(pieces, axis=axis)] - start = _normalize_index(start, dim_size) - end = _normalize_slice_end(end, dim_size) - length = max(0, end - start) - return [g.slice(x, axis=axis, start=start, length=length)] - - if op == "split_with_sizes": - x = _tensor(env, node.inputs[0]) - axis = _normalize_dim(int(node.attrs.get("axis", -1)), len(x.shape)) - sizes = tuple(int(v) for v in node.attrs["sizes"]) - start = 0 - outputs: list[Tensor] = [] - for size in sizes: - outputs.append(g.slice(x, axis=axis, start=start, length=int(size))) - start += int(size) - return [outputs] - - if op == "chunk": - x = _tensor(env, node.inputs[0]) - axis = _normalize_dim(int(node.attrs.get("axis", 0)), len(x.shape)) - chunks = int(node.attrs["chunks"]) - if chunks <= 0: - raise NotImplementedError(f"chunk requires a positive chunk count, got {chunks}") - dim_size = int(x.shape[axis]) - if dim_size <= 0: - return [[g.slice(x, axis=axis, start=0, length=0)]] - chunk_size = (dim_size + chunks - 1) // chunks - outputs: list[Tensor] = [] - start = 0 - while start < dim_size: - length = min(chunk_size, dim_size - start) - outputs.append(g.slice(x, axis=axis, start=start, length=length)) - start += length - return [outputs] - - if op in {"unbind", "aten.unbind.int"}: - x = _tensor(env, node.inputs[0]) - axis = _normalize_dim(int(node.attrs.get("axis", 0)), len(x.shape)) - outputs: list[Tensor] = [] - out_shape = tuple(int(dim) for index, dim in enumerate(x.shape) if index != axis) - for index in range(int(x.shape[axis])): - piece = g.slice(x, axis=axis, start=index, length=1) - outputs.append(g.reshape(piece, out_shape)) - return [outputs] - - if op == "ones": - shape = tuple(int(v) for v in node.attrs.get("shape", ())) - if not shape: - raise NotImplementedError(f"ones requires a static shape, got {shape}") - torch_dtype = _materialize_constant_torch_dtype(node.attrs.get("dtype")) - return [_materialize_constant_tensor(g, torch.ones(shape, dtype=torch_dtype))] - - if op == "tril": - x = _tensor(env, node.inputs[0]) - shape = tuple(int(dim) for dim in x.shape) - if len(shape) < 2: - raise NotImplementedError(f"tril expects rank >= 2, got shape {shape}") - diagonal = int(node.attrs.get("diagonal", 0)) - mask_dtype = _torch_dtype_for_graph_dtype(x.dtype) - mask = torch.ones(shape, dtype=mask_dtype).tril(diagonal=diagonal) - mask_tensor = _materialize_constant_tensor(g, mask) - return [_lower_binary_op(g, x, mask_tensor, "multiply")] - - if op == "unfold": - x = _tensor(env, node.inputs[0]) - rank = len(x.shape) - axis = _normalize_dim(int(node.attrs["dimension"]), rank) - size = int(node.attrs["size"]) - step = int(node.attrs["step"]) - if size <= 0: - raise NotImplementedError(f"unfold requires a positive size, got {size}") - if step <= 0: - raise NotImplementedError(f"unfold requires a positive step, got {step}") - dim_size = int(x.shape[axis]) - if size > dim_size: - raise NotImplementedError( - f"unfold size {size} exceeds dimension size {dim_size} for axis {axis}" - ) - - window_starts = list(range(0, dim_size - size + 1, step)) - if not window_starts: - raise NotImplementedError( - f"unfold produced no windows for axis={axis}, size={size}, step={step}, dim_size={dim_size}" - ) - - base_shape = [int(dim) for dim in x.shape] - trailing_shape = [base_shape[idx] for idx in range(rank) if idx != axis] + [size] - expanded_shape = tuple(trailing_shape[:axis] + [1] + trailing_shape[axis:]) - permutation = tuple(idx for idx in range(rank) if idx != axis) + (axis,) - - windows: list[Tensor] = [] - for start in window_starts: - piece = g.slice(x, axis=axis, start=start, length=size) - if axis != rank - 1: - piece = g.permute(piece, permutation=permutation) - piece = g.reshape(piece, expanded_shape) - windows.append(piece) - if len(windows) == 1: - return [windows[0]] - return [g.cat(windows, axis=axis)] - - if op == "pad": - x = _tensor(env, node.inputs[0]) - mode = str(node.attrs.get("mode", "constant")) - if mode != "constant": - raise NotImplementedError(f"pad mode is unsupported: {mode}") - pads = tuple(int(v) for v in node.attrs.get("pads", ())) - if len(pads) % 2 != 0: - raise NotImplementedError(f"pad expects an even-length pads tuple, got {pads}") - value = float(node.attrs.get("value", 0.0)) - - current = x - current_shape = list(int(dim) for dim in x.shape) - pad_dims = len(pads) // 2 - if pad_dims > len(current_shape): - raise NotImplementedError( - f"pad rank mismatch: pads={pads} for input shape {tuple(current_shape)}" - ) - - for pad_index in range(pad_dims): - before = pads[2 * pad_index] - after = pads[2 * pad_index + 1] - axis = len(current_shape) - 1 - pad_index - if before < 0 or after < 0: - raise NotImplementedError(f"negative pad is unsupported: {pads}") - pieces = [] - torch_dtype = _torch_dtype_for_graph_dtype(current.dtype) - if before > 0: - left_shape = list(current_shape) - left_shape[axis] = before - left = _materialize_constant_tensor( - g, - torch.full(tuple(left_shape), value, dtype=torch_dtype), - ) - pieces.append(left) - pieces.append(current) - if after > 0: - right_shape = list(current_shape) - right_shape[axis] = after - right = _materialize_constant_tensor( - g, - torch.full(tuple(right_shape), value, dtype=torch_dtype), - ) - pieces.append(right) - if len(pieces) > 1: - current = g.cat(pieces, axis=axis) - current_shape[axis] += before + after - return [current] - - if op == "index": - x = _tensor(env, node.inputs[0]) - axis = _normalize_dim(int(node.attrs.get("axis", 0)), len(x.shape)) - index_value = _normalize_index(int(node.attrs["index_value"]), x.shape[axis]) - return [g.index(x, index_value=index_value, axis=axis)] - - if op == "gather": - return [ - g.gather( - _tensor(env, node.inputs[0]), - _tensor(env, node.inputs[1]), - axis=int(node.attrs.get("axis", 0)), - ) - ] - - if op == "embedding": - embedding_tensor = _tensor(env, node.inputs[0]) - indices_tensor = _tensor(env, node.inputs[1]) - _debug_embedding_lowering(node, embedding_tensor, indices_tensor) - return [g.embedding_from_tensor(embedding_tensor, indices_tensor)] - - if op == "conv_module": - input_index = 0 - x_nlc = _tensor(env, node.inputs[input_index]) - input_index += 1 - - pointwise1_weight = _tensor(env, node.inputs[input_index]) - input_index += 1 - pointwise1_bias = None - if bool(node.attrs.get("has_pointwise1_bias", False)): - pointwise1_bias = _tensor(env, node.inputs[input_index]) - input_index += 1 - - depthwise_weight = _tensor(env, node.inputs[input_index]) - input_index += 1 - depthwise_bias = None - if bool(node.attrs.get("has_depthwise_bias", False)): - depthwise_bias = _tensor(env, node.inputs[input_index]) - input_index += 1 - - batch_norm_weight = _tensor(env, node.inputs[input_index]) - batch_norm_bias = _tensor(env, node.inputs[input_index + 1]) - batch_norm_running_mean = _tensor(env, node.inputs[input_index + 2]) - batch_norm_running_var = _tensor(env, node.inputs[input_index + 3]) - input_index += 4 - - pointwise2_weight = _tensor(env, node.inputs[input_index]) - input_index += 1 - pointwise2_bias = None - if bool(node.attrs.get("has_pointwise2_bias", False)): - pointwise2_bias = _tensor(env, node.inputs[input_index]) - - if len(x_nlc.shape) != 3: - raise NotImplementedError(f"conv_module expects rank-3 NLC input, got {x_nlc.shape}") - kernel_size = int(node.attrs.get("depthwise_kernel_size", 0)) - padding = int(node.attrs.get("depthwise_padding", 0)) - if len(depthwise_weight.shape) != 3 or kernel_size != 9 or padding != 4: - raise NotImplementedError( - f"conv_module currently supports same depthwise kernel-9 lowering, got " - f"kernel_size={kernel_size} padding={padding} weight_shape={depthwise_weight.shape}" - ) - - current = g.conv1d_pointwise(x_nlc, pointwise1_weight, bias=pointwise1_bias) - current = g.glu(current, axis=-1) - current = g.conv1d_same_depthwise_k9(current, depthwise_weight, bias=depthwise_bias) - current = g.batch_norm( - current, - batch_norm_weight, - batch_norm_bias, - batch_norm_running_mean, - batch_norm_running_var, - axis=len(current.shape) - 1, - eps=float(node.attrs.get("eps", 1e-5)), - ) - current = g.silu(current) - current = g.conv1d_pointwise(current, pointwise2_weight, bias=pointwise2_bias) - return [current] - - if op == "conv1d": - x = _tensor(env, node.inputs[0]) - weight = _tensor(env, node.inputs[1]) - bias = _tensor(env, node.inputs[2]) if len(node.inputs) > 2 else None - stride = int(node.attrs.get("stride", 1)) - padding = int(node.attrs.get("padding", 0)) - dilation = int(node.attrs.get("dilation", 1)) - groups = int(node.attrs.get("groups", 1)) - - if _should_lower_conv1d_with_internal_conv_cache(ir, node, x, weight): - return _lower_conv1d_with_internal_conv_cache( - g, - ir, - env, - node, - x=x, - weight=weight, - bias=bias, - ) - - if ( - len(x.shape) == 3 - and len(weight.shape) == 3 - and weight.shape[2] == 1 - and stride == 1 - and padding == 0 - and dilation == 1 - and groups == 1 - ): - x_nlc = g.permute(x, (0, 2, 1)) - out_nlc = g.conv1d_pointwise(x_nlc, weight, bias=bias) - return [g.permute(out_nlc, (0, 2, 1))] - - if ( - len(x.shape) == 3 - and len(weight.shape) == 3 - and groups == x.shape[1] == weight.shape[0] - and weight.shape[1] == 1 - and weight.shape[2] == 9 - and stride == 1 - and padding == 4 - and dilation == 1 - ): - x_nlc = g.permute(x, (0, 2, 1)) - out_nlc = g.conv1d_same_depthwise_k9(x_nlc, weight, bias=bias) - return [g.permute(out_nlc, (0, 2, 1))] - - if ( - len(x.shape) == 3 - and len(weight.shape) == 3 - and groups == x.shape[1] == weight.shape[0] - and weight.shape[1] == 1 - and stride == 1 - and padding == dilation * max(weight.shape[2] - 1, 0) - ): - x_nlc = g.permute(x, (0, 2, 1)) - out_nlc = g.conv1d_causal(x_nlc, weight, kernel_size=weight.shape[2], dilation=dilation) - if bias is not None: - bias_reshaped = g.reshape(bias, (1, 1, int(weight.shape[0]))) - out_nlc, bias_reshaped = _legalize_elementwise_binary_inputs(g, out_nlc, bias_reshaped) - out_nlc = g.add(out_nlc, bias_reshaped) - return [g.permute(out_nlc, (0, 2, 1))] - - if ( - len(x.shape) == 3 - and len(weight.shape) == 3 - and groups == 1 - and dilation == 1 - and weight.shape[2] == 3 - and padding == 1 - and stride in {1, 2} - ): - out = g.conv1d_k3(x, weight, stride=stride) - if bias is not None: - bias_reshaped = g.reshape(bias, (1, int(weight.shape[0]), 1)) - out, bias_reshaped = _legalize_elementwise_binary_inputs(g, out, bias_reshaped) - out = g.add(out, bias_reshaped) - return [out] - - if dilation != 1: - raise NotImplementedError(f"conv1d with dilation != 1 is unsupported by generic lowering: {dilation}") - if padding != 0: - if len(x.shape) != 3: - raise NotImplementedError(f"generic conv1d padding expects rank-3 input, got {x.shape}") - batch_size, channels, _ = (int(dim) for dim in x.shape) - pad_dtype = _torch_dtype_for_graph_dtype(x.dtype) - pieces: list[Tensor] = [] - if padding > 0: - left = _materialize_constant_tensor( - g, - torch.zeros((batch_size, channels, padding), dtype=pad_dtype), - ) - pieces.append(left) - pieces.append(x) - if padding > 0: - right = _materialize_constant_tensor( - g, - torch.zeros((batch_size, channels, padding), dtype=pad_dtype), - ) - pieces.append(right) - x = g.cat(pieces, axis=2) - if groups != 1: - if len(x.shape) != 3 or len(weight.shape) != 3: - raise NotImplementedError( - f"grouped conv1d lowering expects rank-3 input and weight, got {x.shape} and {weight.shape}" - ) - input_channels = int(x.shape[1]) - if input_channels % groups != 0: - raise NotImplementedError( - f"grouped conv1d requires input channels divisible by groups, got C_in={input_channels}, groups={groups}" - ) - output_channels = int(weight.shape[0]) - if output_channels % groups != 0: - raise NotImplementedError( - f"grouped conv1d requires output channels divisible by groups, got C_out={output_channels}, groups={groups}" - ) - input_channels_per_group = input_channels // groups - output_channels_per_group = output_channels // groups - if int(weight.shape[1]) != input_channels_per_group: - raise NotImplementedError( - "grouped conv1d weight has incompatible per-group input channels: " - f"weight C_in/group={int(weight.shape[1])}, expected {input_channels_per_group}" - ) - outputs: list[Tensor] = [] - for group_index in range(groups): - x_group = g.slice( - x, - axis=1, - start=group_index * input_channels_per_group, - length=input_channels_per_group, - ) - weight_group = g.slice( - weight, - axis=0, - start=group_index * output_channels_per_group, - length=output_channels_per_group, - ) - bias_group = None - if bias is not None: - bias_group = g.slice( - bias, - axis=0, - start=group_index * output_channels_per_group, - length=output_channels_per_group, - ) - outputs.append(g.conv1d(x_group, weight_group, bias=bias_group, stride=stride)) - if len(outputs) == 1: - return [outputs[0]] - return [g.cat(outputs, axis=1)] - return [g.conv1d(x, weight, bias=bias, stride=stride)] - - if op == "conv2d": - x = _tensor(env, node.inputs[0]) - weight = _tensor(env, node.inputs[1]) - bias = _tensor(env, node.inputs[2]) if len(node.inputs) > 2 else None - stride = int(node.attrs.get("stride", 1)) - padding = int(node.attrs.get("padding", 0)) - dilation = int(node.attrs.get("dilation", 1)) - groups = int(node.attrs.get("groups", 1)) - - if len(x.shape) != 4 or len(weight.shape) != 4: - raise NotImplementedError(f"conv2d lowering expects rank-4 tensors, got {x.shape} and {weight.shape}") - if dilation != 1: - raise NotImplementedError(f"conv2d with dilation != 1 is unsupported: {dilation}") - - if stride == 2 and padding == 0 and groups == 1 and weight.shape[2:] == (3, 3): - batch_size, channels, height, width = (int(dim) for dim in x.shape) - pad_dtype = _torch_dtype_for_graph_dtype(x.dtype) - top = _materialize_constant_tensor( - g, - torch.zeros((batch_size, channels, 1, width), dtype=pad_dtype), - ) - bottom = _materialize_constant_tensor( - g, - torch.zeros((batch_size, channels, 1, width), dtype=pad_dtype), - ) - x = g.cat([top, x, bottom], axis=2) - left = _materialize_constant_tensor( - g, - torch.zeros((batch_size, channels, height + 2, 1), dtype=pad_dtype), - ) - right = _materialize_constant_tensor( - g, - torch.zeros((batch_size, channels, height + 2, 1), dtype=pad_dtype), - ) - x = g.cat([left, x, right], axis=3) - y = g.conv2d_k3s2p1(x, weight, bias=bias) - output_height = ((height - 3) // 2) + 1 - output_width = ((width - 3) // 2) + 1 - y = g.slice(y, axis=2, start=1, length=output_height) - y = g.slice(y, axis=3, start=1, length=output_width) - return [y] - - if stride == 2 and padding == 1 and groups == 1 and weight.shape[2:] == (3, 3): - return [g.conv2d_k3s2p1(x, weight, bias=bias)] - - if ( - stride == 2 - and padding == 1 - and groups == x.shape[1] == weight.shape[0] - and weight.shape[1] == 1 - and weight.shape[2:] == (3, 3) - ): - return [g.conv2d_depthwise_k3s2p1(x, weight, bias=bias)] - - if ( - stride == 1 - and padding == 0 - and groups == 1 - and weight.shape[2:] == (1, 1) - ): - return [g.conv2d_pointwise_1x1(x, weight, bias=bias)] - - if stride == 1 and padding == 1 and groups == 1 and weight.shape[2:] == (3, 3): - return [g.conv2d_k3s1p1(x, weight, bias=bias)] - - return [g.conv2d(x, weight, bias=bias, stride=stride, padding=padding, dilation=dilation, groups=groups)] - - if op == "layer_norm": - x = _tensor(env, node.inputs[0]) - weight = _tensor(env, node.inputs[1]) - bias = _tensor(env, node.inputs[2]) if len(node.inputs) > 2 else None - return [g.layer_norm(x, weight, bias=bias, eps=float(node.attrs["eps"]))] - - if op == "rms_norm": - x = _tensor(env, node.inputs[0]) - weight = _tensor(env, node.inputs[1]) - weight_offset = float(node.attrs.get("weight_offset", 0.0) or 0.0) - if weight_offset != 0.0: - weight = g.scalar_add(weight, weight_offset) - reshape_back: tuple[int, ...] | None = None - if len(x.shape) > 2: - reshape_back = tuple(int(dim) for dim in x.shape) - x = _flatten_to_2d_for_linear(g, x) - out = g.rms_norm(x, weight, eps=float(node.attrs["eps"])) - if reshape_back is not None: - out = g.reshape(out, reshape_back) - return [out] - - if op == "rope": - rope_input = _tensor(env, node.inputs[0]) - if _rope_input_is_bhsd(ir, node.inputs[0]): - rope_input = g.permute(rope_input, (0, 2, 1, 3)) - rope_out = g.rope( - rope_input, - float(node.attrs["theta"]), - position_offset=int(node.attrs.get("position_offset", 0)), - ) - return [g.permute(rope_out, (0, 2, 1, 3))] - return [ - g.rope( - rope_input, - float(node.attrs["theta"]), - position_offset=int(node.attrs.get("position_offset", 0)), - ) - ] - - if op == "rel_pos_bias": - query = _attention_tensor(env, node.inputs[0]) - if len(query.shape) != 4: - raise NotImplementedError(f"rel_pos_bias expects rank-4 query input, got {query.shape}") - relative_key = _tensor(env, node.inputs[1]) - if len(relative_key.shape) != 4: - raise NotImplementedError(f"rel_pos_bias expects rank-4 relative_key input, got {relative_key.shape}") - if int(query.shape[2]) == int(relative_key.shape[2]) and int(query.shape[3]) == int(relative_key.shape[3]): - pass - elif int(query.shape[1]) == int(relative_key.shape[2]) and int(query.shape[3]) == int(relative_key.shape[3]): - query = g.permute(query, (0, 2, 1, 3)) - else: - raise NotImplementedError( - "rel_pos_bias query/relative_key layout mismatch: " - f"query={query.shape}, relative_key={relative_key.shape}" - ) - return [g.rel_pos_bias(query, relative_key, float(node.attrs.get("scale", 1.0)))] - - if op in {"gated_deltanet_prefill", "gated_deltanet_decode"}: - x = _tensor(env, node.inputs[0]) - qkv_weight = _tensor(env, node.inputs[1]) - a_weight = _tensor(env, node.inputs[2]) - b_weight = _tensor(env, node.inputs[3]) - norm_weight = _tensor(env, node.inputs[4]) - - input_index = 5 - z_weight = None - if bool(node.attrs.get("has_z", False)): - z_weight = _tensor(env, node.inputs[input_index]) - input_index += 1 - dt_bias = None - if bool(node.attrs.get("has_dt_bias", False)): - dt_bias = _tensor(env, node.inputs[input_index]) - input_index += 1 - a_log = None - if bool(node.attrs.get("has_a_log", False)): - a_log = _tensor(env, node.inputs[input_index]) - input_index += 1 - conv_weight = None - if bool(node.attrs.get("has_conv", False)): - conv_weight = _tensor(env, node.inputs[input_index]) - - if len(x.shape) != 3: - raise NotImplementedError(f"{op} currently expects rank-3 normalized input, got {x.shape}") - - batch_size, seq_len, hidden_dim = (int(dim) for dim in x.shape) - if batch_size != 1: - raise NotImplementedError(f"{op} currently supports batch size 1, got {x.shape}") - - num_k_heads = int(node.attrs["num_k_heads"]) - num_v_heads = int(node.attrs["num_v_heads"]) - key_dim = int(node.attrs["key_dim"]) - value_dim = int(node.attrs["value_dim"]) - eps = float(node.attrs.get("eps", 1e-6)) - chunk_size = int(node.attrs.get("chunk_size", 64)) - - mixed_qkv = _matmul_with_quantized_rhs_legalization( - g, - _flatten_to_2d_for_linear(g, x), - qkv_weight, - pretransposed_rhs=True, - ) - mixed_qkv_dim = int(qkv_weight.shape[0]) - mixed_qkv = g.reshape(mixed_qkv, (batch_size, seq_len, mixed_qkv_dim)) - if conv_weight is not None: - mixed_qkv = _lower_gated_deltanet_conv1d( - g, - ir, - env, - node, - mixed_qkv=mixed_qkv, - conv_weight=conv_weight, - batch_size=batch_size, - seq_len=seq_len, - mixed_qkv_dim=mixed_qkv_dim, - op=op, - ) - mixed_qkv = g.silu(mixed_qkv) - - q_proj_dim = num_k_heads * key_dim - v_proj_dim = num_v_heads * value_dim - k_proj_dim = num_k_heads * key_dim - q_proj = g.slice(mixed_qkv, axis=2, start=0, length=q_proj_dim) - k_proj = g.slice(mixed_qkv, axis=2, start=q_proj_dim, length=k_proj_dim) - v_proj = g.slice(mixed_qkv, axis=2, start=q_proj_dim + k_proj_dim, length=v_proj_dim) - - q_4d = g.reshape(q_proj, (batch_size, seq_len, num_k_heads, key_dim)) - k_4d = g.reshape(k_proj, (batch_size, seq_len, num_k_heads, key_dim)) - v_4d = g.reshape(v_proj, (batch_size, seq_len, num_v_heads, value_dim)) - - q_norm = g.sum(g.multiply(q_4d, q_4d), axis=3) - q_norm = g.scalar_sqrt(g.scalar_add(q_norm, eps)) - q_norm = g.reshape(q_norm, (batch_size, seq_len, num_k_heads, 1)) - q_4d = g.divide(q_4d, q_norm) - - k_norm = g.sum(g.multiply(k_4d, k_4d), axis=3) - k_norm = g.scalar_sqrt(g.scalar_add(k_norm, eps)) - k_norm = g.reshape(k_norm, (batch_size, seq_len, num_k_heads, 1)) - k_4d = g.divide(k_4d, k_norm) - - a_logits = _matmul_with_quantized_rhs_legalization( - g, - _flatten_to_2d_for_linear(g, x), - a_weight, - pretransposed_rhs=True, - ) - a_logits = g.reshape(a_logits, (batch_size, seq_len, int(a_weight.shape[0]))) - b_logits = _matmul_with_quantized_rhs_legalization( - g, - _flatten_to_2d_for_linear(g, x), - b_weight, - pretransposed_rhs=True, - ) - b_logits = g.reshape(b_logits, (batch_size, seq_len, int(b_weight.shape[0]))) - - if dt_bias is not None: - dt_bias_2d = g.reshape(dt_bias, (1, int(dt_bias.shape[0]))) - a_logits, dt_bias_2d = _legalize_elementwise_binary_inputs(g, a_logits, dt_bias_2d) - a_logits = g.add(a_logits, dt_bias_2d) - a_softplus = _lower_softplus(g, a_logits) - - if a_log is not None: - a_log_2d = g.reshape(a_log, (1, int(a_log.shape[0]))) - neg_exp_a = g.scalar_multiply(g.scalar_exp(a_log_2d), -1.0) - neg_exp_a, a_softplus = _legalize_elementwise_binary_inputs(g, neg_exp_a, a_softplus) - gate_log = g.multiply(neg_exp_a, a_softplus) - else: - gate_log = g.scalar_multiply(a_softplus, -1.0) - beta = g.sigmoid(b_logits) - - cache_layer_key, initial_state = _lower_gated_deltanet_initial_state( - g, - ir, - env, - node, - batch_size=batch_size, - key_dim=key_dim, - num_v_heads=num_v_heads, - value_dim=value_dim, - seq_len=seq_len, - op=op, - ) - - deltanet_scale = 1.0 / math.sqrt(float(key_dim)) - if seq_len == 1: - deltanet_out = g.gated_deltanet_decode(q_4d, k_4d, v_4d, gate_log, beta, initial_state, deltanet_scale) - else: - deltanet_out = g.gated_deltanet_prefill(q_4d, k_4d, v_4d, gate_log, beta, initial_state, chunk_size, deltanet_scale) - if cache_layer_key is not None: - final_state = g.slice(deltanet_out, axis=1, start=seq_len, length=key_dim) - g.recurrent_cache_write(final_state, initial_state) - cache_entries: list[tuple[str, Tensor, Tensor]] = env.setdefault("__internal_kv_cache_state_entries", []) # type: ignore[assignment] - cache_entries.append((cache_layer_key, initial_state, initial_state)) - - y_4d = g.slice(deltanet_out, axis=1, start=0, length=seq_len) - y_2d = g.reshape(y_4d, (seq_len * num_v_heads, value_dim)) - - if z_weight is not None: - z_proj = _matmul_with_quantized_rhs_legalization( - g, - _flatten_to_2d_for_linear(g, x), - z_weight, - pretransposed_rhs=True, - ) - z_proj = g.reshape(z_proj, (seq_len * num_v_heads, value_dim)) - y_2d = g.multiply(g.rms_norm(y_2d, norm_weight, eps=eps), g.silu(z_proj)) - - return [g.reshape(y_2d, (batch_size, seq_len, num_v_heads * value_dim))] - - if op == "group_norm": - return [ - g.group_norm( - _tensor(env, node.inputs[0]), - _tensor(env, node.inputs[1]), - _tensor(env, node.inputs[2]), - num_groups=int(node.attrs["num_groups"]), - eps=float(node.attrs["eps"]), - ) - ] - - if op == "batch_norm": - axis = int(node.attrs.get("axis", 1)) - return [ - g.batch_norm( - _tensor(env, node.inputs[0]), - _tensor(env, node.inputs[1]), - _tensor(env, node.inputs[2]), - _tensor(env, node.inputs[3]), - _tensor(env, node.inputs[4]), - axis=axis, - eps=float(node.attrs["eps"]), - ) - ] - - if op == "identity": - return [env[node.inputs[0]]] - - if op == "contiguous": - return [_tensor(env, node.inputs[0])] - - if op == "advanced_index": - lowered = _try_lower_advanced_index(g, node, env, ir) - if lowered is not None: - return [lowered] - raise NotImplementedError(f"unsupported advanced_index pattern for node {node.id}") - - if op in {"masked_scatter", "aten.masked_scatter.default"}: - base = _tensor(env, node.inputs[0]) - mask = _tensor(env, node.inputs[1]) - source = _tensor(env, node.inputs[2]) - if source.dtype != base.dtype: - source = _ensure_tensor_dtype(g, source, base.dtype) - return [g.masked_scatter(base, mask, source)] - - if op == "masked_fill": - base = _tensor(env, node.inputs[0]) - mask = _lower_compare_op(g, env[node.inputs[1]], 0.0, "not_equal") - inverse_mask = g.scalar_add(g.scalar_multiply(mask, -1.0), 1.0) - fill_value = _clamp_scalar_for_dtype(float(node.attrs["value"]), base.dtype) - fill_term = g.scalar_multiply(_ensure_tensor_dtype(g, mask, base.dtype), fill_value) - base_term = _lower_binary_op(g, base, _ensure_tensor_dtype(g, inverse_mask, base.dtype), "multiply") - return [_lower_binary_op(g, fill_term, base_term, "add")] - - if op == "getitem": - source = env[node.inputs[0]] - index = int(node.attrs["index"]) - if isinstance(source, (tuple, list)): - return [source[index]] - if isinstance(source, Tensor): - normalized_index = _normalize_index(index, int(source.shape[0])) - return [g.index(source, normalized_index, axis=0)] - raise NotImplementedError(f"getitem source is not tuple/list for node {node.id}") - - raise NotImplementedError(f"unsupported IR op in lowering: {op}") - - -def _tensor(env: dict[str, Any], value_id: str) -> Tensor: - try: - value = env[value_id] - except KeyError as exc: - raise NotImplementedError(f"missing IR value during lowering: {value_id}") from exc - if not isinstance(value, Tensor): - raise TypeError(f"expected lowered tensor for {value_id}, got {type(value).__name__}") - return value - - -def _tensor_or_materialized_alias(g: Graph, env: dict[str, Any], value_id: str) -> Tensor: - value = env.get(value_id) - if isinstance(value, BroadcastAlias): - return _materialize_broadcast_alias(g, value) - return _tensor(env, value_id) - - -def _materialize_broadcast_alias(g: Graph, value: BroadcastAlias) -> Tensor: - if value.kind != "gqa_repeat_kv": - raise TypeError(f"unsupported broadcast alias: {value.kind}") - base_shape = tuple(int(dim) for dim in value.tensor.shape) - logical_shape = tuple(int(dim) for dim in value.logical_shape) - if len(base_shape) != 4 or len(logical_shape) not in {4, 5}: - raise TypeError(f"gqa_repeat_kv alias expects 4D->4D/5D shape, got {base_shape} -> {logical_shape}") - batch, kv_heads, seq_len, head_dim = base_shape - if len(logical_shape) == 5: - if logical_shape[0] != batch or logical_shape[1] != kv_heads or logical_shape[3] != seq_len or logical_shape[4] != head_dim: - raise TypeError(f"gqa_repeat_kv alias shape mismatch: {base_shape} -> {logical_shape}") - base = g.reshape(value.tensor, (batch, kv_heads, 1, seq_len, head_dim)) - return _lower_repeat(g, base, (1, 1, logical_shape[2], 1, 1)) - if logical_shape[0] != batch or logical_shape[2] != seq_len or logical_shape[3] != head_dim: - raise TypeError(f"gqa_repeat_kv alias shape mismatch: {base_shape} -> {logical_shape}") - if logical_shape[1] % max(kv_heads, 1) != 0: - raise TypeError(f"gqa_repeat_kv head count mismatch: {base_shape} -> {logical_shape}") - repeats = logical_shape[1] // max(kv_heads, 1) - expanded = g.reshape(value.tensor, (batch, kv_heads, 1, seq_len, head_dim)) - repeated = _lower_repeat(g, expanded, (1, 1, repeats, 1, 1)) - return g.reshape(repeated, logical_shape) - - -def _attention_tensor(env: dict[str, Any], value_id: str) -> Tensor: - try: - value = env[value_id] - except KeyError as exc: - raise NotImplementedError(f"missing IR value during lowering: {value_id}") from exc - if isinstance(value, BroadcastAlias): - if value.kind != "gqa_repeat_kv": - raise TypeError(f"unsupported broadcast alias for attention input {value_id}: {value.kind}") - return value.tensor - if not isinstance(value, Tensor): - raise TypeError(f"expected lowered tensor for {value_id}, got {type(value).__name__}") - return value - - -def _resolve_attention_scale(node: IRNode, query: Tensor) -> float: - scale = float(node.attrs.get("scale", 0.0) or 0.0) - if scale > 0.0: - return scale - if len(query.shape) >= 1 and int(query.shape[-1]) > 0: - return float(int(query.shape[-1]) ** -0.5) - return 1.0 - - -def _normalize_attention_mask_for_cactus(g: Graph, mask: Tensor, query: Tensor) -> Tensor: - """Convert broadcastable HF masks into the native Cactus attention contract.""" - - mask_shape = tuple(int(dim) for dim in mask.shape) - query_shape = tuple(int(dim) for dim in query.shape) - if len(query_shape) != 4: - return mask - - batch_size, seq_len, num_heads, _ = query_shape - if len(mask_shape) == 2 and batch_size == 1 and mask_shape[0] == seq_len: - return g.reshape(mask, (1, mask_shape[0], mask_shape[1])) - - if len(mask_shape) == 4: - if ( - mask_shape[0] == batch_size - and mask_shape[1] == 1 - and mask_shape[2] == seq_len - ): - return g.reshape(mask, (mask_shape[0], mask_shape[2], mask_shape[3])) - if ( - mask_shape[0] == batch_size - and mask_shape[1] == num_heads - and mask_shape[2] == seq_len - ): - return mask - - return mask - - -def _lower_projected_attention_tensor( - g: Graph, - hidden: Tensor, - weight: Tensor, - bias: Tensor | None, - target_shape: tuple[int, ...], -) -> Tensor: - if len(target_shape) != 4: - raise NotImplementedError(f"projected attention tensor expects rank-4 target shape, got {target_shape}") - - linear_input = _flatten_to_2d_for_linear(g, hidden) - out = _matmul_with_quantized_rhs_legalization(g, linear_input, weight, pretransposed_rhs=True) - if bias is not None: - out = g.add(out, bias) - return g.reshape(out, target_shape) - - -def _normalize_attention_add_tensor(g: Graph, tensor: Tensor, target_shape: tuple[int, ...]) -> Tensor: - if tuple(tensor.shape) == target_shape: - return tensor - - if len(tensor.shape) != 4 or len(target_shape) != 4: - raise NotImplementedError( - f"unsupported attention add tensor shape {tuple(tensor.shape)} for target {target_shape}" - ) - - batch, seq_len, heads, head_dim = target_shape - shape = tuple(int(v) for v in tensor.shape) - - if shape[0] in (1, batch) and shape[1] == heads and shape[2] in (1, seq_len) and shape[3] == head_dim: - return g.permute(tensor, (0, 2, 1, 3)) - if shape[0] in (1, batch) and shape[1] in (1, seq_len) and shape[2] == heads and shape[3] == head_dim: - return tensor - - raise NotImplementedError( - f"unsupported attention add tensor shape {shape} for target {target_shape}" - ) - - -def _rope_input_is_bhsd(ir: IRGraph, value_id: str) -> bool: - value = ir.values.get(value_id) - if value is None or value.producer is None: - return False - node = ir.nodes.get(value.producer) - if node is None: - return False - if node.op == "permute": - permutation = tuple(int(dim) for dim in node.attrs.get("permutation", ())) - return permutation == (0, 2, 1, 3) - if node.op == "transpose": - return int(node.attrs.get("dim0", -1)) == 1 and int(node.attrs.get("dim1", -1)) == 2 - return False - - -def _lower_softplus(g: Graph, x: Tensor) -> Tensor: - # Stable fp16 softplus: relu(x) + log(1 + exp(-abs(x))). - abs_x = g.abs(x) - neg_abs_x = g.scalar_multiply(abs_x, -1.0) - exp_term = g.scalar_exp(neg_abs_x) - log_term = g.scalar_log(g.scalar_add(exp_term, 1.0)) - return g.add(g.relu(x), log_term) - - -def _normalize_dim(dim: int, rank: int) -> int: - if dim < 0: - dim += rank - return dim - - -def _normalize_reduction_axes(axis: Any, rank: int) -> tuple[int, ...]: - if axis is None: - return tuple(range(rank)) - if isinstance(axis, int): - return (_normalize_dim(axis, rank),) - if isinstance(axis, (list, tuple)): - normalized: list[int] = [] - seen: set[int] = set() - for raw_axis in axis: - if not isinstance(raw_axis, int): - raise NotImplementedError("reduction axes must be integers") - reduced_axis = _normalize_dim(raw_axis, rank) - if reduced_axis in seen: - continue - seen.add(reduced_axis) - normalized.append(reduced_axis) - if normalized: - return tuple(sorted(normalized)) - raise NotImplementedError("reduction axes must be an int, a tuple/list of ints, or None") - - -def _lower_reduction(g: Graph, op: str, x: Tensor, *, axes: tuple[int, ...]) -> Tensor: - if not axes: - return x - - fn = getattr(g, op) - if len(axes) == 1: - return fn(x, axes[0]) - - if op == "variance": - flattened = g.flatten(x, start_dim=axes[0], end_dim=axes[-1]) - collapsed_axes = set(axes) - expected_axes = set(range(axes[0], axes[-1] + 1)) - if collapsed_axes != expected_axes: - raise NotImplementedError("variance currently requires contiguous multi-axis reductions") - return fn(flattened, axes[0]) - - reduced = x - for axis in sorted(axes, reverse=True): - reduced = fn(reduced, axis) - return reduced - - -def _normalize_index(index: int, dim_size: int) -> int: - if index < 0: - index += dim_size - return index - - -def _static_shape(shape: Any) -> tuple[int, ...] | None: - if not isinstance(shape, tuple): - return None - dims: list[int] = [] - for dim in shape: - try: - dims.append(int(dim)) - except Exception: - return None - return tuple(dims) - - -def _ir_value_dtype(ir: IRGraph, value_id: str) -> str | None: - value = ir.values.get(value_id) - if value is None or value.dtype is None: - return None - return str(value.dtype).strip().lower() - - -def _is_integer_index_dtype(dtype: str | None) -> bool: - return dtype in {"uint8", "uint16", "uint32", "uint64", "int8", "int16", "int32", "int64"} - - -def _constant_tensor_from_ir(ir: IRGraph, value_id: str) -> torch.Tensor | None: - const = ir.constants.get(value_id) - if isinstance(const, torch.nn.Parameter): - const = const.detach() - if isinstance(const, torch.Tensor): - return const.detach().cpu() - return None - - -def _constant_scalar_from_ir(ir: IRGraph, value_id: str) -> float | None: - const = _constant_tensor_from_ir(ir, value_id) - if const is None: - return None - if const.numel() != 1: - raise NotImplementedError( - f"clamp bound {value_id!r} must be a scalar constant, got shape={tuple(const.shape)}" - ) - return float(const.reshape(()).item()) - - -def _try_lower_identity_broadcast_advanced_index( - g: Graph, - node: IRNode, - env: dict[str, Any], - ir: IRGraph, -) -> Tensor | None: - if len(node.inputs) != 3: - return None - - source = _tensor(env, node.inputs[0]) - source_shape = tuple(int(dim) for dim in source.shape) - if len(source_shape) != 2 or int(source_shape[0]) != 1: - return None - - output_shape = _static_shape(ir.values[node.outputs[0]].shape if node.outputs and node.outputs[0] in ir.values else None) - if output_shape is None: - return None - - batch_index = _constant_tensor_from_ir(ir, node.inputs[1]) - token_index = _constant_tensor_from_ir(ir, node.inputs[2]) - if batch_index is None or token_index is None: - index_values = [env.get(value_id) for value_id in node.inputs[1:]] - index_shapes = [ - tuple(int(dim) for dim in value.shape) - for value in index_values - if isinstance(value, Tensor) - ] - source_elems = 1 - for dim in source_shape: - source_elems *= int(dim) - output_elems = 1 - for dim in output_shape: - output_elems *= int(dim) - if ( - int(source_shape[0]) == 1 - and output_elems == source_elems - and index_shapes - and all(shape and all(int(dim) in {1, int(source_shape[1])} for dim in shape) for shape in index_shapes) - ): - return g.reshape(source, output_shape) - return None - - try: - batch_index = torch.broadcast_to(batch_index.to(torch.int64), output_shape) - token_index = torch.broadcast_to(token_index.to(torch.int64), output_shape) - except RuntimeError: - return None - - if torch.count_nonzero(batch_index).item() != 0: - return None - - flat_index = token_index.reshape(-1) - expected = torch.arange(int(source_shape[1]), dtype=torch.int64) - if flat_index.numel() != expected.numel() or not torch.equal(flat_index, expected): - return None - - return g.reshape(source, output_shape) - - -def _try_lower_embedding_advanced_index( - g: Graph, - node: IRNode, - env: dict[str, Any], - ir: IRGraph, -) -> Tensor | None: - if len(node.inputs) != 2: - return None - if not _is_integer_index_dtype(_ir_value_dtype(ir, node.inputs[1])): - return None - - source = _tensor(env, node.inputs[0]) - indices = _tensor(env, node.inputs[1]) - source_shape = tuple(int(dim) for dim in source.shape) - if len(source_shape) != 2: - return None - - output_shape = _static_shape(ir.values[node.outputs[0]].shape if node.outputs and node.outputs[0] in ir.values else None) - gathered = g.embedding_from_tensor(source, indices) - if output_shape is not None and tuple(int(dim) for dim in gathered.shape) != tuple(int(dim) for dim in output_shape): - gathered = g.reshape(gathered, output_shape) - return gathered - - -def _try_lower_prefix_mask_advanced_index( - g: Graph, - node: IRNode, - env: dict[str, Any], - ir: IRGraph, -) -> Tensor | None: - if len(node.inputs) != 2: - return None - if _is_integer_index_dtype(_ir_value_dtype(ir, node.inputs[1])): - return None - - source = _tensor(env, node.inputs[0]) - mask = _tensor(env, node.inputs[1]) - source_shape = tuple(int(dim) for dim in source.shape) - mask_shape = tuple(int(dim) for dim in mask.shape) - if not mask_shape or len(mask_shape) > len(source_shape): - return None - if source_shape[: len(mask_shape)] != mask_shape: - return None - - output_shape = _static_shape(ir.values[node.outputs[0]].shape if node.outputs and node.outputs[0] in ir.values else None) - if output_shape: - prefix_elems = 1 - for dim in mask_shape: - prefix_elems *= int(dim) - trailing_shape = source_shape[len(mask_shape) :] - flattened_shape = (prefix_elems, *trailing_shape) - flattened = source if source_shape == flattened_shape else g.reshape(source, flattened_shape) - expected_prefix = int(output_shape[0]) - if 0 <= expected_prefix <= prefix_elems: - selected = g.slice(flattened, axis=0, start=0, length=expected_prefix) - if tuple(int(dim) for dim in selected.shape) != tuple(int(dim) for dim in output_shape): - selected = g.reshape(selected, output_shape) - return selected - - selected = g.masked_select_prefix(source, mask) - if output_shape is None or not output_shape: - return selected - if len(output_shape) != len(selected.shape): - return selected - expected_prefix = int(output_shape[0]) - if expected_prefix < int(selected.shape[0]): - selected = g.slice(selected, axis=0, start=0, length=expected_prefix) - return selected - - -def _try_lower_advanced_index( - g: Graph, - node: IRNode, - env: dict[str, Any], - ir: IRGraph, -) -> Tensor | None: - lowered = _try_lower_identity_broadcast_advanced_index(g, node, env, ir) - if lowered is not None: - return lowered - lowered = _try_lower_prefix_mask_advanced_index(g, node, env, ir) - if lowered is not None: - return lowered - return _try_lower_embedding_advanced_index(g, node, env, ir) - - -def _normalize_slice_end(end: int, dim_size: int) -> int: - if end < 0: - end += dim_size - return max(0, min(end, dim_size)) - - -def _legalize_elementwise_binary_inputs(g: Graph, lhs: Tensor, rhs: Tensor) -> tuple[Tensor, Tensor]: - if lhs.shape == rhs.shape: - return lhs, rhs - - lhs_rank = len(lhs.shape) - rhs_rank = len(rhs.shape) - - if lhs_rank > rhs_rank: - rhs = _reshape_for_trailing_broadcast(g, rhs, lhs.shape) - elif rhs_rank > lhs_rank: - lhs = _reshape_for_trailing_broadcast(g, lhs, rhs.shape) - - if lhs.shape == rhs.shape: - return lhs, rhs - - return lhs, rhs - - -def _is_scalar_like(value: Any) -> bool: - return isinstance(value, (int, float, bool)) - - -def _lower_binary_op(g: Graph, lhs_value: Any, rhs_value: Any, op: str) -> Tensor: - if isinstance(lhs_value, Tensor) and isinstance(rhs_value, Tensor): - lhs, rhs = _legalize_elementwise_binary_inputs(g, lhs_value, rhs_value) - target_dtype = Graph.FP16 if lhs.dtype == Graph.FP16 and rhs.dtype == Graph.FP16 else Graph.FP32 - lhs = _ensure_tensor_dtype(g, lhs, target_dtype) - rhs = _ensure_tensor_dtype(g, rhs, target_dtype) - if op == "add": - return g.add(lhs, rhs) - if op == "subtract": - return g.subtract(lhs, rhs) - if op == "multiply": - return g.multiply(lhs, rhs) - if op == "divide": - return g.divide(lhs, rhs) - raise NotImplementedError(f"unsupported binary op: {op}") - - if isinstance(lhs_value, Tensor) and _is_scalar_like(rhs_value): - lhs_value = _ensure_scalar_math_tensor(g, lhs_value) - scalar = float(rhs_value) - if op == "add": - return g.scalar_add(lhs_value, scalar) - if op == "subtract": - return g.scalar_subtract(lhs_value, scalar) - if op == "multiply": - return g.scalar_multiply(lhs_value, scalar) - if op == "divide": - return g.scalar_divide(lhs_value, scalar) - raise NotImplementedError(f"unsupported binary op: {op}") - - if _is_scalar_like(lhs_value) and isinstance(rhs_value, Tensor): - rhs_value = _ensure_scalar_math_tensor(g, rhs_value) - scalar = float(lhs_value) - if op == "add": - return g.scalar_add(rhs_value, scalar) - if op == "subtract": - return g.scalar_add(g.scalar_multiply(rhs_value, -1.0), scalar) - if op == "multiply": - return g.scalar_multiply(rhs_value, scalar) - if op == "divide": - raise NotImplementedError("scalar/tensor divide is not directly supported by Cactus graph ops") - raise NotImplementedError(f"unsupported binary op: {op}") - - raise TypeError( - f"unsupported lowered operand types for {op}: " - f"{type(lhs_value).__name__}, {type(rhs_value).__name__}" - ) - - -def _lower_compare_op(g: Graph, lhs_value: Any, rhs_value: Any, op: str) -> Tensor: - if op == "equal": - not_equal = _lower_compare_op(g, lhs_value, rhs_value, "not_equal") - return g.scalar_not_equal(not_equal, 1.0) - - if op == "greater": - delta = _lower_binary_op(g, lhs_value, rhs_value, "subtract") - delta = _ensure_fp16_tensor(g, delta) - return g.scalar_not_equal(g.relu(delta), 0.0) - - if op == "less": - return _lower_compare_op(g, rhs_value, lhs_value, "greater") - - if op == "greater_equal": - less = _lower_compare_op(g, lhs_value, rhs_value, "less") - return g.scalar_not_equal(less, 1.0) - - if op == "less_equal": - greater = _lower_compare_op(g, lhs_value, rhs_value, "greater") - return g.scalar_not_equal(greater, 1.0) - - if op != "not_equal": - raise NotImplementedError(f"unsupported compare op: {op}") - - if isinstance(lhs_value, Tensor) and isinstance(rhs_value, Tensor): - lhs, rhs = _legalize_elementwise_binary_inputs(g, lhs_value, rhs_value) - target_dtype = Graph.FP16 if lhs.dtype == Graph.FP16 and rhs.dtype == Graph.FP16 else Graph.FP32 - lhs = _ensure_tensor_dtype(g, lhs, target_dtype) - rhs = _ensure_tensor_dtype(g, rhs, target_dtype) - return g.not_equal(lhs, rhs) - - if isinstance(lhs_value, Tensor) and _is_scalar_like(rhs_value): - lhs_value = _ensure_tensor_dtype( - g, - lhs_value, - Graph.FP16 if lhs_value.dtype == Graph.FP16 else Graph.FP32, - ) - return g.scalar_not_equal(lhs_value, float(rhs_value)) - - if _is_scalar_like(lhs_value) and isinstance(rhs_value, Tensor): - rhs_value = _ensure_tensor_dtype( - g, - rhs_value, - Graph.FP16 if rhs_value.dtype == Graph.FP16 else Graph.FP32, - ) - return g.scalar_not_equal(rhs_value, float(lhs_value)) - - if _is_scalar_like(lhs_value) and _is_scalar_like(rhs_value): - result = 1.0 if float(lhs_value) != float(rhs_value) else 0.0 - return _materialize_constant_tensor(g, torch.tensor([result], dtype=torch.float32)) - - raise TypeError( - f"unsupported lowered operand types for {op}: " - f"{type(lhs_value).__name__}, {type(rhs_value).__name__}" - ) - - -def _reshape_for_trailing_broadcast(g: Graph, tensor: Tensor, target_shape: tuple[int, ...]) -> Tensor: - tensor_shape = tuple(tensor.shape) - target_rank = len(target_shape) - tensor_rank = len(tensor_shape) - - if tensor_rank > target_rank: - return tensor - - padded_shape = (1,) * (target_rank - tensor_rank) + tensor_shape - - # Only legalize cases that are valid trailing broadcasts, e.g.: - # (H,) -> (1, H), (1, 1, H), etc. - for src_dim, tgt_dim in zip(padded_shape, target_shape): - if src_dim != 1 and src_dim != tgt_dim: - return tensor - - if padded_shape == tensor_shape: - return tensor - - return g.reshape(tensor, padded_shape) - - -def _ensure_fp16_tensor(g: Graph, tensor: Tensor) -> Tensor: - if tensor.dtype == Graph.FP16: - return tensor - return g.precision_cast(tensor, Graph.FP16) - - -def _ensure_scalar_math_tensor(g: Graph, tensor: Tensor) -> Tensor: - return tensor - - -def _ensure_tensor_dtype(g: Graph, tensor: Tensor, dtype: int) -> Tensor: - if tensor.dtype == dtype: - return tensor - return g.precision_cast(tensor, dtype) - - -def _cat_with_legalized_dtype(g: Graph, tensors: list[Tensor], *, axis: int, dtype: int | None) -> Tensor: - if not tensors: - raise NotImplementedError("cat requires at least one tensor") - if dtype is None: - dtypes = {tensor.dtype for tensor in tensors} - dtype = tensors[0].dtype if len(dtypes) == 1 else Graph.FP16 - legalized = [_ensure_tensor_dtype(g, tensor, dtype) for tensor in tensors] - return g.cat(legalized, axis=axis) - - -def _invert_permutation(permutation: tuple[int, ...]) -> tuple[int, ...]: - inverse = [0] * len(permutation) - for index, source_axis in enumerate(permutation): - inverse[int(source_axis)] = index - return tuple(inverse) - - -def _lower_static_strided_slice_via_gather( - g: Graph, - x: Tensor, - *, - axis: int, - indices: list[int], -) -> Tensor: - if not indices: - return g.slice(x, axis=axis, start=0, length=0) - - indices_tensor = _materialize_constant_tensor(g, torch.tensor(indices, dtype=torch.float32)) - if axis == 0: - return g.gather(x, indices_tensor) - - permutation = (axis, *[dim for dim in range(len(x.shape)) if dim != axis]) - transposed = g.permute(x, permutation) - gathered = g.gather(transposed, indices_tensor) - inverse_permutation = _invert_permutation(permutation) - return g.permute(gathered, inverse_permutation) - - -def _lower_where_branch_term(g: Graph, branch_value: Any, mask: Tensor, *, dtype: int | None) -> Tensor: - if dtype is not None: - mask = _ensure_tensor_dtype(g, mask, dtype) - if isinstance(branch_value, Tensor): - branch_mask = mask if dtype is None else _ensure_tensor_dtype(g, mask, branch_value.dtype) - return _lower_binary_op(g, branch_value, branch_mask, "multiply") - return g.scalar_multiply(mask, _clamp_scalar_for_dtype(float(branch_value), mask.dtype)) - - -def _clamp_scalar_for_dtype(value: float, dtype: int | None) -> float: - if dtype != Graph.FP16: - return value - if math.isnan(value): - return value - if value == math.inf: - return 65504.0 - if value == -math.inf: - return -65504.0 - return max(-65504.0, min(65504.0, value)) - - -def _lower_where_op(g: Graph, node: IRNode, env: dict[str, Any]) -> Tensor: - condition = _lower_compare_op(g, env[node.inputs[0]], 0.0, "not_equal") - false_mask = g.scalar_add(g.scalar_multiply(condition, -1.0), 1.0) - - input_index = 1 - if bool(node.attrs.get("true_is_scalar", False)): - true_value: Any = float(node.attrs["true_value"]) - else: - true_value = env[node.inputs[input_index]] - input_index += 1 - - if bool(node.attrs.get("false_is_scalar", False)): - false_value: Any = float(node.attrs["false_value"]) - else: - false_value = env[node.inputs[input_index]] - - result_dtype: int | None = None - if isinstance(true_value, Tensor): - result_dtype = true_value.dtype - elif isinstance(false_value, Tensor): - result_dtype = false_value.dtype - - true_term = _lower_where_branch_term(g, true_value, condition, dtype=result_dtype) - false_term = _lower_where_branch_term(g, false_value, false_mask, dtype=result_dtype) - return _lower_binary_op(g, true_term, false_term, "add") - - -def _flatten_to_2d_for_linear(g: Graph, tensor: Tensor) -> Tensor: - shape = tuple(tensor.shape) - if len(shape) <= 2: - return tensor - leading = 1 - for dim in shape[:-1]: - leading *= int(dim) - return g.view(tensor, (leading, int(shape[-1]))) - - -def _broadcast_batch_shape(lhs_batch: tuple[int, ...], rhs_batch: tuple[int, ...]) -> tuple[int, ...] | None: - rank = max(len(lhs_batch), len(rhs_batch)) - lhs_padded = (1,) * (rank - len(lhs_batch)) + lhs_batch - rhs_padded = (1,) * (rank - len(rhs_batch)) + rhs_batch - result: list[int] = [] - for lhs_dim, rhs_dim in zip(lhs_padded, rhs_padded): - if lhs_dim == rhs_dim: - result.append(int(lhs_dim)) - elif lhs_dim == 1: - result.append(int(rhs_dim)) - elif rhs_dim == 1: - result.append(int(lhs_dim)) - else: - return None - return tuple(result) - - -def _project_broadcast_coord(coord: tuple[int, ...], operand_batch: tuple[int, ...]) -> tuple[int, ...]: - if not operand_batch: - return () - offset = len(coord) - len(operand_batch) - projected: list[int] = [] - for axis, dim in enumerate(operand_batch): - value = coord[offset + axis] - projected.append(0 if int(dim) == 1 else int(value)) - return tuple(projected) - - -def _index_batch_prefix(g: Graph, tensor: Tensor, coord: tuple[int, ...]) -> Tensor: - for index_value in coord: - tensor = g.index(tensor, index_value, axis=0) - return tensor - - -def _lower_static_batched_matmul( - g: Graph, - lhs: Tensor, - rhs: Tensor, - *, - output_dtype: int | None = None, -) -> Tensor | None: - lhs_shape = tuple(int(dim) for dim in lhs.shape) - rhs_shape = tuple(int(dim) for dim in rhs.shape) - if len(lhs_shape) < 3 or len(rhs_shape) < 3: - return None - if lhs_shape[-1] != rhs_shape[-2]: - return None - - lhs_batch = lhs_shape[:-2] - rhs_batch = rhs_shape[:-2] - batch_shape = _broadcast_batch_shape(lhs_batch, rhs_batch) - if batch_shape is None: - return None - - output_matrix_shape = (int(lhs_shape[-2]), int(rhs_shape[-1])) - rank = len(batch_shape) - - def _build(dim: int, coord_prefix: tuple[int, ...]) -> Tensor: - if dim == rank: - lhs_coord = _project_broadcast_coord(coord_prefix, lhs_batch) - rhs_coord = _project_broadcast_coord(coord_prefix, rhs_batch) - lhs_slice = _index_batch_prefix(g, lhs, lhs_coord) - rhs_slice = _index_batch_prefix(g, rhs, rhs_coord) - out = _matmul_with_quantized_rhs_legalization( - g, - lhs_slice, - rhs_slice, - output_dtype=output_dtype, - ) - return g.reshape(out, (1,) * rank + output_matrix_shape) - - pieces = [_build(dim + 1, coord_prefix + (index_value,)) for index_value in range(int(batch_shape[dim]))] - if len(pieces) == 1: - return pieces[0] - return g.cat(pieces, axis=dim) - - return _build(0, ()) - - -def _should_lower_gemma4_decoder_attention_without_kernel(ir: IRGraph, node: IRNode) -> bool: - if os.environ.get("CACTUS_GEMMA4_DECODER_MANUAL_ATTENTION", "0") != "1": - return False - family = str(ir.meta.get("adapter_family") or ir.meta.get("family") or "").strip().lower() - component = str(ir.meta.get("component", "") or "").strip().lower() - if family != "gemma4" or component != "decoder": - return False - return node.op in {"attention", "scaled_dot_product_attention"} - - -def _lower_gemma4_decoder_attention_without_kernel( - g: Graph, - ir: IRGraph, - env: dict[str, Any], - node: IRNode, -) -> Tensor: - query = _attention_tensor(env, node.inputs[0]) - key = _attention_tensor(env, node.inputs[1]) - value = _attention_tensor(env, node.inputs[2]) - - key_transposed = g.permute(key, (0, 1, 3, 2)) - scores = _lower_static_batched_matmul(g, query, key_transposed, output_dtype=Graph.FP16) - if scores is None: - raise NotImplementedError("Gemma4 decoder attention requires static batched matmul support") - - scale = float(node.attrs.get("scale", 1.0)) - if scale != 1.0: - scores = g.scalar_multiply(_ensure_scalar_math_tensor(g, scores), scale) - - additive_mask: Tensor | None = None - if len(node.inputs) > 3: - mask_tensor = _ensure_fp16_tensor(g, _tensor(env, node.inputs[3])) - if bool(node.attrs.get("additive_mask", False)): - additive_mask = mask_tensor - else: - additive_mask = _lower_binary_op(g, mask_tensor, 1.0, "subtract") - additive_mask = g.scalar_multiply(_ensure_scalar_math_tensor(g, additive_mask), 1.0e4) - else: - query_shape = tuple(int(dim) for dim in query.shape) - key_shape = tuple(int(dim) for dim in key.shape) - if len(query_shape) >= 4 and len(key_shape) >= 4: - query_seq = int(query_shape[-2]) - key_seq = int(key_shape[-2]) - additive_mask = _materialize_gemma4_attention_mask( - g, - query_seq=query_seq, - key_seq=key_seq, - is_causal=bool(node.attrs.get("is_causal", False)), - window_size=int(node.attrs.get("window_size", 0)), - dtype=torch.float16, - ) - - if additive_mask is not None: - scores = _lower_binary_op(g, scores, additive_mask, "add") - - probs = g.softmax(scores, axis=-1) - probs = _ensure_tensor_dtype(g, probs, value.dtype) - out = _lower_static_batched_matmul(g, probs, value, output_dtype=value.dtype) - if out is None: - raise NotImplementedError("Gemma4 decoder attention output requires static batched matmul support") - - output_value = ir.values.get(node.outputs[0]) - output_dtype = _map_ir_dtype(output_value.dtype) if output_value is not None and output_value.dtype is not None else out.dtype - if out.dtype != output_dtype: - out = g.precision_cast(out, output_dtype) - return out - - -def _materialize_gemma4_attention_mask( - g: Graph, - *, - query_seq: int, - key_seq: int, - is_causal: bool, - window_size: int, - dtype: torch.dtype, -) -> Tensor | None: - if query_seq <= 0 or key_seq <= 0: - return None - if not is_causal and window_size <= 0: - return None - - q_index = np.arange(query_seq, dtype=np.int32)[:, None] - k_index = np.arange(key_seq, dtype=np.int32)[None, :] - allowed = np.ones((query_seq, key_seq), dtype=np.bool_) - - if is_causal: - allowed &= k_index <= q_index - if window_size > 0 and window_size < key_seq: - allowed &= k_index >= (q_index - (window_size - 1)) - - mask = np.where(allowed, 0.0, -1.0e4).astype(np.float16) - tensor = torch.from_numpy(mask.reshape(1, 1, query_seq, key_seq)).to(dtype=dtype) - return _materialize_constant_tensor(g, tensor) - - -def _legalize_matmul_inputs( - g: Graph, - lhs: Tensor, - rhs: Tensor, - node: IRNode, -) -> tuple[Tensor, Tensor, tuple[int, ...]] | None: - lhs_shape = tuple(lhs.shape) - rhs_shape = tuple(rhs.shape) - if len(lhs_shape) <= 2 and len(rhs_shape) <= 2: - return None - - if len(lhs_shape) > 2 and len(rhs_shape) == 2 and lhs_shape[-1] == rhs_shape[0]: - leading = 1 - for dim in lhs_shape[:-1]: - leading *= int(dim) - lhs_2d = g.reshape(lhs, (leading, int(lhs_shape[-1]))) - output_shape = lhs_shape[:-1] + (int(rhs_shape[1]),) - return lhs_2d, rhs, output_shape - - # Cactus matmul is 2D-only. Legalize the narrow case where both operands have - # only singleton leading dims, e.g. rotary helper matmuls: - # (1, M, K) @ (1, K, N) -> reshape to (M, K) @ (K, N) -> reshape back. - if any(dim != 1 for dim in lhs_shape[:-2]): - return None - if any(dim != 1 for dim in rhs_shape[:-2]): - return None - - lhs_2d_shape = lhs_shape[-2:] - rhs_2d_shape = rhs_shape[-2:] - if lhs_2d_shape[-1] != rhs_2d_shape[0]: - return None - - output_shape = node.meta.get("shape") - if not isinstance(output_shape, tuple): - output_shape = lhs_shape[:-2] + (lhs_2d_shape[0], rhs_2d_shape[1]) - output_shape = tuple(int(v) for v in output_shape) - - lhs_2d = g.reshape(lhs, lhs_2d_shape) - rhs_2d = g.reshape(rhs, rhs_2d_shape) - return lhs_2d, rhs_2d, output_shape - - -def _resolve_expand_shape(input_shape: tuple[int, ...], requested_shape: tuple[int, ...]) -> tuple[int, ...]: - if len(requested_shape) < len(input_shape): - raise NotImplementedError(f"expand cannot reduce rank: {input_shape} -> {requested_shape}") - - padded_input = (1,) * (len(requested_shape) - len(input_shape)) + input_shape - resolved: list[int] = [] - for in_dim, req_dim in zip(padded_input, requested_shape): - if req_dim == -1: - resolved.append(in_dim) - continue - if req_dim < -1: - raise NotImplementedError(f"invalid expand dimension: {req_dim}") - resolved.append(int(req_dim)) - return tuple(resolved) - - -def _lower_repeat(g: Graph, x: Tensor, repeats: tuple[int, ...]) -> Tensor: - if not repeats: - return x - - input_shape = tuple(int(dim) for dim in x.shape) - if len(repeats) < len(input_shape): - raise NotImplementedError(f"repeat cannot reduce rank: {input_shape} with repeats={repeats}") - - if len(repeats) > len(input_shape): - padded_shape = (1,) * (len(repeats) - len(input_shape)) + input_shape - x = g.reshape(x, padded_shape) - - result = x - for axis, factor in enumerate(repeats): - if factor < 0: - raise NotImplementedError(f"repeat does not support negative factors: {repeats}") - if factor == 0: - raise NotImplementedError(f"repeat does not support zero factors: {repeats}") - if factor == 1: - continue - result = g.cat([result] * factor, axis=axis) - return result - - -def _resolve_reshape_shape(input_shape: tuple[int, ...], requested_shape: tuple[int, ...]) -> tuple[int, ...]: - resolved = [int(v) for v in requested_shape] - unknown_indices = [idx for idx, dim in enumerate(resolved) if dim == -1] - if len(unknown_indices) > 1: - raise NotImplementedError(f"reshape with multiple inferred dimensions is unsupported: {requested_shape}") - if not unknown_indices: - return tuple(resolved) - - input_elements = 1 - for dim in input_shape: - input_elements *= int(dim) - - known_elements = 1 - for dim in resolved: - if dim != -1: - known_elements *= int(dim) - - if known_elements == 0 or input_elements % known_elements != 0: - raise NotImplementedError(f"cannot infer reshape target {requested_shape} from input shape {input_shape}") - - resolved[unknown_indices[0]] = input_elements // known_elements - return tuple(resolved) - - -def _match_gqa_expand_alias(ir: IRGraph, expand_node: IRNode) -> str | None: - target_shape = tuple(int(v) for v in expand_node.attrs["shape"]) - if len(target_shape) != 5: - return None - - current_value_id = expand_node.inputs[0] - base_value_id: str | None = None - - while True: - producer_id = ir.values[current_value_id].producer - if producer_id is None: - return None - producer = ir.nodes[producer_id] - if producer.op == "slice": - axis = int(producer.attrs.get("axis", -1)) - input_shape = ir.values[producer.inputs[0]].shape - if input_shape is None: - return None - normalized_axis = _normalize_dim(axis, len(input_shape)) - start = int(producer.attrs.get("start", 0)) - end = int(producer.attrs.get("end", input_shape[normalized_axis])) - if start != 0 or end < input_shape[normalized_axis]: - return None - current_value_id = producer.inputs[0] - continue - if producer.op == "unsqueeze": - base_value_id = producer.inputs[0] - base_shape = ir.values[base_value_id].shape - if base_shape is None: - return None - unsqueezed_dim = _normalize_dim(int(producer.attrs.get("dim", 0)), len(base_shape) + 1) - if unsqueezed_dim != 2: - return None - break - if producer.op in {"reshape", "view"}: - base_value_id = producer.inputs[0] - base_shape = ir.values[base_value_id].shape - view_shape = ir.values[current_value_id].shape - if base_shape is None or view_shape is None: - return None - if len(base_shape) != 4 or len(view_shape) != 5: - return None - if tuple(int(v) for v in view_shape) != ( - int(base_shape[0]), - int(base_shape[1]), - 1, - int(base_shape[2]), - int(base_shape[3]), - ): - return None - break - return None - - if base_value_id is None: - return None - base_shape = ir.values[base_value_id].shape - if base_shape is None or len(base_shape) != 4: - return None - - expected_target = ( - int(base_shape[0]), - int(base_shape[1]), - int(target_shape[2]), - int(base_shape[2]), - int(base_shape[3]), - ) - if target_shape != expected_target: - return None - if int(target_shape[2]) <= 1: - return None - - return base_value_id - - -def _map_ir_dtype(dtype: str) -> int: - if dtype == "bf16": - return Graph.FP32 - if dtype == "fp16": - return Graph.FP16 - if dtype in ("fp32", "fp64"): - return Graph.FP32 - if dtype == "int8": - return Graph.INT8 - if dtype in ("int32", "int64", "bool"): - return Graph.FP32 - - raise NotImplementedError(f"unsupported IR dtype: {dtype}") - - -def _map_ir_or_torch_dtype(dtype: Any) -> int: - if dtype is None: - raise NotImplementedError("missing dtype for precision_cast") - if isinstance(dtype, str): - if dtype.startswith("torch."): - return _map_torch_dtype(getattr(torch, dtype.split(".", 1)[1])) - return _map_ir_dtype(dtype) - return _map_torch_dtype(dtype) - - -def _map_torch_dtype(dtype: Any) -> int: - if dtype == torch.bfloat16: - return Graph.FP32 - if dtype == torch.float16: - return Graph.FP16 - if dtype == torch.float32: - return Graph.FP32 - if dtype == torch.float64: - return Graph.FP32 - if dtype == torch.int8: - return Graph.INT8 - if dtype in (torch.int16, torch.int32, torch.int64, torch.bool): - return Graph.FP32 - raise NotImplementedError(f"unsupported torch dtype: {dtype}") - - -def _materialize_constant_tensor(g: Graph, tensor_value: torch.Tensor) -> Tensor: - graph_dtype = _map_torch_dtype(tensor_value.dtype) - materialized = tensor_value.detach().cpu() - if graph_dtype == Graph.FP32 and materialized.dtype not in (torch.float32, torch.float64): - materialized = materialized.to(torch.float32) - elif graph_dtype == Graph.FP16 and materialized.dtype != torch.float16: - materialized = materialized.to(torch.float16) - elif graph_dtype == Graph.INT8 and materialized.dtype != torch.int8: - materialized = materialized.to(torch.int8) - - tensor = g.input(shape=tuple(materialized.shape), dtype=graph_dtype) - g.set_input(tensor, materialized, dtype=graph_dtype) - materialized_constants = getattr(g, "_transpile_materialized_constants", None) - if isinstance(materialized_constants, list): - materialized_constants.append(tensor) - return tensor - - -def _materialize_constant_torch_dtype(dtype: Any) -> torch.dtype: - if dtype is None: - return torch.float32 - if isinstance(dtype, str): - if dtype.startswith("torch."): - return getattr(torch, dtype.split(".", 1)[1]) - if dtype in ("fp16", "bf16"): - return torch.float16 - if dtype in ("fp32", "fp64", "int16", "int32", "int64", "bool"): - return torch.float32 - if dtype == "int8": - return torch.int8 - if isinstance(dtype, torch.dtype): - return dtype - raise NotImplementedError(f"unsupported dtype for constant materialization: {dtype}") - - -def _torch_dtype_for_graph_dtype(dtype: int) -> torch.dtype: - if dtype == Graph.FP16: - return torch.float16 - if dtype == Graph.FP32: - return torch.float32 - if dtype == Graph.INT8: - return torch.int8 - raise NotImplementedError(f"unsupported graph dtype for constant materialization: {dtype}") diff --git a/python/cactus/transpile/media_limits.py b/python/cactus/transpile/media_limits.py deleted file mode 100644 index 8b7910cf9..000000000 --- a/python/cactus/transpile/media_limits.py +++ /dev/null @@ -1,35 +0,0 @@ -from __future__ import annotations - -import os - -_DEFAULT_IMAGE_SIZE = 256 - - -def static_image_size() -> int: - raw = os.environ.get("CACTUS_TRANSPILER_IMAGE_SIZE", str(_DEFAULT_IMAGE_SIZE)) - try: - return max(1, int(raw)) - except (TypeError, ValueError): - return _DEFAULT_IMAGE_SIZE - - -def resize_static_image(image: object) -> object: - size = static_image_size() - target = (size, size) - if getattr(image, "size", None) == target or not hasattr(image, "resize"): - return image - try: - from PIL import Image # type: ignore - from PIL import ImageOps # type: ignore - - resample = Image.Resampling.BILINEAR - except AttributeError: # pragma: no cover - resample = Image.BILINEAR # type: ignore[name-defined] - from PIL import ImageOps # type: ignore - except Exception: - return image.resize(target) - - try: - return ImageOps.pad(image, target, method=resample, color=(255, 255, 255)) - except Exception: - return image.resize(target, resample=resample) diff --git a/python/cactus/transpile/model_adapters.py b/python/cactus/transpile/model_adapters.py deleted file mode 100644 index 57e07fe4c..000000000 --- a/python/cactus/transpile/model_adapters.py +++ /dev/null @@ -1,6326 +0,0 @@ -from __future__ import annotations - -from contextlib import contextmanager -from dataclasses import dataclass -import inspect -import math -import struct -import os -from pathlib import Path -from typing import Callable - -import numpy as np -import torch -import torch.nn.functional as F - -from cactus.transpile.component_pipeline import ComponentModuleSpec -from cactus.convert.cactus_adapters.tensor_io import CACTUS_MAGIC -from cactus.convert.cactus_adapters.tensor_io import align_offset -from cactus.convert.model_adapters.naming import gemma4_scale_factor - - -_GEMMA4_SAFE_TEXT_MLP_PRODUCT_SCALE = 1.0 / 64.0 -_UNSET = object() -try: - from transformers.models.gemma4.modeling_gemma4 import create_bidirectional_mask as _GEMMA4_CREATE_BIDIRECTIONAL_MASK # type: ignore -except Exception: - try: - from transformers.masking_utils import create_bidirectional_mask as _GEMMA4_CREATE_BIDIRECTIONAL_MASK # type: ignore - except Exception: - _GEMMA4_CREATE_BIDIRECTIONAL_MASK = None - - -class UnsupportedComponentsError(RuntimeError): - pass - - -@dataclass -class CanonicalizedModel: - module: torch.nn.Module - task: str - family: str - input_names: tuple[str, ...] = () - - -@dataclass(frozen=True) -class _Gemma4NativeMergeSegment: - kind: str - input_start: int - length: int - feature_start: int = 0 - - -@dataclass(frozen=True) -class _Gemma4NativeMergePlan: - segments: tuple[_Gemma4NativeMergeSegment, ...] - pli_token_ids: tuple[int, ...] - - -def _model_name_or_path(model: torch.nn.Module) -> str: - value = getattr(model, "name_or_path", None) - if isinstance(value, str) and value: - return value - config = getattr(model, "config", None) - value = getattr(config, "_name_or_path", None) - if isinstance(value, str) and value: - return value - return "" - - -def _transpile_graph_meta(model: torch.nn.Module, *, adapter_family: str, adapter_type: str, input_names: tuple[str, ...]) -> dict[str, object]: - meta: dict[str, object] = { - "adapter_family": adapter_family, - "adapter_type": adapter_type, - "model_name_or_path": _model_name_or_path(model), - "input_names": input_names, - } - # Rope-table precompute falls back to this on components that don't reserve a KV cache. - for candidate in _config_candidates_for_context_length(model): - for key in ("max_position_embeddings", "context_length", "model_max_length", "n_positions"): - value = _config_int_value(candidate, key) - if value is not None: - meta["max_position_embeddings"] = value - return meta - return meta - - -def _extract_tensor_output(output: object, *, preferred_field: str | None = None) -> torch.Tensor: - if preferred_field is not None: - value = getattr(output, preferred_field, None) - if isinstance(value, torch.Tensor): - return value - - if isinstance(output, torch.Tensor): - return output - if isinstance(output, tuple) and output and isinstance(output[0], torch.Tensor): - return output[0] - - for field_name in ("last_hidden_state", "logits"): - value = getattr(output, field_name, None) - if isinstance(value, torch.Tensor): - return value - - raise TypeError(f"could not extract tensor output from {type(output).__name__}") - - -def _gemma4_get_placeholder_masks( - get_placeholder_mask: Callable[..., object], - *, - token_type_ids: torch.Tensor | None, - input_ids: torch.Tensor, - inputs_embeds: torch.Tensor, -) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: - kwargs: dict[str, torch.Tensor] = { - "input_ids": input_ids, - "inputs_embeds": inputs_embeds, - } - try: - signature = inspect.signature(get_placeholder_mask) - except (TypeError, ValueError): - signature = None - if token_type_ids is not None and signature is not None: - parameters = signature.parameters - accepts_kwargs = any( - parameter.kind == inspect.Parameter.VAR_KEYWORD - for parameter in parameters.values() - ) - if accepts_kwargs or "token_type_ids" in parameters: - kwargs["token_type_ids"] = token_type_ids - - accepts_token_type_ids = "token_type_ids" in kwargs - result = get_placeholder_mask(**kwargs) - if not isinstance(result, tuple) or len(result) != 3: - raise TypeError( - "Gemma4 get_placeholder_mask must return " - "three placeholder masks" - ) - first_mask, second_mask, third_mask = result - if not ( - isinstance(first_mask, torch.Tensor) - and isinstance(second_mask, torch.Tensor) - and isinstance(third_mask, torch.Tensor) - ): - raise TypeError("Gemma4 get_placeholder_mask returned non-tensor masks") - if accepts_token_type_ids: - text_mask, image_mask, audio_mask = first_mask, second_mask, third_mask - else: - image_mask, video_mask, audio_mask = first_mask, second_mask, third_mask - text_mask = ~(image_mask | video_mask | audio_mask) - return text_mask, image_mask, audio_mask - - -def _module_or_config_attr(module: object, name: str, default: object | None = None) -> object: - value = getattr(module, name, None) - if value is not None: - return value - config = getattr(module, "config", None) - value = getattr(config, name, None) - if value is not None: - return value - return default - - -def _gemma4_get_per_layer_inputs( - backbone: torch.nn.Module, - input_ids: torch.Tensor, - inputs_embeds: torch.Tensor, -) -> torch.Tensor: - get_per_layer_inputs = getattr(backbone, "get_per_layer_inputs") - try: - signature = inspect.signature(get_per_layer_inputs) - except (TypeError, ValueError): - signature = None - if signature is not None and "inputs_embeds" not in signature.parameters: - return get_per_layer_inputs(input_ids) - return get_per_layer_inputs(input_ids, inputs_embeds) - - -def _select_last_active_token(hidden_or_logits: torch.Tensor, attention_mask: torch.Tensor | None) -> torch.Tensor: - if attention_mask is None or hidden_or_logits.ndim < 3: - return hidden_or_logits[:, -1:, ...] - - if attention_mask.ndim != 2: - raise ValueError(f"expected 2D attention mask, got shape {tuple(attention_mask.shape)}") - if attention_mask.shape[1] != hidden_or_logits.shape[1]: - raise ValueError( - "attention mask / hidden sequence length mismatch: " - f"{tuple(attention_mask.shape)} vs {tuple(hidden_or_logits.shape)}" - ) - - trailing_zero = attention_mask[:, :1] - attention_mask[:, :1] - shifted_mask = torch.cat((attention_mask[:, 1:], trailing_zero), dim=1) - last_active_mask = torch.logical_and(attention_mask != 0, shifted_mask == 0) - expanded_mask = last_active_mask.to(dtype=hidden_or_logits.dtype) - for _ in range(hidden_or_logits.ndim - 2): - expanded_mask = expanded_mask.unsqueeze(-1) - return (hidden_or_logits * expanded_mask).sum(dim=1, keepdim=True) - - -def _select_last_non_pad_token( - hidden_or_logits: torch.Tensor, - input_ids: torch.Tensor | None, - *, - pad_token_id: int | None, -) -> torch.Tensor: - if hidden_or_logits.ndim < 3: - return hidden_or_logits - if not isinstance(input_ids, torch.Tensor) or input_ids.ndim != 2: - return hidden_or_logits[:, -1:, ...] - if input_ids.shape[1] != hidden_or_logits.shape[1]: - raise ValueError( - "input id / hidden sequence length mismatch: " - f"{tuple(input_ids.shape)} vs {tuple(hidden_or_logits.shape)}" - ) - if pad_token_id is None: - return hidden_or_logits[:, -1:, ...] - attention_mask = (input_ids != int(pad_token_id)).to(dtype=torch.int64) - return _select_last_active_token(hidden_or_logits, attention_mask) - - -def _tile_to_length(tensor: torch.Tensor, length: int) -> torch.Tensor: - reps = -(-length // int(tensor.shape[1])) - return torch.cat([tensor] * reps, dim=1)[:, :length].contiguous() - - -def _resolve_model_pad_token_id(model: torch.nn.Module) -> int | None: - config = getattr(model, "config", None) - for attr_name in ("pad_token_id", "eos_token_id", "bos_token_id"): - value = getattr(config, attr_name, None) - if value is not None: - return int(value) - generation_config = getattr(model, "generation_config", None) - for attr_name in ("pad_token_id", "eos_token_id", "bos_token_id"): - value = getattr(generation_config, attr_name, None) - if value is not None: - return int(value) - return None - - -def _filter_supported_kwargs(module: torch.nn.Module, kwargs: dict[str, object]) -> dict[str, object]: - try: - signature = inspect.signature(module.forward) - except (TypeError, ValueError): - return kwargs - accepted = signature.parameters - supports_var_kwargs = any( - parameter.kind == inspect.Parameter.VAR_KEYWORD - for parameter in accepted.values() - ) - if supports_var_kwargs: - return kwargs - return {name: value for name, value in kwargs.items() if name in accepted} - - -def _module_floating_dtype(module: torch.nn.Module) -> torch.dtype | None: - for parameter in module.parameters(): - if parameter.is_floating_point(): - return parameter.dtype - for buffer in module.buffers(): - if buffer.is_floating_point(): - return buffer.dtype - return None - - -def _module_device(module: torch.nn.Module) -> torch.device | None: - for parameter in module.parameters(): - return parameter.device - for buffer in module.buffers(): - return buffer.device - return None - - -def _unique_modules(*candidates: object) -> tuple[torch.nn.Module, ...]: - modules: list[torch.nn.Module] = [] - seen: set[int] = set() - for candidate in candidates: - if not isinstance(candidate, torch.nn.Module): - continue - candidate_id = id(candidate) - if candidate_id in seen: - continue - seen.add(candidate_id) - modules.append(candidate) - return tuple(modules) - - -def _gemma4_vision_modules(multimodal_backbone: torch.nn.Module) -> tuple[torch.nn.Module, ...]: - return _unique_modules( - getattr(multimodal_backbone, "vision_tower", None), - getattr(multimodal_backbone, "embed_vision", None), - ) - - -def _gemma4_audio_modules(multimodal_backbone: torch.nn.Module) -> tuple[torch.nn.Module, ...]: - return _unique_modules( - getattr(multimodal_backbone, "audio_tower", None), - getattr(multimodal_backbone, "embed_audio", None), - ) - - -def _gemma4_text_config(multimodal_backbone: torch.nn.Module) -> object: - config = getattr(multimodal_backbone, "config", None) - get_text_config = getattr(config, "get_text_config", None) - if callable(get_text_config): - try: - return get_text_config() - except Exception: - pass - return config - - -def _parse_cache_context_length(value: str | int | None) -> int | None: - if value is None: - return None - if isinstance(value, str): - normalized = value.strip().lower() - if normalized in {"", "auto"}: - return None - parsed = int(normalized) - else: - parsed = int(value) - if parsed <= 0: - raise ValueError("--cache-context-length must be a positive integer or auto") - return parsed - - -def _config_int_value(config: object, key: str) -> int | None: - if config is None: - return None - if isinstance(config, dict): - raw_value = config.get(key) - else: - raw_value = getattr(config, key, None) - if raw_value is None: - return None - parsed = int(raw_value) - return parsed if parsed > 0 else None - - -def _config_candidates_for_context_length(model: torch.nn.Module) -> tuple[object, ...]: - candidates: list[object] = [] - seen: set[int] = set() - - def _add(candidate: object) -> None: - if candidate is None: - return - candidate_id = id(candidate) - if candidate_id in seen: - return - seen.add(candidate_id) - candidates.append(candidate) - - for owner in ( - model, - getattr(model, "model", None), - getattr(getattr(model, "model", None), "language_model", None), - ): - config = getattr(owner, "config", None) - _add(config) - get_text_config = getattr(config, "get_text_config", None) - if callable(get_text_config): - _add(get_text_config()) - for attr in ("text_config", "llm_config", "language_config", "decoder_config"): - _add(getattr(config, attr, None)) - - return tuple(candidates) - - -def _cache_context_length( - model: torch.nn.Module, - *, - input_seq_len: int, - cache_context_length: str | int | None, - fallback_extra_tokens: int, -) -> int: - explicit_length = _parse_cache_context_length(cache_context_length) - if explicit_length is not None: - return explicit_length - - for candidate in _config_candidates_for_context_length(model): - for key in ("max_position_embeddings", "context_length", "model_max_length", "n_positions"): - value = _config_int_value(candidate, key) - if value is not None: - return value - return input_seq_len + int(fallback_extra_tokens) - - -def _max_cache_seq_len( - model: torch.nn.Module, - input_ids: torch.Tensor, - cache_context_length: str | int | None, - *, - fallback_extra_tokens: int, -) -> int: - return max(1024, _cache_context_length( - model, - input_seq_len=int(input_ids.shape[1]), - cache_context_length=cache_context_length, - fallback_extra_tokens=fallback_extra_tokens, - )) - - -def _gemma4_special_token_ids(multimodal_backbone: torch.nn.Module) -> tuple[int, int, int]: - config = getattr(multimodal_backbone, "config", None) - text_config = _gemma4_text_config(multimodal_backbone) - image_token_id = int(getattr(config, "image_token_id", 0) or 0) - audio_token_id = int(getattr(config, "audio_token_id", 0) or 0) - pad_token_id = getattr(text_config, "pad_token_id", None) - if pad_token_id is None: - pad_token_id = getattr(config, "pad_token_id", 0) - return image_token_id, audio_token_id, int(pad_token_id or 0) - - -@contextmanager -def _temporary_cpu_float32_modules(modules: tuple[torch.nn.Module, ...]): - promoted: list[tuple[torch.nn.Module, torch.dtype]] = [] - for module in modules: - device = _module_device(module) - dtype = _module_floating_dtype(module) - if device is None or device.type != "cpu" or dtype is None or dtype == torch.float32: - continue - promoted.append((module, dtype)) - module.to(dtype=torch.float32) - try: - yield - finally: - for module, dtype in reversed(promoted): - module.to(dtype=dtype) - - -def _torch_is_compiling() -> bool: - dynamo = getattr(torch, "_dynamo", None) - is_compiling = getattr(dynamo, "is_compiling", None) - if callable(is_compiling): - try: - return bool(is_compiling()) - except Exception: - return False - return False - - -def _gemma4_cpu_safe_text_mlp_enabled(layer: torch.nn.Module, hidden_states: torch.Tensor) -> bool: - if hidden_states.device.type != "cpu" or hidden_states.dtype != torch.float16: - return False - return _gemma4_text_mlp_manual_available(layer) - - -def _gemma4_text_mlp_manual_available(layer: torch.nn.Module) -> bool: - if bool(getattr(layer, "enable_moe_block", False)): - return False - mlp = getattr(layer, "mlp", None) - if mlp is None: - return False - for attr_name in ("gate_proj", "up_proj", "down_proj", "act_fn"): - if not hasattr(mlp, attr_name): - return False - return True - - -def _gemma4_cpu_safe_text_mlp_forward(mlp: torch.nn.Module, hidden_states: torch.Tensor) -> torch.Tensor: - gate = F.linear( - hidden_states, - mlp.gate_proj.weight, - mlp.gate_proj.bias, - ) - up = F.linear( - hidden_states, - mlp.up_proj.weight, - mlp.up_proj.bias, - ) - # Gemma4 text MLPs can overflow in FP16 at the gated product. - # Scaling one branch here keeps the product and down-projection finite. - # The following post-feedforward RMSNorm is scale-invariant, so this - # constant factor cancels back out in the normalized output. - activated = mlp.act_fn(gate) * _GEMMA4_SAFE_TEXT_MLP_PRODUCT_SCALE - return F.linear( - activated * up, - mlp.down_proj.weight, - mlp.down_proj.bias, - ) - - -def _gemma4_text_attention_forward( - attn: torch.nn.Module, - hidden_states: torch.Tensor, - *, - position_embeddings: torch.Tensor | tuple[torch.Tensor, torch.Tensor] | None, - attention_mask: torch.Tensor | None, - position_ids: torch.LongTensor | None, - past_key_values: object | None, - use_cache: bool, - shared_kv_states: dict[int, tuple[torch.Tensor, torch.Tensor]] | None, -) -> torch.Tensor: - if shared_kv_states is None: - attn_out, _ = attn( - hidden_states=hidden_states, - position_embeddings=position_embeddings, - attention_mask=attention_mask, - position_ids=position_ids, - past_key_values=past_key_values, - use_cache=use_cache, - ) - return attn_out - - from transformers.models.gemma4.modeling_gemma4 import ALL_ATTENTION_FUNCTIONS # type: ignore - from transformers.models.gemma4.modeling_gemma4 import apply_rotary_pos_emb # type: ignore - from transformers.models.gemma4.modeling_gemma4 import eager_attention_forward # type: ignore - - input_shape = hidden_states.shape[:-1] - hidden_shape = (*input_shape, -1, attn.head_dim) - if position_embeddings is None: - raise ValueError("Gemma4 text attention requires position embeddings") - cos, sin = position_embeddings - - query_states = attn.q_proj(hidden_states).view(hidden_shape) - query_states = attn.q_norm(query_states) - query_states = apply_rotary_pos_emb(query_states, cos, sin, unsqueeze_dim=2) - query_states = query_states.transpose(1, 2) - - kv_shared_key = getattr(attn, "layer_type", getattr(attn, "kv_shared_layer_index", None)) - if bool(getattr(attn, "is_kv_shared_layer", False)) and kv_shared_key in shared_kv_states: - key_states, value_states = shared_kv_states[kv_shared_key] - key_states = key_states.to(query_states.device) - value_states = value_states.to(query_states.device) - else: - key_states = attn.k_proj(hidden_states).view(hidden_shape) - value_states = attn.v_proj(hidden_states).view(hidden_shape) if attn.v_proj is not None else key_states - - key_states = attn.k_norm(key_states) - key_states = apply_rotary_pos_emb(key_states, cos, sin, unsqueeze_dim=2) - key_states = key_states.transpose(1, 2) - - value_states = attn.v_norm(value_states) - value_states = value_states.transpose(1, 2) - - layer_idx = getattr(attn, "layer_idx", None) - if bool(getattr(attn, "store_full_length_kv", False)): - shared_kv_states[kv_shared_key] = (key_states, value_states) - elif layer_idx is not None and not bool(getattr(attn, "is_kv_shared_layer", False)): - shared_kv_states[int(layer_idx)] = (key_states, value_states) - - attention_interface = eager_attention_forward - if getattr(attn.config, "_attn_implementation", "eager") != "eager": - attention_interface = ALL_ATTENTION_FUNCTIONS[attn.config._attn_implementation] - - attn_output, _ = attention_interface( - attn, - query_states, - key_states, - value_states, - attention_mask, - dropout=attn.attention_dropout if attn.training else 0.0, - scaling=attn.scaling, - sliding_window=attn.sliding_window, - ) - attn_output = attn_output.reshape(*input_shape, -1).contiguous() - return attn.o_proj(attn_output) - - -def _gemma4_text_decoder_layer_forward( - layer: torch.nn.Module, - hidden_states: torch.Tensor, - *, - position_embeddings: torch.Tensor | tuple[torch.Tensor, torch.Tensor] | None = None, - per_layer_input: torch.Tensor | None = None, - attention_mask: torch.Tensor | None = None, - position_ids: torch.LongTensor | None = None, - past_key_values: object | None = None, - use_cache: bool = False, - shared_kv_states: dict[int, tuple[torch.Tensor, torch.Tensor]] | None = None, -) -> torch.Tensor: - manual_attention_required = shared_kv_states is not None - if not manual_attention_required and not _gemma4_cpu_safe_text_mlp_enabled(layer, hidden_states): - return layer( - hidden_states, - position_embeddings=position_embeddings, - per_layer_input=per_layer_input, - attention_mask=attention_mask, - position_ids=position_ids, - past_key_values=past_key_values, - use_cache=use_cache, - ) - if not _gemma4_text_mlp_manual_available(layer): - raise RuntimeError("Gemma4 KV sharing requires a supported text MLP structure") - - residual = hidden_states - - hidden_states = layer.input_layernorm(hidden_states) - hidden_states = _gemma4_text_attention_forward( - layer.self_attn, - hidden_states, - position_embeddings=position_embeddings, - attention_mask=attention_mask, - position_ids=position_ids, - past_key_values=past_key_values, - use_cache=use_cache, - shared_kv_states=shared_kv_states, - ) - hidden_states = layer.post_attention_layernorm(hidden_states) - hidden_states = residual + hidden_states - - residual = hidden_states - hidden_states = layer.pre_feedforward_layernorm(hidden_states) - mlp_out = _gemma4_cpu_safe_text_mlp_forward(layer.mlp, hidden_states) - hidden_states = layer.post_feedforward_layernorm(mlp_out) - hidden_states = residual.float() + hidden_states - hidden_states = hidden_states.to(dtype=residual.dtype) - - if getattr(layer, "hidden_size_per_layer_input", 0): - residual = hidden_states - hidden_states = layer.per_layer_input_gate(hidden_states) - hidden_states = layer.act_fn(hidden_states) - if per_layer_input is None: - raise ValueError("Gemma4 layer expected per_layer_input but none was provided") - hidden_states = torch.multiply(hidden_states, per_layer_input) - hidden_states = layer.per_layer_projection(hidden_states) - hidden_states = layer.post_per_layer_input_norm(hidden_states) - hidden_states = residual + hidden_states - - hidden_states = hidden_states * layer.layer_scalar - return hidden_states - - -def _gemma4_text_backbone_forward( - backbone: torch.nn.Module, - *, - inputs_embeds: torch.Tensor, - per_layer_inputs: torch.Tensor | None, - causal_mask_mapping: dict[str, torch.Tensor], - position_ids: torch.LongTensor, - layer_start: int = 0, - layer_end: int | None = None, - apply_norm: bool = True, - shared_kv_states: dict[int, tuple[torch.Tensor, torch.Tensor]] | None | object = _UNSET, - capture_layer_index: int | None = None, -) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]: - hidden_states = inputs_embeds - captured_hidden: torch.Tensor | None = None - layer_types = tuple(dict.fromkeys(getattr(backbone.config, "layer_types", ()))) - position_embeddings = { - layer_type: backbone.rotary_emb(hidden_states, position_ids, layer_type) - for layer_type in layer_types - } - if shared_kv_states is _UNSET: - num_shared_layers = int(getattr(backbone.config, "num_kv_shared_layers", 0) or 0) - shared_kv_states = {} if num_shared_layers > 0 else None - - config_layer_types = tuple(getattr(backbone.config, "layer_types", ())) - end = backbone.config.num_hidden_layers if layer_end is None else min(int(layer_end), int(backbone.config.num_hidden_layers)) - for layer_index in range(max(0, int(layer_start)), end): - decoder_layer = backbone.layers[layer_index] - layer_per_input = None - if per_layer_inputs is not None: - layer_per_input = per_layer_inputs[:, :, decoder_layer.layer_idx, :] - attention_type = getattr( - decoder_layer, - "attention_type", - config_layer_types[layer_index] if layer_index < len(config_layer_types) else "full_attention", - ) - hidden_states = _gemma4_text_decoder_layer_forward( - decoder_layer, - hidden_states, - per_layer_input=layer_per_input, - attention_mask=causal_mask_mapping[attention_type], - position_embeddings=position_embeddings[attention_type], - position_ids=position_ids, - past_key_values=None, - use_cache=False, - shared_kv_states=shared_kv_states, - ) - if capture_layer_index is not None and layer_index == int(capture_layer_index): - captured_hidden = hidden_states - - output_hidden = backbone.norm(hidden_states) if apply_norm else hidden_states - if capture_layer_index is None: - return output_hidden - if captured_hidden is None: - captured_hidden = hidden_states - return output_hidden, captured_hidden - - -def _gemma4_strip_audio_padding(audio_output: object) -> torch.Tensor: - audio_features = getattr(audio_output, "pooler_output", None) - if not isinstance(audio_features, torch.Tensor): - raise TypeError("Gemma4 audio output did not expose tensor pooler_output") - - audio_mask_from_encoder = getattr(audio_output, "attention_mask", None) - if isinstance(audio_mask_from_encoder, torch.Tensor): - return audio_features[audio_mask_from_encoder].unsqueeze(0) - - audio_mask_from_encoder = getattr(audio_output, "audio_mel_mask", None) - if not isinstance(audio_mask_from_encoder, torch.Tensor): - return audio_features - all_real_tokens: list[torch.Tensor] = [] - for encodings, padding_mask in zip(audio_features, audio_mask_from_encoder, strict=True): - all_real_tokens.append(encodings[~padding_mask]) - return torch.cat(all_real_tokens, dim=0).unsqueeze(0) - - -def _gemma4_rms_norm_no_scale(hidden_states: torch.Tensor, *, eps: float) -> torch.Tensor: - hidden_states_f32 = hidden_states.float() - mean_squared = hidden_states_f32.pow(2).mean(-1, keepdim=True) + eps - return hidden_states_f32 * torch.pow(mean_squared, -0.5) - - -def _gemma4_rms_norm(hidden_states: torch.Tensor, weight: torch.Tensor, *, eps: float) -> torch.Tensor: - normed = _gemma4_rms_norm_no_scale(hidden_states, eps=eps) - return normed * weight.to(device=normed.device, dtype=normed.dtype) - - -def _gemma4_layer_norm_eps(multimodal_backbone: torch.nn.Module) -> float: - config = getattr(multimodal_backbone, "config", None) - text_config = None - get_text_config = getattr(config, "get_text_config", None) - if callable(get_text_config): - try: - text_config = get_text_config() - except Exception: - text_config = None - for candidate in (text_config, config): - value = getattr(candidate, "rms_norm_eps", None) - if value is not None: - return float(value) - value = getattr(candidate, "layer_norm_eps", None) - if value is not None: - return float(value) - return 1e-6 - - -def _gemma4_load_cactus_fp_tensor(path: Path) -> torch.Tensor: - with path.open("rb") as handle: - header = handle.read(84) - if len(header) < 84 or header[:4] != CACTUS_MAGIC: - raise RuntimeError(f"Gemma4 Cactus tensor is missing a valid header: {path}") - - alignment = max(1, int(struct.unpack_from(" 0) - precision = int(struct.unpack_from(" 0 else 0 - data_offset = ( - align_offset(scales_offset + scales_bytes, alignment) - if scales_bytes > 0 - else aligned_header - ) - data_count = byte_size // np.dtype(dtype).itemsize - array = np.memmap(path, mode="r", dtype=dtype, offset=data_offset, shape=(data_count,)) - tensor = torch.from_numpy(np.array(array, copy=True)) - if shape: - tensor = tensor.reshape(shape) - return tensor - - -def _gemma4_load_vision_post_proj_norm(weights_dir: str | Path | None) -> torch.Tensor | None: - if weights_dir is None: - return None - path = Path(weights_dir).expanduser() / "embed_vision_post_proj_norm.weights" - if not path.exists(): - return None - tensor = _gemma4_load_cactus_fp_tensor(path) - if tensor.ndim != 1: - raise RuntimeError(f"Gemma4 vision post-projection norm must be 1D, got shape={tuple(tensor.shape)}") - return tensor.float() - - -def _gemma4_can_use_native_like_vision_features(multimodal_backbone: torch.nn.Module) -> bool: - vision_tower = getattr(multimodal_backbone, "vision_tower", None) - embed_vision = getattr(multimodal_backbone, "embed_vision", None) - return ( - isinstance(vision_tower, torch.nn.Module) - and isinstance(embed_vision, torch.nn.Module) - and hasattr(vision_tower, "patch_embedder") - and hasattr(vision_tower, "encoder") - and hasattr(embed_vision, "embedding_projection") - ) - - -def _gemma4_can_use_native_like_audio_features(multimodal_backbone: torch.nn.Module) -> bool: - audio_tower = getattr(multimodal_backbone, "audio_tower", None) - embed_audio = getattr(multimodal_backbone, "embed_audio", None) - layers = getattr(audio_tower, "layers", None) if isinstance(audio_tower, torch.nn.Module) else None - conformer = getattr(audio_tower, "conformer", None) if isinstance(audio_tower, torch.nn.Module) else None - return ( - isinstance(audio_tower, torch.nn.Module) - and isinstance(embed_audio, torch.nn.Module) - and hasattr(audio_tower, "subsample_conv_projection") - and (isinstance(layers, torch.nn.ModuleList) or isinstance(conformer, torch.nn.ModuleList)) - and hasattr(embed_audio, "embedding_projection") - ) - - -def _gemma4_pool_vision_hidden_native_like( - vision_tower: torch.nn.Module, - hidden_states: torch.Tensor, - pixel_position_ids: torch.Tensor, - *, - image_soft_token_counts: tuple[int, ...] | None = None, - image_pool_shapes: tuple[tuple[int, int, int], ...] | None = None, -) -> torch.Tensor: - output_length = int(_module_or_config_attr(vision_tower, "default_output_length", 280) or 280) - pooling_kernel_size = int(_module_or_config_attr(vision_tower, "pooling_kernel_size", 3) or 3) - padding_positions = (pixel_position_ids == -1).all(dim=-1) - pooled_batches: list[torch.Tensor] = [] - for row_idx in range(int(hidden_states.shape[0])): - hidden_row = hidden_states[row_idx] - position_row = pixel_position_ids[row_idx] - padding_row = padding_positions[row_idx] - if image_pool_shapes is not None and row_idx < len(image_pool_shapes): - grid_h, grid_w, pooled_count = image_pool_shapes[row_idx] - valid_patch_count = int(grid_h) * int(grid_w) - channels = int(hidden_row.shape[-1]) - valid_hidden = hidden_row[:valid_patch_count].float() - pooled = valid_hidden.reshape( - int(grid_h) // pooling_kernel_size, - pooling_kernel_size, - int(grid_w) // pooling_kernel_size, - pooling_kernel_size, - channels, - ).mean(dim=(1, 3)) - pooled_batches.append(pooled.reshape(int(pooled_count), channels)) - continue - - # Avoid boolean advanced indexing here. Gemma4 image padding is not - # guaranteed to be prefix-shaped, and the generic lowerer optimizes - # some masks as prefix slices. Zero-weighting padded patches preserves - # native pooling semantics while staying easy to lower. - clamped_positions = position_row.clamp(min=0) - max_x = clamped_positions[:, 0].max() + 1 - kernel_positions = torch.div(clamped_positions, pooling_kernel_size, rounding_mode="floor") - kernel_indices = kernel_positions[:, 0] + torch.div( - max_x, - pooling_kernel_size, - rounding_mode="floor", - ) * kernel_positions[:, 1] - valid_patch_weights = torch.logical_not(padding_row).float().unsqueeze(-1) - weights = ( - F.one_hot(kernel_indices.long(), output_length).float() - * valid_patch_weights - / float(pooling_kernel_size**2) - ) - pooled_full = weights.transpose(0, 1) @ hidden_row.float() - if image_soft_token_counts is not None and row_idx < len(image_soft_token_counts): - pooled_batches.append(pooled_full[: int(image_soft_token_counts[row_idx])]) - else: - valid_bins = torch.logical_not((weights == 0).all(dim=0)) - pooled_batches.append(pooled_full[valid_bins]) - pooled_hidden = pooled_batches[0] if len(pooled_batches) == 1 else torch.cat(pooled_batches, dim=0) - hidden_size = int(_module_or_config_attr(vision_tower, "hidden_size", pooled_hidden.shape[-1]) or pooled_hidden.shape[-1]) - pooled_hidden = pooled_hidden * math.sqrt(float(hidden_size)) - if bool(_module_or_config_attr(vision_tower, "standardize", False)): - std_bias = getattr(vision_tower, "std_bias", None) - std_scale = getattr(vision_tower, "std_scale", None) - if std_bias is not None and std_scale is not None: - pooled_hidden = (pooled_hidden - std_bias) * std_scale - return pooled_hidden - - -def _gemma4_vision_encoder_hidden_states( - vision_encoder: torch.nn.Module, - inputs_embeds: torch.Tensor, - attention_mask: torch.Tensor, - pixel_position_ids: torch.Tensor | None, -) -> torch.Tensor: - config = getattr(vision_encoder, "config", None) - hidden_states = inputs_embeds - rotary_emb = getattr(vision_encoder, "rotary_emb") - layers = getattr(vision_encoder, "layers") - num_layers = int(getattr(config, "num_hidden_layers", len(layers))) - layer_types = tuple(str(value) for value in (getattr(config, "layer_types", ()) or ())) - if layer_types and all(hasattr(rotary_emb, f"{layer_type}_inv_freq") for layer_type in layer_types): - # Newer Gemma4 builds use per-layer vision RoPE tables keyed by the - # layer attention type instead of a single vision RoPE table. - attention_mask_4d = (attention_mask.unsqueeze(-1) * attention_mask.unsqueeze(1)).unsqueeze(1) - position_embeddings_by_type = { - layer_type: rotary_emb(hidden_states, pixel_position_ids, layer_type) - for layer_type in layer_types - } - for decoder_layer in layers[:num_layers]: - layer_type = str(getattr(decoder_layer, "attention_type", layer_types[0])) - hidden_states = decoder_layer( - hidden_states, - attention_mask=attention_mask_4d, - position_embeddings=position_embeddings_by_type[layer_type], - position_ids=pixel_position_ids, - ) - return hidden_states - - if _GEMMA4_CREATE_BIDIRECTIONAL_MASK is None: - raise RuntimeError("Gemma4 bidirectional mask helper is unavailable in this transformers install") - attention_mask = _GEMMA4_CREATE_BIDIRECTIONAL_MASK( - config=config, - inputs_embeds=inputs_embeds, - attention_mask=attention_mask, - ) - position_embeddings = rotary_emb(hidden_states, pixel_position_ids) - for decoder_layer in layers[:num_layers]: - hidden_states = decoder_layer( - hidden_states, - attention_mask=attention_mask, - position_embeddings=position_embeddings, - position_ids=pixel_position_ids, - ) - return hidden_states - - -def _gemma4_convert_audio_mask_to_blocked_static( - mask_4d: torch.Tensor, - *, - chunk_size: int, - left_context: int, - right_context: int, -) -> torch.Tensor: - """Build Gemma4's blocked local audio mask without tensor index expand. - - Hugging Face's helper uses an expanded gather-index tensor. The v2 graph - runtime does not have a first-class expand op, so we express the same fixed - chunk extraction as static slices and cats. - """ - _, _, seq_len, _ = mask_4d.shape - seq_len_int = int(seq_len) - num_blocks = (seq_len_int + int(chunk_size) - 1) // int(chunk_size) - padded_seq_len = num_blocks * int(chunk_size) - pad_amount = padded_seq_len - seq_len_int - if pad_amount: - mask_4d = F.pad(mask_4d, (0, pad_amount, 0, pad_amount), value=False) - mask_4d = F.pad(mask_4d, (int(left_context), int(right_context)), value=False) - - context_size = int(chunk_size) + int(left_context) + int(right_context) - blocks: list[torch.Tensor] = [] - for block_idx in range(num_blocks): - start = block_idx * int(chunk_size) - block = mask_4d[ - :, - :, - start : start + int(chunk_size), - start : start + context_size, - ] - blocks.append(block.unsqueeze(2)) - return torch.cat(blocks, dim=2) - - -def _gemma4_compute_native_like_image_features( - multimodal_backbone: torch.nn.Module, - pixel_values: torch.Tensor, - pixel_position_ids: torch.Tensor | None, - *, - post_proj_norm_weight: torch.Tensor | None = None, - image_soft_token_counts: tuple[int, ...] | None = None, - image_pool_shapes: tuple[tuple[int, int, int], ...] | None = None, -) -> torch.Tensor: - if pixel_position_ids is None: - raise TypeError("Gemma4 native-like vision feature path requires pixel_position_ids") - vision_tower = getattr(multimodal_backbone, "vision_tower", None) - embed_vision = getattr(multimodal_backbone, "embed_vision", None) - if not isinstance(vision_tower, torch.nn.Module) or not isinstance(embed_vision, torch.nn.Module): - raise TypeError("Gemma4 multimodal backbone is missing native-like vision modules") - - padding_positions = (pixel_position_ids == -1).all(dim=-1) - vision_inputs = vision_tower.patch_embedder( - pixel_values, - pixel_position_ids, - padding_positions, - ) - vision_hidden = _gemma4_vision_encoder_hidden_states( - vision_tower.encoder, - vision_inputs, - ~padding_positions, - pixel_position_ids, - ) - pooled_hidden = _gemma4_pool_vision_hidden_native_like( - vision_tower, - vision_hidden, - pixel_position_ids, - image_soft_token_counts=image_soft_token_counts, - image_pool_shapes=image_pool_shapes, - ) - projection = getattr(embed_vision, "embedding_projection", None) - if not isinstance(projection, torch.nn.Linear): - raise TypeError("Gemma4 vision embedder is missing embedding_projection") - pre_projection_norm = getattr(embed_vision, "embedding_pre_projection_norm", None) - if isinstance(pre_projection_norm, torch.nn.Module): - pooled_hidden = pre_projection_norm(pooled_hidden) - pooled_hidden = pooled_hidden.to(dtype=projection.weight.dtype) - projected = F.linear( - pooled_hidden, - projection.weight, - projection.bias, - ) - if post_proj_norm_weight is not None and post_proj_norm_weight.numel() > 0: - return _gemma4_rms_norm( - projected, - post_proj_norm_weight, - eps=_gemma4_layer_norm_eps(multimodal_backbone), - ) - return projected - - -_GEMMA4_NPU_PROJ_WEIGHTS = ( - ("embed_audio", "embed_audio_proj.weights"), - ("embed_vision", "embed_vision_proj.weights"), -) - - -@contextmanager -def _gemma4_npu_projection_reparam(module: torch.nn.Module): - backbone = getattr(module, "multimodal_backbone", None) - scaled: list[tuple[torch.nn.Linear, float]] = [] - for attr, cactus_name in _GEMMA4_NPU_PROJ_WEIGHTS: - projection = getattr(getattr(backbone, attr, None), "embedding_projection", None) - factor = gemma4_scale_factor(cactus_name) - if isinstance(projection, torch.nn.Linear) and factor != 1.0: - scaled.append((projection, factor)) - for projection, factor in scaled: - projection.weight.data.mul_(factor) - try: - yield - finally: - for projection, factor in scaled: - projection.weight.data.div_(factor) - - -def _gemma4_compute_native_like_audio_features( - multimodal_backbone: torch.nn.Module, - input_features: torch.Tensor, - input_features_mask: torch.Tensor, -) -> torch.Tensor: - audio_tower = getattr(multimodal_backbone, "audio_tower", None) - embed_audio = getattr(multimodal_backbone, "embed_audio", None) - if not isinstance(audio_tower, torch.nn.Module) or not isinstance(embed_audio, torch.nn.Module): - raise TypeError("Gemma4 multimodal backbone is missing native-like audio modules") - - config = getattr(audio_tower, "config", None) - if config is None: - raise TypeError("Gemma4 audio tower is missing config") - hidden_states, output_mask = audio_tower.subsample_conv_projection(input_features, input_features_mask) - - if not hasattr(audio_tower, "rel_pos_enc"): - # Newer Gemma4 audio towers keep relative-position handling inside each - # conformer block. Follow that native path with fixed-shape masks. - chunk_size = int(getattr(config, "conf_attention_chunk_size", getattr(config, "attention_chunk_size", 12))) - right_context = int(getattr(config, "conf_attention_context_right", getattr(config, "attention_context_right", 0))) - left_context = int(getattr(config, "conf_attention_context_left", getattr(config, "attention_context_left", 13))) - max_past_horizon = max(0, left_context - 1) - upper_diagonal = max_past_horizon + right_context - context_size = chunk_size + max_past_horizon + right_context - lower_causal_mask = torch.tril( - torch.ones((context_size, chunk_size), dtype=torch.bool, device=hidden_states.device), - diagonal=0, - ).T - upper_causal_mask = torch.tril( - torch.ones((chunk_size, context_size), dtype=torch.bool, device=hidden_states.device), - diagonal=upper_diagonal, - ) - causal_valid_mask = ( - torch.ones((chunk_size, context_size), dtype=torch.bool, device=hidden_states.device) - * lower_causal_mask - * upper_causal_mask - ) - - layers = getattr(audio_tower, "conformer", None) - if not isinstance(layers, torch.nn.ModuleList): - raise TypeError("Gemma4 audio tower is missing conformer layers") - num_layers = int(getattr(config, "num_hidden_layers", len(layers))) - for encoder_layer in layers[:num_layers]: - hidden_states = encoder_layer(hidden_states, output_mask, causal_valid_mask) - - reduction_factor = int(getattr(config, "conf_reduction_factor", 1) or 1) - if reduction_factor > 1: - hidden_states = hidden_states[:, ::reduction_factor] - if output_mask is not None: - output_mask = output_mask[:, ::reduction_factor] - - output_proj = getattr(audio_tower, "output_proj", None) - audio_encodings = output_proj(hidden_states) if isinstance(output_proj, torch.nn.Linear) else hidden_states - else: - position_embeddings = audio_tower.rel_pos_enc(hidden_states) - seq_len = hidden_states.shape[1] - query_positions = torch.arange(seq_len, device=hidden_states.device)[:, None] - key_positions = torch.arange(seq_len, device=hidden_states.device)[None, :] - distance = query_positions - key_positions - left_context = int(getattr(config, "attention_context_left", 13)) - 1 - right_context = int(getattr(config, "attention_context_right", 0)) - local_mask = ((distance >= 0) & (distance < left_context)) | ((distance < 0) & ((-distance) < right_context)) - attention_mask = local_mask.unsqueeze(0).unsqueeze(0) - if output_mask is not None: - attention_mask = attention_mask & output_mask[:, None, None, :].to(dtype=torch.bool) - attention_mask = _gemma4_convert_audio_mask_to_blocked_static( - attention_mask, - chunk_size=int(getattr(config, "attention_chunk_size", 12)), - left_context=left_context, - right_context=right_context, - ) - - layers = getattr(audio_tower, "layers", None) - if not isinstance(layers, torch.nn.ModuleList): - layers = getattr(audio_tower, "conformer", None) - if not isinstance(layers, torch.nn.ModuleList): - raise TypeError("Gemma4 audio tower is missing layers/conformer") - num_layers = int(getattr(config, "num_hidden_layers", len(layers))) - for encoder_layer in layers[:num_layers]: - hidden_states = encoder_layer( - hidden_states, - attention_mask=attention_mask, - position_embeddings=position_embeddings, - ) - - output_proj = getattr(audio_tower, "output_proj", None) - if not isinstance(output_proj, torch.nn.Linear): - raise TypeError("Gemma4 audio tower is missing output_proj") - audio_encodings = output_proj(hidden_states) - - if not callable(getattr(embed_audio, "forward", None)): - raise TypeError("Gemma4 audio embedder is not callable") - projected = embed_audio(inputs_embeds=audio_encodings) - # The transpiled bundle is shape-specialized from representative media. - # Native Gemma4 audio preprocessing emits an unpadded feature tensor for - # that media, so the post-subsampling sequence is already the real token - # sequence. Returning the dense sequence avoids dynamic boolean indexing in - # torch.export and matches the prompt's static audio soft-token count. - return projected - - -def _gemma4_feature_token_count(features: torch.Tensor | None) -> int: - if features is None or features.numel() == 0: - return 0 - if features.ndim == 3: - return int(features.shape[1]) - if features.ndim == 2: - return int(features.shape[0]) - if features.ndim >= 1: - return int(features.reshape(-1, features.shape[-1]).shape[0]) - return 0 - - -def _gemma4_static_image_soft_token_counts( - pixel_position_ids: torch.Tensor | None, - *, - pooling_kernel_size: int, -) -> tuple[int, ...] | None: - if pixel_position_ids is None or pixel_position_ids.ndim != 3: - return None - counts: list[int] = [] - for positions in pixel_position_ids.detach().cpu(): - valid = positions[(positions != -1).any(dim=-1)] - if valid.numel() == 0: - counts.append(0) - continue - max_x = int(valid[:, 0].max().item()) + 1 - max_y = int(valid[:, 1].max().item()) + 1 - counts.append((max_x // int(pooling_kernel_size)) * (max_y // int(pooling_kernel_size))) - return tuple(counts) - - -def _gemma4_static_image_pool_shapes( - pixel_position_ids: torch.Tensor | None, - *, - pooling_kernel_size: int, -) -> tuple[tuple[int, int, int], ...] | None: - if pixel_position_ids is None or pixel_position_ids.ndim != 3: - return None - shapes: list[tuple[int, int, int]] = [] - for positions in pixel_position_ids.detach().cpu(): - valid = positions[(positions != -1).any(dim=-1)] - if valid.numel() == 0: - shapes.append((0, 0, 0)) - continue - grid_w = int(valid[:, 0].max().item()) + 1 - grid_h = int(valid[:, 1].max().item()) + 1 - pooled_count = (grid_w // int(pooling_kernel_size)) * (grid_h // int(pooling_kernel_size)) - shapes.append((grid_h, grid_w, pooled_count)) - return tuple(shapes) - - -def _gemma4_build_native_merge_plan( - multimodal_backbone: torch.nn.Module, - input_ids: torch.Tensor, - *, - image_feature_count: int, - audio_feature_count: int, -) -> _Gemma4NativeMergePlan: - if input_ids.ndim != 2 or int(input_ids.shape[0]) != 1: - raise ValueError( - "Gemma4 native multimodal merge currently expects a static batch-1 input_ids tensor, " - f"got shape {tuple(input_ids.shape)}" - ) - - image_token_id, audio_token_id, pad_token_id = _gemma4_special_token_ids(multimodal_backbone) - token_ids = [int(token) for token in input_ids.detach().cpu().reshape(-1).tolist()] - segments: list[_Gemma4NativeMergeSegment] = [] - pli_token_ids: list[int] = [] - image_offset = 0 - audio_offset = 0 - index = 0 - - while index < len(token_ids): - token_id = token_ids[index] - is_image = image_token_id != 0 and token_id == image_token_id - is_audio = audio_token_id != 0 and token_id == audio_token_id - - if is_image or is_audio: - region_start = index - while index < len(token_ids) and token_ids[index] == token_id: - index += 1 - placeholder_count = index - region_start - if is_image: - insert_count = min(placeholder_count, max(0, image_feature_count - image_offset)) - if insert_count > 0: - segments.append( - _Gemma4NativeMergeSegment( - kind="image", - input_start=region_start, - length=insert_count, - feature_start=image_offset, - ) - ) - pli_token_ids.extend([pad_token_id] * insert_count) - image_offset += insert_count - else: - insert_count = min(placeholder_count, max(0, audio_feature_count - audio_offset)) - if insert_count > 0: - segments.append( - _Gemma4NativeMergeSegment( - kind="audio", - input_start=region_start, - length=insert_count, - feature_start=audio_offset, - ) - ) - pli_token_ids.extend([pad_token_id] * insert_count) - audio_offset += insert_count - continue - - text_start = index - while index < len(token_ids): - next_token = token_ids[index] - if image_token_id != 0 and next_token == image_token_id: - break - if audio_token_id != 0 and next_token == audio_token_id: - break - index += 1 - text_tokens = token_ids[text_start:index] - if text_tokens: - segments.append( - _Gemma4NativeMergeSegment( - kind="text", - input_start=text_start, - length=len(text_tokens), - ) - ) - pli_token_ids.extend(text_tokens) - - if not segments: - raise RuntimeError("Gemma4 native multimodal merge built no input segments") - return _Gemma4NativeMergePlan( - segments=tuple(segments), - pli_token_ids=tuple(pli_token_ids), - ) - - -def _gemma4_feature_sequence( - features: torch.Tensor, - *, - batch_size: int, - feature_dim: int, -) -> torch.Tensor: - if features.ndim == 3: - if int(features.shape[0]) == batch_size: - return features - if int(features.shape[0]) == 1: - return features.expand(batch_size, -1, -1) - if features.ndim == 2: - if int(features.shape[-1]) != feature_dim: - raise ValueError(f"Gemma4 feature dim mismatch: expected {feature_dim}, got {tuple(features.shape)}") - return features.unsqueeze(0).expand(batch_size, -1, -1) - if features.ndim >= 1 and int(features.shape[-1]) == feature_dim: - flattened = features.reshape(-1, feature_dim) - return flattened.unsqueeze(0).expand(batch_size, -1, -1) - raise ValueError(f"unsupported Gemma4 feature tensor shape: {tuple(features.shape)}") - - -def _gemma4_text_embedding_scale(embedding: torch.nn.Module, fallback_scale: float) -> float: - """Return the extra scale needed after calling the HF embedding module. - - Native Cactus scales raw token embedding weights by sqrt(hidden_dim). The - HF Gemma4 embedding module already applies that scale in forward(), so the - transpiler must not multiply a second time when it traces the HF module. - """ - for attr_name in ("scalar_embed_scale", "embed_scale"): - value = getattr(embedding, attr_name, None) - if isinstance(value, torch.Tensor): - try: - if value.numel() > 0 and abs(float(value.reshape(-1)[0].item()) - float(fallback_scale)) < 1e-3: - return 1.0 - except Exception: - continue - elif isinstance(value, (float, int)) and abs(float(value) - float(fallback_scale)) < 1e-3: - return 1.0 - return float(fallback_scale) - - -def _gemma4_apply_native_merge_plan( - model: torch.nn.Module, - *, - input_ids: torch.Tensor, - image_features: torch.Tensor, - audio_features: torch.Tensor, - merge_plan: _Gemma4NativeMergePlan, - pli_token_ids: torch.Tensor, -) -> tuple[torch.Tensor, torch.Tensor]: - embedding = model.get_input_embeddings() - if not isinstance(embedding, torch.nn.Module): - raise TypeError("Gemma4 model is missing input embeddings") - batch_size = int(input_ids.shape[0]) - feature_dim = int(getattr(embedding, "embedding_dim", 0) or image_features.shape[-1] or audio_features.shape[-1]) - target_dtype = _module_floating_dtype(embedding) or image_features.dtype - image_sequence = _gemma4_feature_sequence( - image_features, - batch_size=batch_size, - feature_dim=feature_dim, - ).to(device=input_ids.device, dtype=target_dtype) - audio_sequence = _gemma4_feature_sequence( - audio_features, - batch_size=batch_size, - feature_dim=feature_dim, - ).to(device=input_ids.device, dtype=target_dtype) - - embedded_segments: list[torch.Tensor] = [] - model_config = getattr(model, "config", None) - text_config = getattr(model_config, "text_config", None) - hidden_scale = float(getattr(model_config, "hidden_size", 0) or 0) - if hidden_scale <= 0.0: - hidden_scale = float(getattr(text_config, "hidden_size", 0) or 0) - if hidden_scale <= 0.0: - hidden_scale = float(feature_dim) - hidden_scale = float(hidden_scale) ** 0.5 - text_extra_scale = _gemma4_text_embedding_scale(embedding, hidden_scale) - for segment in merge_plan.segments: - if segment.kind == "text": - text_tokens = input_ids[:, segment.input_start : segment.input_start + segment.length] - text_embeds = embedding(text_tokens) - if text_extra_scale != 1.0: - text_embeds = text_embeds * text_extra_scale - embedded_segments.append(text_embeds) - elif segment.kind == "image": - embedded_segments.append( - image_sequence[:, segment.feature_start : segment.feature_start + segment.length, :] - ) - elif segment.kind == "audio": - embedded_segments.append( - audio_sequence[:, segment.feature_start : segment.feature_start + segment.length, :] - ) - else: - raise RuntimeError(f"unknown Gemma4 merge segment kind: {segment.kind!r}") - - inputs_embeds = torch.cat(embedded_segments, dim=1) - if pli_token_ids.numel() != inputs_embeds.shape[1]: - raise RuntimeError( - "Gemma4 native merge PLI token count mismatch: " - f"{int(pli_token_ids.numel())} vs {int(inputs_embeds.shape[1])}" - ) - pli_tokens = pli_token_ids.to(device=input_ids.device, dtype=input_ids.dtype).unsqueeze(0) - if batch_size != 1: - pli_tokens = pli_tokens.expand(batch_size, -1) - return inputs_embeds, pli_tokens - - -def _gemma4_remap_sequence_tensor( - tensor: torch.Tensor | None, - *, - merge_plan: _Gemma4NativeMergePlan, -) -> torch.Tensor | None: - if tensor is None: - return None - if tensor.ndim != 2: - raise ValueError(f"Gemma4 merge remap expects a rank-2 tensor, got shape {tuple(tensor.shape)}") - remapped_segments: list[torch.Tensor] = [] - for segment in merge_plan.segments: - remapped_segments.append( - tensor[:, segment.input_start : segment.input_start + segment.length] - ) - if not remapped_segments: - return tensor[:, :0] - return torch.cat(remapped_segments, dim=1) - - -def _gemma4_build_standard_causal_mask_mapping( - *, - create_causal_mask: Callable[..., torch.Tensor] | None, - create_sliding_window_causal_mask: Callable[..., torch.Tensor] | None, - config: object, - inputs_embeds: torch.Tensor, - attention_mask: torch.Tensor | None, - position_ids: torch.Tensor, -) -> dict[str, torch.Tensor]: - if not callable(create_causal_mask) or not callable(create_sliding_window_causal_mask): - raise RuntimeError("Gemma4 standard causal mask helpers are unavailable") - mask_kwargs = { - "config": config, - "inputs_embeds": inputs_embeds, - "attention_mask": attention_mask, - "past_key_values": None, - "position_ids": position_ids, - } - return { - "full_attention": create_causal_mask(**mask_kwargs), - "sliding_attention": create_sliding_window_causal_mask(**mask_kwargs), - } - - -def _infer_input_names(module: torch.nn.Module, *, preferred: tuple[str, ...]) -> tuple[str, ...]: - try: - signature = inspect.signature(module.forward) - except (TypeError, ValueError): - return preferred[:1] - - control_names = { - "self", - "return_dict", - "use_cache", - "past_key_values", - "cache_position", - "position_ids", - "labels", - "decoder_input_ids", - "decoder_attention_mask", - "output_attentions", - "output_hidden_states", - } - available = [ - name - for name, parameter in signature.parameters.items() - if name not in control_names - and parameter.kind in ( - inspect.Parameter.POSITIONAL_ONLY, - inspect.Parameter.POSITIONAL_OR_KEYWORD, - inspect.Parameter.KEYWORD_ONLY, - ) - ] - matched = [name for name in preferred if name in available] - if matched: - return tuple(matched) - if available: - return tuple(available[: min(2, len(available))]) - return preferred[:1] - - -class BoundInputAdapter(torch.nn.Module): - def __init__(self, model: torch.nn.Module, *, input_names: tuple[str, ...], family: str, metadata_task: str): - super().__init__() - self.model = model - self.input_names = tuple(input_names) - self.family = family - self.metadata_task = metadata_task - - def _kwargs_from_bound_inputs(self, *bound_inputs: torch.Tensor | None) -> dict[str, torch.Tensor]: - provided = tuple(bound_inputs) - if len(self.input_names) > len(provided): - raise ValueError( - f"adapter expected at most {len(provided)} bound inputs, got {len(self.input_names)} names" - ) - kwargs: dict[str, torch.Tensor] = {} - for index, name in enumerate(self.input_names): - value = provided[index] - if value is None: - raise ValueError(f"missing required bound input {index} for {name}") - kwargs[name] = value - return kwargs - - def get_transpile_metadata(self): - return { - "graph": { - **_transpile_graph_meta( - self.model, - adapter_family=self.family, - adapter_type=type(self).__name__, - input_names=self.input_names, - ), - "task": self.metadata_task, - } - } - - -class CausalLMLogitsAdapter(torch.nn.Module): - def __init__(self, model: torch.nn.Module, *, pad_token_id: int | None = None): - super().__init__() - self.model = model - self.backbone = getattr(model, "model", None) - self.lm_head = getattr(model, "lm_head", None) - self.pad_token_id = pad_token_id if pad_token_id is not None else _resolve_model_pad_token_id(model) - - def forward(self, input_ids: torch.Tensor): - backbone = self.backbone - lm_head = self.lm_head - if isinstance(backbone, torch.nn.Module) and isinstance(lm_head, torch.nn.Module): - backbone_kwargs: dict[str, object] = { - "input_ids": input_ids, - "attention_mask": (input_ids != int(self.pad_token_id)).long() - if self.pad_token_id is not None - else None, - "use_cache": False, - "return_dict": False, - } - outputs = backbone(**_filter_supported_kwargs(backbone, backbone_kwargs)) - hidden_states = _extract_tensor_output(outputs, preferred_field="last_hidden_state") - hidden_states = _select_last_non_pad_token( - hidden_states, - input_ids, - pad_token_id=self.pad_token_id, - ) - return lm_head(hidden_states) - - outputs = self.model( - input_ids=input_ids, - use_cache=False, - return_dict=False, - ) - logits = _extract_tensor_output(outputs, preferred_field="logits") - return _select_last_non_pad_token( - logits, - input_ids, - pad_token_id=self.pad_token_id, - ) - - def get_transpile_metadata(self): - return { - "graph": { - **_transpile_graph_meta( - self.model, - adapter_family="generic", - adapter_type=type(self).__name__, - input_names=("input_ids",), - ), - } - } - - -class Lfm2CausalLMLogitsAdapter(torch.nn.Module): - def __init__(self, model: torch.nn.Module, *, pad_token_id: int | None = None): - super().__init__() - self.model = model - model_root = getattr(model, "model", None) - language_model = getattr(model_root, "language_model", None) - self.backbone = language_model if isinstance(language_model, torch.nn.Module) else model_root - self.lm_head = getattr(model, "lm_head", None) - self.pad_token_id = pad_token_id if pad_token_id is not None else _resolve_model_pad_token_id(model) - from transformers.models.lfm2.modeling_lfm2 import create_causal_mask # type: ignore - - self._create_causal_mask = create_causal_mask - - def forward(self, input_ids: torch.Tensor): - backbone = self.backbone - lm_head = self.lm_head - if not isinstance(backbone, torch.nn.Module) or not isinstance(lm_head, torch.nn.Module): - raise TypeError("LFM2 causal logits adapter requires backbone and lm_head modules") - - inputs_embeds = backbone.embed_tokens(input_ids) - attention_mask = ( - (input_ids != int(self.pad_token_id)).to(dtype=torch.int64) - if self.pad_token_id is not None - else None - ) - position_ids = torch.arange(inputs_embeds.shape[1], device=inputs_embeds.device).unsqueeze(0) - causal_mask = self._create_causal_mask( - config=backbone.config, - inputs_embeds=inputs_embeds, - attention_mask=attention_mask, - past_key_values=None, - position_ids=position_ids, - ) - - hidden_states = inputs_embeds - position_embeddings = backbone.rotary_emb(hidden_states, position_ids=position_ids) - layer_types = tuple(getattr(backbone.config, "layer_types", ())) - linear_attention = attention_mask if inputs_embeds.shape[1] != 1 else None - - for layer_index, decoder_layer in enumerate(backbone.layers[: backbone.config.num_hidden_layers]): - layer_type = layer_types[layer_index] if layer_index < len(layer_types) else "full_attention" - layer_mask = causal_mask if layer_type == "full_attention" else linear_attention - hidden_states = decoder_layer( - hidden_states, - attention_mask=layer_mask, - position_embeddings=position_embeddings, - position_ids=position_ids, - past_key_values=None, - use_cache=False, - ) - - hidden_states = backbone.embedding_norm(hidden_states) - hidden_states = _select_last_non_pad_token( - hidden_states, - input_ids, - pad_token_id=self.pad_token_id, - ) - return lm_head(hidden_states) - - def get_transpile_metadata(self): - return { - "graph": { - **_transpile_graph_meta( - self.model, - adapter_family="lfm2", - adapter_type=type(self).__name__, - input_names=("input_ids",), - ), - "num_hidden_layers": int(self.backbone.config.num_hidden_layers), - } - } - - -class Lfm2CausalLMStepAdapter(torch.nn.Module): - def __init__(self, model: torch.nn.Module, *, pad_token_id: int | None = None): - super().__init__() - self.model = model - model_root = getattr(model, "model", None) - language_model = getattr(model_root, "language_model", None) - self.backbone = language_model if isinstance(language_model, torch.nn.Module) else model_root - self.lm_head = getattr(model, "lm_head", None) - self.pad_token_id = pad_token_id if pad_token_id is not None else _resolve_model_pad_token_id(model) - - def forward(self, input_ids: torch.Tensor, position_ids: torch.Tensor) -> torch.Tensor: - backbone = self.backbone - lm_head = self.lm_head - if not isinstance(backbone, torch.nn.Module) or not isinstance(lm_head, torch.nn.Module): - raise TypeError("LFM2 causal step adapter requires backbone and lm_head modules") - - inputs_embeds = backbone.embed_tokens(input_ids) - text_position_ids = position_ids.to(dtype=torch.int64) - causal_mask = None - - hidden_states = inputs_embeds - position_embeddings = backbone.rotary_emb(hidden_states, position_ids=text_position_ids) - layer_types = tuple(getattr(backbone.config, "layer_types", ())) - - for layer_index, decoder_layer in enumerate(backbone.layers[: backbone.config.num_hidden_layers]): - layer_type = layer_types[layer_index] if layer_index < len(layer_types) else "full_attention" - layer_mask = causal_mask if layer_type == "full_attention" else None - hidden_states = decoder_layer( - hidden_states, - attention_mask=layer_mask, - position_embeddings=position_embeddings, - position_ids=text_position_ids, - past_key_values=None, - use_cache=False, - ) - - hidden_states = backbone.embedding_norm(hidden_states) - return lm_head(hidden_states[:, -1:, :]) - - def get_transpile_metadata(self): - layer_types = tuple(str(value) for value in (getattr(self.backbone.config, "layer_types", ()) or ())) - return { - "graph": { - **_transpile_graph_meta( - self.model, - adapter_family="lfm2", - adapter_type=type(self).__name__, - input_names=("input_ids", "position_ids"), - ), - "num_hidden_layers": int(self.backbone.config.num_hidden_layers), - "layer_types": layer_types, - } - } - - -def _lfm2_vl_model_root(model: torch.nn.Module) -> torch.nn.Module: - root = getattr(model, "model", None) - if not isinstance(root, torch.nn.Module): - raise TypeError("LFM2-VL adapter requires a model.model module") - return root - - -def _lfm2_language_backbone(model: torch.nn.Module) -> torch.nn.Module: - root = getattr(model, "model", None) - language_model = getattr(root, "language_model", None) - backbone = language_model if isinstance(language_model, torch.nn.Module) else root - if not isinstance(backbone, torch.nn.Module): - raise TypeError("LFM2 adapter requires a language model backbone") - return backbone - - -class Lfm2VlVisionEncoderAdapter(torch.nn.Module): - def __init__( - self, - model: torch.nn.Module, - *, - weights_dir: str | None = None, - ): - super().__init__() - self.model = model - self.weights_dir = weights_dir - root = _lfm2_vl_model_root(model) - vision_tower = getattr(root, "vision_tower", None) - if not isinstance(vision_tower, torch.nn.Module): - raise TypeError("LFM2-VL model is missing a vision_tower module") - vision_model = getattr(vision_tower, "vision_model", None) - if not isinstance(vision_model, torch.nn.Module): - raise TypeError("LFM2-VL vision_tower is missing vision_model") - - self.embeddings = vision_model.embeddings - self.encoder = vision_model.encoder - self.post_layernorm = vision_model.post_layernorm - - def forward( - self, - pixel_values: torch.Tensor, - pixel_attention_mask: torch.Tensor, - positional_embeddings: torch.Tensor, - ) -> torch.Tensor: - target_dtype = self.embeddings.patch_embedding.weight.dtype - hidden_states = self.embeddings.patch_embedding(pixel_values.to(dtype=target_dtype)) - hidden_states = hidden_states + positional_embeddings.to( - device=hidden_states.device, - dtype=hidden_states.dtype, - ) - masked_bias = -30000.0 - seq_len = hidden_states.shape[1] - key_valid = pixel_attention_mask[:, None, None, :].to(dtype=hidden_states.dtype) - encoder_attention_mask = (key_valid * (-masked_bias) + masked_bias).expand( - hidden_states.shape[0], 1, seq_len, seq_len - ) - for encoder_layer in self.encoder.layers: - hidden_states = encoder_layer(hidden_states, encoder_attention_mask) - hidden_states = self.post_layernorm(hidden_states) - return hidden_states - - def get_transpile_metadata(self): - return { - "graph": { - **_transpile_graph_meta( - self.model, - adapter_family="lfm2_vl", - adapter_type=type(self).__name__, - input_names=("pixel_values", "pixel_attention_mask", "positional_embeddings"), - ), - "weights_dir": self.weights_dir, - } - } - - -class Lfm2VlVisionProjectorAdapter(torch.nn.Module): - def __init__( - self, - model: torch.nn.Module, - *, - weights_dir: str | None = None, - ): - super().__init__() - self.model = model - self.weights_dir = weights_dir - root = _lfm2_vl_model_root(model) - projector = getattr(root, "multi_modal_projector", None) - if not isinstance(projector, torch.nn.Module): - raise TypeError("LFM2-VL model is missing multi_modal_projector") - self.projector = projector - - def forward(self, vision_features: torch.Tensor) -> torch.Tensor: - hidden_states = vision_features - if self.projector.use_layer_norm: - hidden_states = self.projector.layer_norm(hidden_states) - hidden_states = self.projector.linear_1(hidden_states) - hidden_states = self.projector.act(hidden_states) - hidden_states = self.projector.linear_2(hidden_states) - return hidden_states - - def get_transpile_metadata(self): - return { - "graph": { - **_transpile_graph_meta( - self.model, - adapter_family="lfm2_vl", - adapter_type=type(self).__name__, - input_names=("vision_features",), - ), - "weights_dir": self.weights_dir, - } - } - - -class Lfm2VlLMEncoderAdapter(torch.nn.Module): - def __init__( - self, - model: torch.nn.Module, - *, - weights_dir: str | None = None, - ): - super().__init__() - self.model = model - self.weights_dir = weights_dir - self.backbone = _lfm2_language_backbone(model) - - def forward( - self, - input_ids: torch.Tensor, - attention_mask: torch.Tensor, - ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: - inputs_embeds = self.backbone.embed_tokens(input_ids) - position_ids = torch.arange(inputs_embeds.shape[1], device=inputs_embeds.device).unsqueeze(0) - return inputs_embeds, attention_mask.to(dtype=torch.int64), position_ids - - def get_transpile_metadata(self): - return { - "graph": { - **_transpile_graph_meta( - self.model, - adapter_family="lfm2_vl", - adapter_type=type(self).__name__, - input_names=("input_ids", "attention_mask"), - ), - "weights_dir": self.weights_dir, - } - } - - -class Lfm2VlLMEncoderStepAdapter(torch.nn.Module): - def __init__(self, model: torch.nn.Module, *, weights_dir: str | None = None): - super().__init__() - self.model = model - self.weights_dir = weights_dir - self.backbone = _lfm2_language_backbone(model) - - def forward( - self, - input_ids: torch.Tensor, - position_ids: torch.Tensor, - ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: - inputs_embeds = self.backbone.embed_tokens(input_ids) - attention_mask = input_ids.to(dtype=torch.int64) * 0 + 1 - return inputs_embeds, attention_mask, position_ids.to(dtype=torch.int64) - - def get_transpile_metadata(self): - return { - "graph": { - **_transpile_graph_meta( - self.model, - adapter_family="lfm2_vl", - adapter_type=type(self).__name__, - input_names=("input_ids", "position_ids"), - ), - "weights_dir": self.weights_dir, - } - } - - -class Lfm2VlLMEncoderTextChunkAdapter(Lfm2VlLMEncoderStepAdapter): - pass - - -class Lfm2VlDecoderAdapter(torch.nn.Module): - def __init__( - self, - model: torch.nn.Module, - *, - weights_dir: str | None = None, - last_token_only: bool = True, - return_hidden: bool = False, - ): - super().__init__() - self.model = model - self.weights_dir = weights_dir - self.last_token_only = bool(last_token_only) - self.return_hidden = bool(return_hidden) - self.backbone = _lfm2_language_backbone(model) - self.lm_head = getattr(model, "lm_head", None) - if not isinstance(self.lm_head, torch.nn.Module): - raise TypeError("LFM2-VL decoder adapter requires an lm_head module") - from transformers.models.lfm2.modeling_lfm2 import create_causal_mask # type: ignore - - self._create_causal_mask = create_causal_mask - - def forward( - self, - inputs_embeds: torch.Tensor, - attention_mask: torch.Tensor, - position_ids: torch.Tensor, - ) -> torch.Tensor: - causal_mask = self._create_causal_mask( - config=self.backbone.config, - inputs_embeds=inputs_embeds, - attention_mask=attention_mask, - past_key_values=None, - position_ids=position_ids, - ) - hidden_states = inputs_embeds - position_embeddings = self.backbone.rotary_emb(hidden_states, position_ids=position_ids) - layer_types = tuple(getattr(self.backbone.config, "layer_types", ())) - linear_attention = attention_mask if inputs_embeds.shape[1] != 1 else None - - for layer_index, decoder_layer in enumerate(self.backbone.layers[: self.backbone.config.num_hidden_layers]): - layer_type = layer_types[layer_index] if layer_index < len(layer_types) else "full_attention" - layer_mask = causal_mask if layer_type == "full_attention" else linear_attention - hidden_states = decoder_layer( - hidden_states, - attention_mask=layer_mask, - position_embeddings=position_embeddings, - position_ids=position_ids, - past_key_values=None, - use_cache=False, - ) - - hidden_states = self.backbone.embedding_norm(hidden_states) - if self.return_hidden: - return hidden_states - if self.last_token_only: - hidden_states = hidden_states[:, -1:, :] - return self.lm_head(hidden_states) - - def get_transpile_metadata(self): - return { - "graph": { - **_transpile_graph_meta( - self.model, - adapter_family="lfm2_vl", - adapter_type=type(self).__name__, - input_names=("inputs_embeds", "attention_mask", "position_ids"), - ), - "weights_dir": self.weights_dir, - } - } - - -class Lfm2VlMultimodalCausalLMLogitsAdapter(torch.nn.Module): - def __init__(self, model: torch.nn.Module, *, input_names: tuple[str, ...]): - super().__init__() - self.model = model - self.input_names = tuple(input_names) - - def forward(self, *args: torch.Tensor) -> torch.Tensor: - kwargs = { - name: value - for name, value in zip(self.input_names, args, strict=True) - } - outputs = self.model( - **kwargs, - use_cache=False, - return_dict=True, - logits_to_keep=1, - ) - return _extract_tensor_output(outputs, preferred_field="logits") - - def get_transpile_metadata(self): - return { - "graph": _transpile_graph_meta( - self.model, - adapter_family="lfm2_vl", - adapter_type=type(self).__name__, - input_names=self.input_names, - ), - } - - -class GemmaCausalLMLogitsAdapter(torch.nn.Module): - def __init__(self, model: torch.nn.Module, *, pad_token_id: int | None = None): - super().__init__() - self.model = model - self.backbone = model.model - self.pad_token_id = pad_token_id if pad_token_id is not None else _resolve_model_pad_token_id(model) - from transformers.models.gemma.modeling_gemma import create_causal_mask # type: ignore - - self._create_causal_mask = create_causal_mask - - def forward(self, input_ids: torch.Tensor): - return self.debug_forward(input_ids)[0] - - def debug_forward(self, input_ids: torch.Tensor) -> tuple[torch.Tensor, list[torch.Tensor]]: - inputs_embeds = self.backbone.embed_tokens(input_ids) - position_ids = torch.arange(inputs_embeds.shape[1], device=inputs_embeds.device).unsqueeze(0) - causal_mask = self._create_causal_mask( - self.backbone.config, - inputs_embeds, - None, - past_key_values=None, - position_ids=position_ids, - ) - - hidden_states = inputs_embeds - position_embeddings = self.backbone.rotary_emb(hidden_states, position_ids=position_ids) - checkpoints: list[torch.Tensor] = [] - - for decoder_layer in self.backbone.layers[: self.backbone.config.num_hidden_layers]: - hidden_states = decoder_layer( - hidden_states, - attention_mask=causal_mask, - position_ids=position_ids, - past_key_values=None, - use_cache=False, - position_embeddings=position_embeddings, - ) - checkpoints.append(hidden_states) - - hidden_states = self.backbone.norm(hidden_states) - hidden_states = _select_last_non_pad_token( - hidden_states, - input_ids, - pad_token_id=self.pad_token_id, - ) - checkpoints.append(hidden_states) - return _gemma4_apply_final_logit_softcapping(self.model, self.model.lm_head(hidden_states)), checkpoints - - def get_transpile_metadata(self): - return { - "graph": { - **_transpile_graph_meta( - self.model, - adapter_family="gemma", - adapter_type=type(self).__name__, - input_names=("input_ids",), - ), - "num_hidden_layers": int(self.backbone.config.num_hidden_layers), - } - } - - -def _gemma3_text_backbone_forward(backbone, inputs_embeds, causal_mask_mapping, position_ids, - checkpoints=None): - hidden_states = inputs_embeds - layer_types = backbone.config.layer_types - position_embeddings = { - layer_type: backbone.rotary_emb(hidden_states, position_ids, layer_type) - for layer_type in set(layer_types) - } - for i, decoder_layer in enumerate(backbone.layers[: backbone.config.num_hidden_layers]): - hidden_states = decoder_layer( - hidden_states, - attention_mask=causal_mask_mapping[layer_types[i]], - position_embeddings=position_embeddings[layer_types[i]], - position_ids=position_ids, - past_key_values=None, - ) - if checkpoints is not None: - checkpoints.append(hidden_states) - return backbone.norm(hidden_states) - - -def _gemma3_decoder_graph_meta(model, backbone, adapter_type, input_names): - sliding_window = getattr(backbone.config, "sliding_window", None) - return { - "graph": { - **_transpile_graph_meta( - model, - adapter_family="gemma3", - adapter_type=adapter_type, - input_names=input_names, - ), - "num_hidden_layers": int(backbone.config.num_hidden_layers), - "layer_types": tuple(getattr(backbone.config, "layer_types", [])), - "sliding_window": None if sliding_window is None else int(sliding_window), - } - } - - -class Gemma3CausalLMLogitsAdapter(torch.nn.Module): - def __init__(self, model: torch.nn.Module, *, pad_token_id: int | None = None): - super().__init__() - self.model = model - self.backbone = model.model - self.pad_token_id = pad_token_id if pad_token_id is not None else _resolve_model_pad_token_id(model) - from transformers.models.gemma3.modeling_gemma3 import create_causal_mask # type: ignore - from transformers.models.gemma3.modeling_gemma3 import create_sliding_window_causal_mask # type: ignore - - self._create_causal_mask = create_causal_mask - self._create_sliding_window_causal_mask = create_sliding_window_causal_mask - - def forward(self, input_ids: torch.Tensor): - return self.debug_forward(input_ids)[0] - - def debug_forward(self, input_ids: torch.Tensor) -> tuple[torch.Tensor, list[torch.Tensor]]: - inputs_embeds = self.backbone.embed_tokens(input_ids) - position_ids = torch.arange(inputs_embeds.shape[1], device=inputs_embeds.device).unsqueeze(0) - causal_mask_mapping = { - "full_attention": self._create_causal_mask( - self.backbone.config, - inputs_embeds, - None, - None, - past_key_values=None, - position_ids=position_ids, - ), - "sliding_attention": self._create_sliding_window_causal_mask( - self.backbone.config, - inputs_embeds, - None, - None, - past_key_values=None, - position_ids=position_ids, - ), - } - - checkpoints: list[torch.Tensor] = [] - hidden_states = _gemma3_text_backbone_forward( - self.backbone, inputs_embeds, causal_mask_mapping, position_ids, checkpoints=checkpoints, - ) - checkpoints.append(hidden_states) - return self.model.lm_head(hidden_states), checkpoints - - def get_transpile_metadata(self): - return _gemma3_decoder_graph_meta(self.model, self.backbone, type(self).__name__, ("input_ids",)) - - -class Gemma3LMEncoderStepAdapter(torch.nn.Module): - def __init__(self, model: torch.nn.Module): - super().__init__() - self.backbone = model.model - - def forward(self, input_ids: torch.Tensor, position_ids: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: - inputs_embeds = self.backbone.embed_tokens(input_ids) - return inputs_embeds.to(torch.float16), position_ids.to(dtype=torch.int64) - - def get_transpile_metadata(self): - return { - "graph": { - **_transpile_graph_meta( - self.backbone, - adapter_family="gemma3", - adapter_type=type(self).__name__, - input_names=("input_ids", "position_ids"), - ), - } - } - - -class Gemma3EmbedsCausalLMStepAdapter(torch.nn.Module): - def __init__(self, model: torch.nn.Module): - super().__init__() - self.model = model - self.backbone = model.model - - def _additive_mask_mapping(self, inputs_embeds: torch.Tensor) -> dict[str, torch.Tensor]: - seq_len = int(inputs_embeds.shape[1]) - zero_mask = torch.zeros( - (1, 1, seq_len, seq_len), - dtype=inputs_embeds.dtype, - device=inputs_embeds.device, - ) - return {layer_type: zero_mask for layer_type in set(self.backbone.config.layer_types)} - - def forward(self, inputs_embeds: torch.Tensor, position_ids: torch.Tensor) -> torch.Tensor: - hidden_states = inputs_embeds.to(self.backbone.norm.weight.dtype) - causal_mask_mapping = self._additive_mask_mapping(hidden_states) - text_position_ids = position_ids.to(dtype=torch.int64) - hidden_states = _gemma3_text_backbone_forward( - self.backbone, hidden_states, causal_mask_mapping, text_position_ids, - ) - logits = self.model.lm_head(hidden_states[:, -1:, :]) - return _gemma4_apply_final_logit_softcapping(self.model, logits) - - def get_transpile_metadata(self): - return _gemma3_decoder_graph_meta(self.model, self.backbone, type(self).__name__, ("inputs_embeds", "position_ids")) - - -class Gemma3EmbedsCausalLMPrefillChunkAdapter(Gemma3EmbedsCausalLMStepAdapter): - def _additive_mask_mapping(self, inputs_embeds: torch.Tensor) -> dict[str, torch.Tensor]: - seq_len = int(inputs_embeds.shape[1]) - blocked_values = torch.ones( - (1, 1, seq_len, seq_len), - dtype=inputs_embeds.dtype, - device=inputs_embeds.device, - ) * torch.finfo(inputs_embeds.dtype).min - allowed_values = torch.zeros( - (1, 1, seq_len, seq_len), - dtype=inputs_embeds.dtype, - device=inputs_embeds.device, - ) - causal = torch.tril( - torch.ones((seq_len, seq_len), dtype=torch.bool, device=inputs_embeds.device), - ).view(1, 1, seq_len, seq_len) - mapping = {"full_attention": torch.where(causal, allowed_values, blocked_values)} - if "sliding_attention" in set(self.backbone.config.layer_types): - window = int(self.backbone.config.sliding_window) - below_window = torch.tril( - torch.ones((seq_len, seq_len), dtype=torch.bool, device=inputs_embeds.device), - diagonal=-window, - ).view(1, 1, seq_len, seq_len) - mapping["sliding_attention"] = torch.where( - below_window, blocked_values, mapping["full_attention"], - ) - return mapping - - -class Gemma4CausalLMLogitsAdapter(torch.nn.Module): - def __init__(self, model: torch.nn.Module, *, pad_token_id: int | None = None): - super().__init__() - self.model = model - model_backbone = model.model - self.backbone = getattr(model_backbone, "language_model", model_backbone) - self.pad_token_id = pad_token_id if pad_token_id is not None else _resolve_model_pad_token_id(model) - from transformers.models.gemma4.modeling_gemma4 import create_causal_mask # type: ignore - from transformers.models.gemma4.modeling_gemma4 import create_sliding_window_causal_mask # type: ignore - - self._create_causal_mask = create_causal_mask - self._create_sliding_window_causal_mask = create_sliding_window_causal_mask - - def forward(self, input_ids: torch.Tensor): - return self.debug_forward(input_ids)[0] - - def debug_forward(self, input_ids: torch.Tensor) -> tuple[torch.Tensor, list[torch.Tensor]]: - inputs_embeds = self.backbone.embed_tokens(input_ids) - per_layer_inputs = None - if self.backbone.hidden_size_per_layer_input: - per_layer_inputs = _gemma4_get_per_layer_inputs(self.backbone, input_ids, inputs_embeds) - per_layer_inputs = self.backbone.project_per_layer_inputs(inputs_embeds, per_layer_inputs) - - position_ids = torch.arange(inputs_embeds.shape[1], device=inputs_embeds.device).unsqueeze(0) - mask_kwargs = { - "config": self.backbone.config, - "inputs_embeds": inputs_embeds, - "attention_mask": None, - "past_key_values": None, - "position_ids": position_ids, - } - causal_mask_mapping = { - "full_attention": self._create_causal_mask(**mask_kwargs), - "sliding_attention": self._create_sliding_window_causal_mask(**mask_kwargs), - } - - hidden_states = inputs_embeds - checkpoints: list[torch.Tensor] = [] - position_embeddings = { - layer_type: self.backbone.rotary_emb(hidden_states, position_ids, layer_type) - for layer_type in self.backbone.unique_layer_types - } - shared_kv_states: dict[str, torch.Tensor] = {} - - for i, decoder_layer in enumerate(self.backbone.layers[: self.backbone.config.num_hidden_layers]): - per_layer_input = per_layer_inputs[:, :, i, :] if per_layer_inputs is not None else None - hidden_states = decoder_layer( - hidden_states, - per_layer_input, - shared_kv_states=shared_kv_states, - position_embeddings=position_embeddings[self.backbone.config.layer_types[i]], - attention_mask=causal_mask_mapping[self.backbone.config.layer_types[i]], - position_ids=position_ids, - past_key_values=None, - ) - checkpoints.append(hidden_states) - - hidden_states = self.backbone.norm(hidden_states) - checkpoints.append(hidden_states) - return self.model.lm_head(hidden_states), checkpoints - - def debug_first_block(self, input_ids: torch.Tensor) -> dict[str, torch.Tensor]: - inputs_embeds = self.backbone.embed_tokens(input_ids) - per_layer_inputs = None - if self.backbone.hidden_size_per_layer_input: - per_layer_inputs = _gemma4_get_per_layer_inputs(self.backbone, input_ids, inputs_embeds) - per_layer_inputs = self.backbone.project_per_layer_inputs(inputs_embeds, per_layer_inputs) - - position_ids = torch.arange(inputs_embeds.shape[1], device=inputs_embeds.device).unsqueeze(0) - mask_kwargs = { - "config": self.backbone.config, - "inputs_embeds": inputs_embeds, - "attention_mask": None, - "past_key_values": None, - "position_ids": position_ids, - } - causal_mask_mapping = { - "full_attention": self._create_causal_mask(**mask_kwargs), - "sliding_attention": self._create_sliding_window_causal_mask(**mask_kwargs), - } - - hidden_states = inputs_embeds - position_embeddings = { - layer_type: self.backbone.rotary_emb(hidden_states, position_ids, layer_type) - for layer_type in self.backbone.unique_layer_types - } - shared_kv_states: dict[str, torch.Tensor] = {} - layer = self.backbone.layers[0] - layer_type = self.backbone.config.layer_types[0] - per_layer_input = per_layer_inputs[:, :, 0, :] if per_layer_inputs is not None else None - - checkpoints: dict[str, torch.Tensor] = {} - - residual = hidden_states - normed = layer.input_layernorm(hidden_states) - checkpoints["pre_attn_norm"] = normed - - attn_out = layer.self_attn( - normed, - position_embeddings=position_embeddings[layer_type], - attention_mask=causal_mask_mapping[layer_type], - position_ids=position_ids, - past_key_values=None, - shared_kv_states=shared_kv_states, - ) - if isinstance(attn_out, tuple): - attn_out = attn_out[0] - checkpoints["attn_o_proj"] = attn_out - - post_attn_norm = layer.post_attention_layernorm(attn_out) - checkpoints["post_attn_norm"] = post_attn_norm - - after_attention = residual + post_attn_norm - checkpoints["after_attention_residual"] = after_attention - - pre_ffn_norm = layer.pre_feedforward_layernorm(after_attention) - checkpoints["pre_ffn_norm"] = pre_ffn_norm - - mlp_out = layer.mlp(pre_ffn_norm) - checkpoints["mlp_down"] = mlp_out - - post_ffn_norm = layer.post_feedforward_layernorm(mlp_out) - checkpoints["post_ffn_norm"] = post_ffn_norm - - after_ffn = after_attention + post_ffn_norm - checkpoints["after_ffn_residual"] = after_ffn - - if per_layer_input is not None: - gated = layer.per_layer_input_gate(after_ffn) - gated = layer.act_fn(gated) - projected = gated * per_layer_input - per_layer_proj = layer.per_layer_projection(projected) - checkpoints["per_layer_input_proj"] = per_layer_proj - post_per_layer_input_norm = layer.post_per_layer_input_norm(per_layer_proj) - checkpoints["post_per_layer_input_norm"] = post_per_layer_input_norm - after_ffn = after_ffn + post_per_layer_input_norm - - layer_scalar = getattr(layer, "layer_scalar", None) - if layer_scalar is not None: - after_ffn = after_ffn * layer_scalar - checkpoints["layer_scalar_out"] = after_ffn - return checkpoints - - def get_transpile_metadata(self): - sliding_window = getattr(self.backbone.config, "sliding_window", None) - layer_types = list(getattr(self.backbone.config, "layer_types", [])) - return { - "graph": { - **_transpile_graph_meta( - self.model, - adapter_family="gemma4", - adapter_type=type(self).__name__, - input_names=("input_ids",), - ), - "num_hidden_layers": int(self.backbone.config.num_hidden_layers), - "layer_types": tuple(layer_types), - "sliding_window": None if sliding_window is None else int(sliding_window), - } - } - - -def _gemma4_apply_final_logit_softcapping(model: torch.nn.Module, logits: torch.Tensor) -> torch.Tensor: - config = getattr(model, "config", None) - text_config = None - get_text_config = getattr(config, "get_text_config", None) - if callable(get_text_config): - try: - text_config = get_text_config() - except Exception: - text_config = None - cap = getattr(text_config, "final_logit_softcapping", None) - if cap is None: - cap = getattr(config, "final_logit_softcapping", None) - if cap is None: - return logits - cap_value = float(cap) - if cap_value <= 0.0: - return logits - return torch.tanh(logits / cap_value) * cap_value - - -class Gemma4MultimodalCausalLMLogitsAdapter(BoundInputAdapter): - def __init__(self, model: torch.nn.Module, *, input_names: tuple[str, ...], weights_dir: str | None = None): - super().__init__( - model, - input_names=input_names, - family="gemma4", - metadata_task="multimodal_causal_lm_logits", - ) - model_backbone = model.model - self.multimodal_backbone = model_backbone - self.backbone = getattr(model_backbone, "language_model", model_backbone) - self.last_token_logits_only = False - self._use_cached_multimodal_features = False - self._native_merge_plan: _Gemma4NativeMergePlan | None = None - self.register_buffer("_cached_image_features", torch.empty(0), persistent=False) - self.register_buffer("_cached_audio_features", torch.empty(0), persistent=False) - self.register_buffer("_native_merge_pli_token_ids", torch.empty(0, dtype=torch.long), persistent=False) - vision_post_proj_norm = _gemma4_load_vision_post_proj_norm(weights_dir) - if vision_post_proj_norm is None: - vision_post_proj_norm = torch.empty(0) - self.register_buffer("_cactus_vision_post_proj_norm", vision_post_proj_norm, persistent=False) - self._capture_cpu_float32_text_modules: list[tuple[torch.nn.Module, torch.dtype]] = [] - self._create_causal_mask_mapping = None - self._create_masks_for_generate = None - self._create_causal_mask = None - self._create_sliding_window_causal_mask = None - try: - from transformers.models.gemma4.modeling_gemma4 import create_causal_mask # type: ignore - from transformers.models.gemma4.modeling_gemma4 import create_causal_mask_mapping # type: ignore - from transformers.models.gemma4.modeling_gemma4 import create_masks_for_generate # type: ignore - from transformers.models.gemma4.modeling_gemma4 import create_sliding_window_causal_mask # type: ignore - - self._create_causal_mask = create_causal_mask - self._create_causal_mask_mapping = create_causal_mask_mapping - self._create_masks_for_generate = create_masks_for_generate - self._create_sliding_window_causal_mask = create_sliding_window_causal_mask - except Exception: - pass - - def _apply_final_logit_softcapping(self, logits: torch.Tensor) -> torch.Tensor: - return _gemma4_apply_final_logit_softcapping(self.model, logits) - - def _capture_text_modules(self) -> tuple[torch.nn.Module, ...]: - return _unique_modules( - self.backbone, - getattr(self.model, "get_input_embeddings", lambda: None)(), - getattr(self.model, "lm_head", None), - ) - - def _compute_image_features( - self, - pixel_values: torch.Tensor, - pixel_position_ids: torch.Tensor | None, - get_image_features: Callable[..., object], - ) -> torch.Tensor: - multimodal_backbone = self.multimodal_backbone - if _gemma4_can_use_native_like_vision_features(multimodal_backbone): - return _gemma4_compute_native_like_image_features( - multimodal_backbone, - pixel_values, - pixel_position_ids, - post_proj_norm_weight=self._cactus_vision_post_proj_norm, - ) - vision_tower = getattr(multimodal_backbone, "vision_tower", None) - embed_vision = getattr(multimodal_backbone, "embed_vision", None) - vision_modules = _gemma4_vision_modules(multimodal_backbone) - if ( - pixel_values.device.type != "cpu" - or pixel_values.dtype != torch.float16 - or len(vision_modules) != 2 - or _torch_is_compiling() - ): - return get_image_features( - pixel_values, - pixel_position_ids, - None, - return_dict=True, - ).pooler_output - - vision_dtype = _module_floating_dtype(vision_tower) - embed_dtype = _module_floating_dtype(embed_vision) - if vision_dtype != torch.float16 and embed_dtype != torch.float16: - return get_image_features( - pixel_values, - pixel_position_ids, - None, - return_dict=True, - ).pooler_output - - # Gemma4's CPU float16 vision path can emit non-finite soft tokens; upcast only - # the static image feature extraction path and restore the original module dtypes. - with _temporary_cpu_float32_modules(vision_modules): - return get_image_features( - pixel_values.float(), - pixel_position_ids, - None, - return_dict=True, - ).pooler_output - - def _compute_audio_features( - self, - input_features: torch.Tensor, - input_features_mask: torch.Tensor, - get_audio_features: Callable[..., object], - ) -> torch.Tensor: - multimodal_backbone = self.multimodal_backbone - if _gemma4_can_use_native_like_audio_features(multimodal_backbone): - return _gemma4_compute_native_like_audio_features( - multimodal_backbone, - input_features, - input_features_mask, - ) - audio_modules = _gemma4_audio_modules(multimodal_backbone) - if ( - input_features.device.type != "cpu" - or input_features.dtype != torch.float16 - or len(audio_modules) != 2 - or _torch_is_compiling() - ): - audio_output = get_audio_features(input_features, input_features_mask, return_dict=True) - return _gemma4_strip_audio_padding(audio_output) - - if all(_module_floating_dtype(module) != torch.float16 for module in audio_modules): - audio_output = get_audio_features(input_features, input_features_mask, return_dict=True) - return _gemma4_strip_audio_padding(audio_output) - - with _temporary_cpu_float32_modules(audio_modules): - audio_output = get_audio_features(input_features.float(), input_features_mask, return_dict=True) - return _gemma4_strip_audio_padding(audio_output) - - def prepare_cpu_float32_capture(self) -> None: - self._capture_cpu_float32_text_modules.clear() - if os.environ.get("CACTUS_GEMMA4_CAPTURE_FP32") != "1": - return - for module in self._capture_text_modules(): - device = _module_device(module) - dtype = _module_floating_dtype(module) - if device is None or device.type != "cpu" or dtype is None or dtype == torch.float32: - continue - self._capture_cpu_float32_text_modules.append((module, dtype)) - module.to(dtype=torch.float32) - - def restore_cpu_float32_capture(self) -> None: - for module, dtype in reversed(self._capture_cpu_float32_text_modules): - module.to(dtype=dtype) - self._capture_cpu_float32_text_modules.clear() - - def _resolve_image_features( - self, - *, - inputs_embeds: torch.Tensor, - pixel_values: torch.Tensor | None, - pixel_position_ids: torch.Tensor | None, - get_image_features: Callable[..., object], - ) -> torch.Tensor | None: - if pixel_values is None: - return None - if self._use_cached_multimodal_features and self._cached_image_features.numel() > 0: - image_features = self._cached_image_features - else: - image_features = self._compute_image_features( - pixel_values, - pixel_position_ids, - get_image_features, - ) - return image_features.to(inputs_embeds.device, inputs_embeds.dtype) - - def _resolve_audio_features( - self, - *, - inputs_embeds: torch.Tensor, - input_features: torch.Tensor | None, - input_features_mask: torch.Tensor | None, - get_audio_features: Callable[..., object], - ) -> torch.Tensor | None: - if input_features is None or input_features_mask is None: - return None - if self._use_cached_multimodal_features and self._cached_audio_features.numel() > 0: - audio_features = self._cached_audio_features - else: - audio_features = self._compute_audio_features( - input_features, - input_features_mask, - get_audio_features, - ) - return audio_features.to(inputs_embeds.device, inputs_embeds.dtype) - - def _prepare_text_backbone_inputs( - self, - *, - input_ids: torch.Tensor, - attention_mask: torch.Tensor | None, - token_type_ids: torch.Tensor | None, - pixel_values: torch.Tensor | None, - pixel_position_ids: torch.Tensor | None, - input_features: torch.Tensor | None, - input_features_mask: torch.Tensor | None, - get_placeholder_mask: Callable[..., object], - get_image_features: Callable[..., object], - get_audio_features: Callable[..., object], - ) -> tuple[torch.Tensor, torch.Tensor | None, dict[str, torch.Tensor], torch.LongTensor]: - inputs_embeds = self.model.get_input_embeddings()(input_ids) - - image_features = self._resolve_image_features( - inputs_embeds=inputs_embeds, - pixel_values=pixel_values, - pixel_position_ids=pixel_position_ids, - get_image_features=get_image_features, - ) - audio_features = self._resolve_audio_features( - inputs_embeds=inputs_embeds, - input_features=input_features, - input_features_mask=input_features_mask, - get_audio_features=get_audio_features, - ) - - if image_features is not None and audio_features is not None and self._native_merge_plan is not None: - inputs_embeds, per_layer_inputs_tokens = _gemma4_apply_native_merge_plan( - self.model, - input_ids=input_ids, - image_features=image_features, - audio_features=audio_features, - merge_plan=self._native_merge_plan, - pli_token_ids=self._native_merge_pli_token_ids, - ) - attention_mask = _gemma4_remap_sequence_tensor( - attention_mask, - merge_plan=self._native_merge_plan, - ) - token_type_ids = _gemma4_remap_sequence_tensor( - token_type_ids, - merge_plan=self._native_merge_plan, - ) - if token_type_ids is None: - raise RuntimeError("Gemma4 native merge requires token_type_ids for multimodal attention masking") - else: - text_mask, image_mask, audio_mask = _gemma4_get_placeholder_masks( - get_placeholder_mask, - token_type_ids=token_type_ids, - input_ids=input_ids, - inputs_embeds=inputs_embeds, - ) - if image_features is not None: - inputs_embeds = inputs_embeds.masked_scatter( - image_mask.unsqueeze(-1).expand_as(inputs_embeds), - image_features, - ) - if audio_features is not None: - inputs_embeds = inputs_embeds.masked_scatter( - audio_mask.unsqueeze(-1).expand_as(inputs_embeds), - audio_features, - ) - per_layer_inputs_tokens = input_ids * text_mask.to(dtype=input_ids.dtype) - - per_layer_inputs = None - text_config = _gemma4_text_config(self.multimodal_backbone) - if getattr(text_config, "hidden_size_per_layer_input", None): - per_layer_inputs = _gemma4_get_per_layer_inputs(self.backbone, per_layer_inputs_tokens, inputs_embeds) - per_layer_inputs = self.backbone.project_per_layer_inputs(inputs_embeds, per_layer_inputs) - - position_ids = torch.arange(inputs_embeds.shape[1], device=inputs_embeds.device).unsqueeze(0) - if self._native_merge_plan is not None: - causal_mask_mapping = _gemma4_build_standard_causal_mask_mapping( - create_causal_mask=self._create_causal_mask, - create_sliding_window_causal_mask=self._create_sliding_window_causal_mask, - config=self.backbone.config, - inputs_embeds=inputs_embeds, - attention_mask=attention_mask, - position_ids=position_ids, - ) - elif getattr(text_config, "use_bidirectional_attention", None) == "vision": - causal_mask_mapping = self._create_causal_mask_mapping( - self.multimodal_backbone.config, - inputs_embeds, - attention_mask, - None, - position_ids, - token_type_ids, - pixel_values, - is_training=self.training, - ) - else: - causal_mask_mapping = self._create_masks_for_generate( - self.multimodal_backbone.config, - inputs_embeds, - attention_mask, - None, - position_ids, - ) - return inputs_embeds, per_layer_inputs, causal_mask_mapping, position_ids - - def _forward_hidden_states( - self, - *, - input_ids: torch.Tensor, - attention_mask: torch.Tensor | None, - token_type_ids: torch.Tensor | None, - pixel_values: torch.Tensor | None, - pixel_position_ids: torch.Tensor | None, - input_features: torch.Tensor | None, - input_features_mask: torch.Tensor | None, - get_placeholder_mask: Callable[..., object], - get_image_features: Callable[..., object], - get_audio_features: Callable[..., object], - ) -> torch.Tensor: - inputs_embeds, per_layer_inputs, causal_mask_mapping, position_ids = self._prepare_text_backbone_inputs( - input_ids=input_ids, - attention_mask=attention_mask, - token_type_ids=token_type_ids, - pixel_values=pixel_values, - pixel_position_ids=pixel_position_ids, - input_features=input_features, - input_features_mask=input_features_mask, - get_placeholder_mask=get_placeholder_mask, - get_image_features=get_image_features, - get_audio_features=get_audio_features, - ) - return _gemma4_text_backbone_forward( - self.backbone, - inputs_embeds=inputs_embeds, - per_layer_inputs=per_layer_inputs, - causal_mask_mapping=causal_mask_mapping, - position_ids=position_ids, - ) - - def prime_static_multimodal_features(self, *bound_inputs: torch.Tensor | None) -> None: - kwargs = self._kwargs_from_bound_inputs(*bound_inputs) - input_ids = kwargs["input_ids"] - token_type_ids = kwargs.get("token_type_ids") - pixel_values = kwargs.get("pixel_values") - pixel_position_ids = kwargs.get("pixel_position_ids") - input_features = kwargs.get("input_features") - input_features_mask = kwargs.get("input_features_mask") - - multimodal_backbone = self.multimodal_backbone - get_placeholder_mask = getattr(multimodal_backbone, "get_placeholder_mask", None) - get_image_features = getattr(multimodal_backbone, "get_image_features", None) - get_audio_features = getattr(multimodal_backbone, "get_audio_features", None) - if not callable(get_placeholder_mask): - return - - with torch.no_grad(): - inputs_embeds = self.model.get_input_embeddings()(input_ids) - _, image_mask, audio_mask = _gemma4_get_placeholder_masks( - get_placeholder_mask, - token_type_ids=token_type_ids, - input_ids=input_ids, - inputs_embeds=inputs_embeds, - ) - - if pixel_values is not None and callable(get_image_features) and image_mask.any(): - image_features = self._compute_image_features( - pixel_values, - pixel_position_ids, - get_image_features, - ) - self._cached_image_features = image_features.detach() - else: - self._cached_image_features = self._cached_image_features.new_empty(0) - - if ( - input_features is not None - and input_features_mask is not None - and callable(get_audio_features) - and audio_mask.any() - ): - self._cached_audio_features = self._compute_audio_features( - input_features, - input_features_mask, - get_audio_features, - ).detach() - else: - self._cached_audio_features = self._cached_audio_features.new_empty(0) - - if self._cached_image_features.numel() > 0 and self._cached_audio_features.numel() > 0: - plan = _gemma4_build_native_merge_plan( - self.multimodal_backbone, - input_ids, - image_feature_count=_gemma4_feature_token_count(self._cached_image_features), - audio_feature_count=_gemma4_feature_token_count(self._cached_audio_features), - ) - self._native_merge_plan = plan - self._native_merge_pli_token_ids = torch.tensor( - plan.pli_token_ids, - dtype=torch.long, - ) - else: - self._native_merge_plan = None - self._native_merge_pli_token_ids = self._native_merge_pli_token_ids.new_empty(0) - - self._use_cached_multimodal_features = True - - def forward(self, *bound_inputs: torch.Tensor | None) -> torch.Tensor: - kwargs = self._kwargs_from_bound_inputs(*bound_inputs) - input_ids = kwargs["input_ids"] - attention_mask = kwargs.get("attention_mask") - token_type_ids = kwargs.get("token_type_ids") - pixel_values = kwargs.get("pixel_values") - pixel_position_ids = kwargs.get("pixel_position_ids") - input_features = kwargs.get("input_features") - input_features_mask = kwargs.get("input_features_mask") - - multimodal_backbone = self.multimodal_backbone - get_placeholder_mask = getattr(multimodal_backbone, "get_placeholder_mask", None) - get_image_features = getattr(multimodal_backbone, "get_image_features", None) - get_audio_features = getattr(multimodal_backbone, "get_audio_features", None) - lm_head = getattr(self.model, "lm_head", None) - if ( - not callable(get_placeholder_mask) - or not callable(get_image_features) - or not callable(get_audio_features) - or not callable(self._create_causal_mask_mapping) - or not callable(self._create_masks_for_generate) - or not isinstance(lm_head, torch.nn.Module) - ): - outputs = self.model( - return_dict=True, - use_cache=False, - **kwargs, - ) - logits = _extract_tensor_output(outputs, preferred_field="logits") - if self.last_token_logits_only and logits.ndim >= 3: - return _select_last_active_token(logits, attention_mask) - return logits - - hidden_states = self._forward_hidden_states( - input_ids=input_ids, - attention_mask=attention_mask, - token_type_ids=token_type_ids, - pixel_values=pixel_values, - pixel_position_ids=pixel_position_ids, - input_features=input_features, - input_features_mask=input_features_mask, - get_placeholder_mask=get_placeholder_mask, - get_image_features=get_image_features, - get_audio_features=get_audio_features, - ) - if self.last_token_logits_only: - hidden_states = _select_last_active_token(hidden_states, attention_mask) - return self._apply_final_logit_softcapping(lm_head(hidden_states)) - - def get_transpile_metadata(self): - sliding_window = getattr(self.backbone.config, "sliding_window", None) - layer_types = list(getattr(self.backbone.config, "layer_types", [])) - return { - "graph": { - **_transpile_graph_meta( - self.model, - adapter_family="gemma4", - adapter_type=type(self).__name__, - input_names=self.input_names, - ), - "task": self.metadata_task, - "num_hidden_layers": int(self.backbone.config.num_hidden_layers), - "layer_types": tuple(layer_types), - "sliding_window": None if sliding_window is None else int(sliding_window), - "last_token_logits_only": bool(self.last_token_logits_only), - } - } - - -_GEMMA4_DECODER_PIPELINE_IO_KEYS = ( - "inputs_embeds", - "per_layer_inputs", - "position_ids", -) - - -class _Gemma4MultimodalComponentBase(torch.nn.Module): - def __init__( - self, - model: torch.nn.Module, - *, - input_names: tuple[str, ...], - weights_dir: str | None = None, - native_merge_plan: _Gemma4NativeMergePlan | None = None, - native_image_soft_token_counts: tuple[int, ...] | None = None, - native_image_pool_shapes: tuple[tuple[int, int, int], ...] | None = None, - ): - super().__init__() - self.model = model - self.input_names = input_names - self.multimodal_backbone = model.model - model_backbone = model.model - self.backbone = getattr(model_backbone, "language_model", model_backbone) - self._native_merge_plan = native_merge_plan - self._native_image_soft_token_counts = native_image_soft_token_counts - self._native_image_pool_shapes = native_image_pool_shapes - self.register_buffer( - "_native_merge_pli_token_ids", - torch.tensor(native_merge_plan.pli_token_ids, dtype=torch.long) - if native_merge_plan is not None - else torch.empty(0, dtype=torch.long), - persistent=False, - ) - vision_post_proj_norm = _gemma4_load_vision_post_proj_norm(weights_dir) - if vision_post_proj_norm is None: - vision_post_proj_norm = torch.empty(0) - self.register_buffer("_cactus_vision_post_proj_norm", vision_post_proj_norm, persistent=False) - self._create_causal_mask_mapping = None - self._create_masks_for_generate = None - self._create_causal_mask = None - self._create_sliding_window_causal_mask = None - try: - from transformers.models.gemma4.modeling_gemma4 import create_causal_mask # type: ignore - from transformers.models.gemma4.modeling_gemma4 import create_causal_mask_mapping # type: ignore - from transformers.models.gemma4.modeling_gemma4 import create_masks_for_generate # type: ignore - from transformers.models.gemma4.modeling_gemma4 import create_sliding_window_causal_mask # type: ignore - - self._create_causal_mask = create_causal_mask - self._create_causal_mask_mapping = create_causal_mask_mapping - self._create_masks_for_generate = create_masks_for_generate - self._create_sliding_window_causal_mask = create_sliding_window_causal_mask - except Exception: - pass - self._capture_modules: list[tuple[torch.nn.Module, torch.dtype]] = [] - - def _base_graph_meta(self, *, adapter_type: str, input_names: tuple[str, ...]) -> dict[str, object]: - sliding_window = getattr(self.backbone.config, "sliding_window", None) - layer_types = list(getattr(self.backbone.config, "layer_types", [])) - return { - **_transpile_graph_meta( - self.model, - adapter_family="gemma4", - adapter_type=adapter_type, - input_names=input_names, - ), - "task": "multimodal_causal_lm_logits", - "num_hidden_layers": int(self.backbone.config.num_hidden_layers), - "layer_types": tuple(layer_types), - "sliding_window": None if sliding_window is None else int(sliding_window), - } - - def _compute_image_features( - self, - pixel_values: torch.Tensor, - pixel_position_ids: torch.Tensor | None, - ) -> torch.Tensor: - if _gemma4_can_use_native_like_vision_features(self.multimodal_backbone): - return _gemma4_compute_native_like_image_features( - self.multimodal_backbone, - pixel_values, - pixel_position_ids, - post_proj_norm_weight=self._cactus_vision_post_proj_norm, - image_soft_token_counts=self._native_image_soft_token_counts, - image_pool_shapes=self._native_image_pool_shapes, - ) - get_image_features = getattr(self.multimodal_backbone, "get_image_features", None) - vision_tower = getattr(self.multimodal_backbone, "vision_tower", None) - embed_vision = getattr(self.multimodal_backbone, "embed_vision", None) - vision_modules = _gemma4_vision_modules(self.multimodal_backbone) - if not callable(get_image_features): - raise TypeError("Gemma4 multimodal backbone is missing get_image_features") - if ( - pixel_values.device.type != "cpu" - or pixel_values.dtype != torch.float16 - or len(vision_modules) != 2 - or _torch_is_compiling() - ): - return get_image_features( - pixel_values, - pixel_position_ids, - None, - return_dict=True, - ).pooler_output - - vision_dtype = _module_floating_dtype(vision_tower) - embed_dtype = _module_floating_dtype(embed_vision) - if vision_dtype != torch.float16 and embed_dtype != torch.float16: - return get_image_features( - pixel_values, - pixel_position_ids, - None, - return_dict=True, - ).pooler_output - - with _temporary_cpu_float32_modules(vision_modules): - return get_image_features( - pixel_values.float(), - pixel_position_ids, - None, - return_dict=True, - ).pooler_output - - def _compute_audio_features( - self, - input_features: torch.Tensor, - input_features_mask: torch.Tensor, - ) -> torch.Tensor: - if _gemma4_can_use_native_like_audio_features(self.multimodal_backbone): - return _gemma4_compute_native_like_audio_features( - self.multimodal_backbone, - input_features, - input_features_mask, - ) - get_audio_features = getattr(self.multimodal_backbone, "get_audio_features", None) - if not callable(get_audio_features): - raise TypeError("Gemma4 multimodal backbone is missing get_audio_features") - audio_modules = _gemma4_audio_modules(self.multimodal_backbone) - if ( - input_features.device.type == "cpu" - and input_features.dtype == torch.float16 - and len(audio_modules) == 2 - and not _torch_is_compiling() - and any(_module_floating_dtype(module) == torch.float16 for module in audio_modules) - ): - with _temporary_cpu_float32_modules(audio_modules): - audio_output = get_audio_features(input_features.float(), input_features_mask, return_dict=True) - return _gemma4_strip_audio_padding(audio_output) - - audio_output = get_audio_features(input_features, input_features_mask, return_dict=True) - return _gemma4_strip_audio_padding(audio_output) - - def _modules_to_prepare_for_capture(self) -> tuple[torch.nn.Module, ...]: - return () - - def prepare_for_capture(self, **_: object) -> None: - self.restore_after_capture() - modules = self._modules_to_prepare_for_capture() - if not modules: - return - promoted: list[tuple[torch.nn.Module, torch.dtype]] = [] - for module in modules: - device = _module_device(module) - dtype = _module_floating_dtype(module) - if device is None or device.type != "cpu" or dtype is None or dtype == torch.float32: - continue - promoted.append((module, dtype)) - module.to(dtype=torch.float32) - self._capture_modules = promoted - - def restore_after_capture(self, **_: object) -> None: - for module, dtype in reversed(self._capture_modules): - module.to(dtype=dtype) - self._capture_modules.clear() - - def _prepare_decoder_inputs( - self, - *, - input_ids: torch.Tensor, - attention_mask: torch.Tensor | None, - token_type_ids: torch.Tensor, - image_features: torch.Tensor, - audio_features: torch.Tensor, - ) -> tuple[torch.Tensor, torch.Tensor, torch.LongTensor]: - if self._native_merge_plan is not None: - inputs_embeds, per_layer_inputs_tokens = _gemma4_apply_native_merge_plan( - self.model, - input_ids=input_ids, - image_features=image_features, - audio_features=audio_features, - merge_plan=self._native_merge_plan, - pli_token_ids=self._native_merge_pli_token_ids, - ) - attention_mask = _gemma4_remap_sequence_tensor( - attention_mask, - merge_plan=self._native_merge_plan, - ) - token_type_ids = _gemma4_remap_sequence_tensor( - token_type_ids, - merge_plan=self._native_merge_plan, - ) - if token_type_ids is None: - raise RuntimeError("Gemma4 native merge requires token_type_ids for multimodal attention masking") - else: - get_placeholder_mask = getattr(self.multimodal_backbone, "get_placeholder_mask", None) - if not callable(get_placeholder_mask): - raise TypeError("Gemma4 multimodal backbone is missing get_placeholder_mask") - - inputs_embeds = self.model.get_input_embeddings()(input_ids) - text_mask, image_mask, audio_mask = _gemma4_get_placeholder_masks( - get_placeholder_mask, - token_type_ids=token_type_ids, - input_ids=input_ids, - inputs_embeds=inputs_embeds, - ) - inputs_embeds = inputs_embeds.masked_scatter( - image_mask.unsqueeze(-1).expand_as(inputs_embeds), - image_features.to(inputs_embeds.device, inputs_embeds.dtype), - ) - inputs_embeds = inputs_embeds.masked_scatter( - audio_mask.unsqueeze(-1).expand_as(inputs_embeds), - audio_features.to(inputs_embeds.device, inputs_embeds.dtype), - ) - per_layer_inputs_tokens = input_ids * text_mask.to(dtype=input_ids.dtype) - - per_layer_inputs: torch.Tensor - text_config = _gemma4_text_config(self.multimodal_backbone) - if getattr(text_config, "hidden_size_per_layer_input", None): - per_layer_inputs = _gemma4_get_per_layer_inputs(self.backbone, per_layer_inputs_tokens, inputs_embeds) - per_layer_inputs = self.backbone.project_per_layer_inputs(inputs_embeds, per_layer_inputs) - else: - per_layer_inputs = inputs_embeds.new_empty((inputs_embeds.shape[0], inputs_embeds.shape[1], 0, 0)) - - position_ids = torch.arange(inputs_embeds.shape[1], device=inputs_embeds.device).unsqueeze(0) - return ( - inputs_embeds, - per_layer_inputs, - position_ids, - ) - - def _prepare_text_decoder_step_inputs( - self, - *, - input_ids: torch.Tensor, - position_ids: torch.LongTensor, - ) -> tuple[torch.Tensor, torch.Tensor, torch.LongTensor]: - embedding = self.model.get_input_embeddings() - if not isinstance(embedding, torch.nn.Module): - raise TypeError("Gemma4 model is missing input embeddings") - inputs_embeds = embedding(input_ids) - - model_config = getattr(self.model, "config", None) - text_config = getattr(model_config, "text_config", None) - hidden_scale = float(getattr(model_config, "hidden_size", 0) or 0) - if hidden_scale <= 0.0: - hidden_scale = float(getattr(text_config, "hidden_size", 0) or 0) - if hidden_scale <= 0.0: - hidden_scale = float(inputs_embeds.shape[-1]) - text_extra_scale = _gemma4_text_embedding_scale(embedding, hidden_scale ** 0.5) - if text_extra_scale != 1.0: - inputs_embeds = inputs_embeds * text_extra_scale - - text_config = _gemma4_text_config(self.multimodal_backbone) - if getattr(text_config, "hidden_size_per_layer_input", None): - per_layer_inputs = _gemma4_get_per_layer_inputs(self.backbone, input_ids, inputs_embeds) - per_layer_inputs = self.backbone.project_per_layer_inputs(inputs_embeds, per_layer_inputs) - else: - per_layer_inputs = inputs_embeds.new_empty((inputs_embeds.shape[0], inputs_embeds.shape[1], 0, 0)) - return inputs_embeds, per_layer_inputs, position_ids - - -class Gemma4VisionEncoderAdapter(_Gemma4MultimodalComponentBase): - def __init__( - self, - model: torch.nn.Module, - *, - weights_dir: str | None = None, - native_image_soft_token_counts: tuple[int, ...] | None = None, - native_image_pool_shapes: tuple[tuple[int, int, int], ...] | None = None, - ): - super().__init__( - model, - input_names=("pixel_values", "pixel_position_ids"), - weights_dir=weights_dir, - native_image_soft_token_counts=native_image_soft_token_counts, - native_image_pool_shapes=native_image_pool_shapes, - ) - - def _modules_to_prepare_for_capture(self) -> tuple[torch.nn.Module, ...]: - return () - - def forward(self, pixel_values: torch.Tensor, pixel_position_ids: torch.Tensor | None) -> torch.Tensor: - return self._compute_image_features(pixel_values, pixel_position_ids) - - def get_transpile_metadata(self): - return { - "graph": self._base_graph_meta( - adapter_type=type(self).__name__, - input_names=("pixel_values", "pixel_position_ids"), - ), - } - - -class Gemma4AudioEncoderAdapter(_Gemma4MultimodalComponentBase): - def __init__(self, model: torch.nn.Module, *, weights_dir: str | None = None): - super().__init__(model, input_names=("input_features", "input_features_mask"), weights_dir=weights_dir) - - def _modules_to_prepare_for_capture(self) -> tuple[torch.nn.Module, ...]: - return () - - def forward(self, input_features: torch.Tensor, input_features_mask: torch.Tensor) -> torch.Tensor: - return self._compute_audio_features(input_features, input_features_mask) - - def get_transpile_metadata(self): - return { - "graph": self._base_graph_meta( - adapter_type=type(self).__name__, - input_names=("input_features", "input_features_mask"), - ), - } - - -class Gemma4LMEncoderAdapter(_Gemma4MultimodalComponentBase): - def __init__( - self, - model: torch.nn.Module, - *, - weights_dir: str | None = None, - native_merge_plan: _Gemma4NativeMergePlan | None = None, - ): - super().__init__( - model, - input_names=("input_ids", "attention_mask", "token_type_ids", "image_features", "audio_features"), - weights_dir=weights_dir, - native_merge_plan=native_merge_plan, - ) - - def forward( - self, - input_ids: torch.Tensor, - attention_mask: torch.Tensor | None, - token_type_ids: torch.Tensor, - image_features: torch.Tensor, - audio_features: torch.Tensor, - ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.LongTensor]: - return self._prepare_decoder_inputs( - input_ids=input_ids, - attention_mask=attention_mask, - token_type_ids=token_type_ids, - image_features=image_features, - audio_features=audio_features, - ) - - def get_transpile_metadata(self): - return { - "graph": self._base_graph_meta( - adapter_type=type(self).__name__, - input_names=("input_ids", "attention_mask", "token_type_ids", "image_features", "audio_features"), - ), - } - - -class Gemma4LMEncoderStepAdapter(_Gemma4MultimodalComponentBase): - def __init__(self, model: torch.nn.Module, *, weights_dir: str | None = None): - super().__init__(model, input_names=("input_ids", "position_ids"), weights_dir=weights_dir) - - def forward( - self, - input_ids: torch.Tensor, - position_ids: torch.LongTensor, - ) -> tuple[torch.Tensor, torch.Tensor, torch.LongTensor]: - return self._prepare_text_decoder_step_inputs( - input_ids=input_ids, - position_ids=position_ids, - ) - - def get_transpile_metadata(self): - return { - "graph": self._base_graph_meta( - adapter_type=type(self).__name__, - input_names=("input_ids", "position_ids"), - ), - } - - -class Gemma4LMEncoderTextChunkAdapter(Gemma4LMEncoderStepAdapter): - pass - - -class Gemma4LMEncoderMediaStepAdapter(_Gemma4MultimodalComponentBase): - def __init__(self, model: torch.nn.Module, *, weights_dir: str | None = None): - super().__init__( - model, - input_names=("inputs_embeds", "input_ids", "position_ids"), - weights_dir=weights_dir, - ) - - def forward( - self, - inputs_embeds: torch.Tensor, - input_ids: torch.Tensor, - position_ids: torch.LongTensor, - ) -> tuple[torch.Tensor, torch.Tensor, torch.LongTensor]: - text_config = _gemma4_text_config(self.multimodal_backbone) - if getattr(text_config, "hidden_size_per_layer_input", None): - per_layer_inputs = _gemma4_get_per_layer_inputs(self.backbone, input_ids, inputs_embeds) - per_layer_inputs = self.backbone.project_per_layer_inputs(inputs_embeds, per_layer_inputs) - else: - per_layer_inputs = inputs_embeds.new_empty((inputs_embeds.shape[0], inputs_embeds.shape[1], 0, 0)) - return inputs_embeds, per_layer_inputs, position_ids - - def get_transpile_metadata(self): - return { - "graph": self._base_graph_meta( - adapter_type=type(self).__name__, - input_names=("inputs_embeds", "input_ids", "position_ids"), - ), - } - - -class Gemma4LMEncoderMediaChunkAdapter(Gemma4LMEncoderMediaStepAdapter): - pass - - -class Gemma4DecoderAdapter(_Gemma4MultimodalComponentBase): - def __init__(self, model: torch.nn.Module, *, weights_dir: str | None = None): - super().__init__(model, input_names=_GEMMA4_DECODER_PIPELINE_IO_KEYS, weights_dir=weights_dir) - - def forward( - self, - inputs_embeds: torch.Tensor, - per_layer_inputs: torch.Tensor, - position_ids: torch.LongTensor, - ) -> torch.Tensor: - normalized_per_layer_inputs = per_layer_inputs - if normalized_per_layer_inputs.numel() == 0: - normalized_per_layer_inputs = None - attention_mask = torch.ones( - position_ids.shape, - dtype=torch.long, - device=position_ids.device, - ) - causal_mask_mapping = _gemma4_build_standard_causal_mask_mapping( - create_causal_mask=self._create_causal_mask, - create_sliding_window_causal_mask=self._create_sliding_window_causal_mask, - config=self.backbone.config, - inputs_embeds=inputs_embeds, - attention_mask=attention_mask, - position_ids=position_ids, - ) - hidden_states = _gemma4_text_backbone_forward( - self.backbone, - inputs_embeds=inputs_embeds, - per_layer_inputs=normalized_per_layer_inputs, - causal_mask_mapping=causal_mask_mapping, - position_ids=position_ids, - ) - logits = self.model.lm_head(hidden_states[:, -1:, :]) - return _gemma4_apply_final_logit_softcapping(self.model, logits) - - def get_transpile_metadata(self): - return { - "graph": self._base_graph_meta( - adapter_type=type(self).__name__, - input_names=_GEMMA4_DECODER_PIPELINE_IO_KEYS, - ) - } - - -class Gemma4DecoderStepAdapter(Gemma4DecoderAdapter): - def forward( - self, - inputs_embeds: torch.Tensor, - per_layer_inputs: torch.Tensor, - position_ids: torch.LongTensor, - ) -> tuple[torch.Tensor, torch.Tensor]: - normalized_per_layer_inputs = per_layer_inputs - if normalized_per_layer_inputs.numel() == 0: - normalized_per_layer_inputs = None - attention_mask = torch.ones( - position_ids.shape, - dtype=torch.long, - device=position_ids.device, - ) - causal_mask_mapping = _gemma4_build_standard_causal_mask_mapping( - create_causal_mask=self._create_causal_mask, - create_sliding_window_causal_mask=self._create_sliding_window_causal_mask, - config=self.backbone.config, - inputs_embeds=inputs_embeds, - attention_mask=attention_mask, - position_ids=position_ids, - ) - hidden_states, probe_hidden = _gemma4_text_backbone_forward( - self.backbone, - inputs_embeds=inputs_embeds, - per_layer_inputs=normalized_per_layer_inputs, - causal_mask_mapping=causal_mask_mapping, - position_ids=position_ids, - capture_layer_index=28, - ) - logits = self.model.lm_head(hidden_states[:, -1:, :]) - logits = _gemma4_apply_final_logit_softcapping(self.model, logits) - return logits, probe_hidden[:, -1:, :] - - -class Gemma4DecoderPrefillChunkAdapter(Gemma4DecoderAdapter): - def forward( - self, - inputs_embeds: torch.Tensor, - per_layer_inputs: torch.Tensor, - position_ids: torch.LongTensor, - ) -> torch.Tensor: - normalized_per_layer_inputs = per_layer_inputs - if normalized_per_layer_inputs.numel() == 0: - normalized_per_layer_inputs = None - attention_mask = torch.ones( - position_ids.shape, - dtype=torch.long, - device=position_ids.device, - ) - causal_mask_mapping = _gemma4_build_standard_causal_mask_mapping( - create_causal_mask=self._create_causal_mask, - create_sliding_window_causal_mask=self._create_sliding_window_causal_mask, - config=self.backbone.config, - inputs_embeds=inputs_embeds, - attention_mask=attention_mask, - position_ids=position_ids, - ) - num_layers = int(getattr(self.backbone.config, "num_hidden_layers", len(self.backbone.layers))) - num_shared_layers = int(getattr(self.backbone.config, "num_kv_shared_layers", 0) or 0) - first_shared_layer = max(0, num_layers - num_shared_layers) - if first_shared_layer <= 0 or first_shared_layer >= num_layers or int(inputs_embeds.shape[1]) <= 1: - hidden_states = _gemma4_text_backbone_forward( - self.backbone, - inputs_embeds=inputs_embeds, - per_layer_inputs=normalized_per_layer_inputs, - causal_mask_mapping=causal_mask_mapping, - position_ids=position_ids, - ) - else: - shared_kv_states: dict[int, tuple[torch.Tensor, torch.Tensor]] = {} - hidden_states = _gemma4_text_backbone_forward( - self.backbone, - inputs_embeds=inputs_embeds, - per_layer_inputs=normalized_per_layer_inputs, - causal_mask_mapping=causal_mask_mapping, - position_ids=position_ids, - layer_end=first_shared_layer, - apply_norm=False, - shared_kv_states=shared_kv_states, - ) - tail_masks = { - key: value[:, :, -1:, :] if value.ndim == 4 else value - for key, value in causal_mask_mapping.items() - } - tail_per_layer_inputs = ( - None - if normalized_per_layer_inputs is None - else normalized_per_layer_inputs[:, -1:, :, :] - ) - hidden_states = _gemma4_text_backbone_forward( - self.backbone, - inputs_embeds=hidden_states[:, -1:, :], - per_layer_inputs=tail_per_layer_inputs, - causal_mask_mapping=tail_masks, - position_ids=position_ids[:, -1:], - layer_start=first_shared_layer, - apply_norm=True, - shared_kv_states=shared_kv_states, - ) - logits = self.model.lm_head(hidden_states[:, -1:, :]) - return _gemma4_apply_final_logit_softcapping(self.model, logits) - - -class Gemma4DecoderEmbedChunkAdapter(Gemma4DecoderAdapter): - """Embedding-readout variant of the prefill chunk: runs the full backbone over - ALL tokens with final norm (no shared-KV-tail last-token shortcut) and returns - the full-sequence last hidden state (no lm_head).""" - - def forward( - self, - inputs_embeds: torch.Tensor, - per_layer_inputs: torch.Tensor, - position_ids: torch.LongTensor, - ) -> torch.Tensor: - normalized_per_layer_inputs = per_layer_inputs - if normalized_per_layer_inputs.numel() == 0: - normalized_per_layer_inputs = None - attention_mask = torch.ones( - position_ids.shape, - dtype=torch.long, - device=position_ids.device, - ) - causal_mask_mapping = _gemma4_build_standard_causal_mask_mapping( - create_causal_mask=self._create_causal_mask, - create_sliding_window_causal_mask=self._create_sliding_window_causal_mask, - config=self.backbone.config, - inputs_embeds=inputs_embeds, - attention_mask=attention_mask, - position_ids=position_ids, - ) - return _gemma4_text_backbone_forward( - self.backbone, - inputs_embeds=inputs_embeds, - per_layer_inputs=normalized_per_layer_inputs, - causal_mask_mapping=causal_mask_mapping, - position_ids=position_ids, - ) - - -def _build_gemma3_causal_lm_component_specs( - model: torch.nn.Module, - *, - named_tensors: dict[str, torch.Tensor], - weights_dir: str | None = None, - components: tuple[str, ...] | None = None, - cache_context_length: str | int | None = None, -) -> list[ComponentModuleSpec] | None: - input_ids = named_tensors.get("input_ids") - if input_ids is None: - return None - pad_token_id = _resolve_model_pad_token_id(model) - requested = tuple(components or ( - "decoder", - "lm_encoder_step", - "lm_encoder_text_chunk", - "decoder_prefill_chunk", - "decoder_step", - )) - requested_set = set(requested) - common_graph_meta = { - "weights_dir": weights_dir, - "task": "causal_lm_logits", - "adapter_family": "gemma3", - } - metadata = {"family": "gemma3", "task": "causal_lm_logits"} - - specs: list[ComponentModuleSpec] = [] - if "decoder" in requested_set: - decoder = Gemma3CausalLMLogitsAdapter(model, pad_token_id=pad_token_id).eval() - specs.append(ComponentModuleSpec( - component="decoder", - module=decoder, - example_inputs=(input_ids,), - input_keys=("input_ids",), - output_keys=("logits",), - graph_meta={**common_graph_meta, "component": "decoder"}, - metadata=metadata, - )) - - cached_components = {"lm_encoder_step", "lm_encoder_text_chunk", "decoder_prefill_chunk", "decoder_step"} - if not (cached_components & requested_set): - return specs - - lm_encoder_step = Gemma3LMEncoderStepAdapter(model).eval() - lm_encoder_text_chunk = Gemma3LMEncoderStepAdapter(model).eval() - decoder_step = Gemma3EmbedsCausalLMStepAdapter(model).eval() - decoder_prefill_chunk = Gemma3EmbedsCausalLMPrefillChunkAdapter(model).eval() - max_cache_seq_len = _max_cache_seq_len(model, input_ids, cache_context_length, fallback_extra_tokens=512) - cached_graph_meta = { - "use_internal_kv_cache": True, - "max_cache_seq_len": max_cache_seq_len, - "cache_sink_size": 4, - } - prefill_chunk_size = max(2, int(os.environ.get("CACTUS_GEMMA3_PREFILL_CHUNK", "128") or "128")) - - step_input_ids = input_ids[:, :1] - step_position_ids = torch.zeros_like(step_input_ids, dtype=torch.int64) - chunk_input_ids = _tile_to_length(input_ids, prefill_chunk_size) - chunk_position_ids = torch.arange( - prefill_chunk_size, dtype=torch.int64, device=input_ids.device, - ).unsqueeze(0) - with torch.no_grad(): - step_embeds, step_pos_out = lm_encoder_step(step_input_ids, step_position_ids) - chunk_embeds, chunk_pos_out = lm_encoder_text_chunk(chunk_input_ids, chunk_position_ids) - - if "lm_encoder_step" in requested_set: - specs.append(ComponentModuleSpec( - component="lm_encoder_step", - module=lm_encoder_step, - example_inputs=(step_input_ids, step_position_ids), - input_keys=("input_ids", "position_ids"), - output_keys=("inputs_embeds", "position_ids"), - graph_meta={**common_graph_meta, "component": "lm_encoder_step"}, - metadata=metadata, - )) - if "lm_encoder_text_chunk" in requested_set: - specs.append(ComponentModuleSpec( - component="lm_encoder_text_chunk", - module=lm_encoder_text_chunk, - example_inputs=(chunk_input_ids, chunk_position_ids), - input_keys=("input_ids", "position_ids"), - output_keys=("inputs_embeds", "position_ids"), - graph_meta={ - **common_graph_meta, - "component": "lm_encoder_text_chunk", - "encoder_chunk_size": prefill_chunk_size, - }, - metadata=metadata, - )) - if "decoder_prefill_chunk" in requested_set: - specs.append(ComponentModuleSpec( - component="decoder_prefill_chunk", - module=decoder_prefill_chunk, - example_inputs=(chunk_embeds, chunk_pos_out), - input_keys=("inputs_embeds", "position_ids"), - output_keys=("logits",), - graph_meta={ - **common_graph_meta, - **cached_graph_meta, - "component": "decoder_prefill_chunk", - "prefill_chunk_size": prefill_chunk_size, - }, - metadata=metadata, - )) - if "decoder_step" in requested_set: - specs.append(ComponentModuleSpec( - component="decoder_step", - module=decoder_step, - example_inputs=(step_embeds, step_pos_out), - input_keys=("inputs_embeds", "position_ids"), - output_keys=("logits",), - graph_meta={ - **common_graph_meta, - **cached_graph_meta, - "component": "decoder_step", - }, - metadata=metadata, - )) - return specs - - -def _build_gemma4_multimodal_component_specs( - model: torch.nn.Module, - *, - named_tensors: dict[str, torch.Tensor], - weights_dir: str | None, - components: tuple[str, ...] | None = None, - cache_context_length: str | int | None = None, - dynamic_batch: bool = False, - max_slots: int = 1, -) -> list[ComponentModuleSpec]: - pixel_values = named_tensors["pixel_values"] - pixel_position_ids = named_tensors.get("pixel_position_ids") - input_features = named_tensors["input_features"] - input_features_mask = named_tensors["input_features_mask"] - input_ids = named_tensors["input_ids"] - attention_mask = named_tensors.get("attention_mask") - token_type_ids = named_tensors["token_type_ids"] - - requested_components = tuple(components or ("vision_encoder", "audio_encoder", "lm_encoder", "decoder")) - requested_set = set(requested_components) - if not requested_set: - return [] - - expanded_components: list[str] = [] - - def _require(component: str) -> None: - if component not in expanded_components: - expanded_components.append(component) - - if "decoder" in requested_set: - _require("vision_encoder") - _require("audio_encoder") - _require("lm_encoder") - _require("decoder_prefill_chunk") - _require("decoder_embed_chunk") - _require("lm_encoder_text_chunk") - _require("lm_encoder_media_chunk") - _require("lm_encoder_step") - _require("lm_encoder_media_step") - _require("decoder_step") - if "lm_encoder" in requested_set: - _require("vision_encoder") - _require("audio_encoder") - _require("lm_encoder") - if "decoder_prefill_chunk" in requested_set: - _require("decoder_prefill_chunk") - if "vision_encoder" in requested_set: - _require("vision_encoder") - if "audio_encoder" in requested_set: - _require("audio_encoder") - if "lm_encoder_step" in requested_set: - _require("lm_encoder_step") - if "lm_encoder_text_chunk" in requested_set: - _require("lm_encoder_text_chunk") - if "lm_encoder_media_step" in requested_set: - _require("lm_encoder_media_step") - if "lm_encoder_media_chunk" in requested_set: - _require("lm_encoder_media_chunk") - if "decoder_step" in requested_set: - _require("lm_encoder_step") - _require("decoder_step") - - vision_tower = getattr(getattr(model, "model", model), "vision_tower", None) - pooling_kernel_size = int(_module_or_config_attr(vision_tower, "pooling_kernel_size", 3) or 3) - native_image_soft_token_counts = _gemma4_static_image_soft_token_counts( - pixel_position_ids, - pooling_kernel_size=pooling_kernel_size, - ) - native_image_pool_shapes = _gemma4_static_image_pool_shapes( - pixel_position_ids, - pooling_kernel_size=pooling_kernel_size, - ) - - vision_encoder = Gemma4VisionEncoderAdapter( - model, - weights_dir=weights_dir, - native_image_soft_token_counts=native_image_soft_token_counts, - native_image_pool_shapes=native_image_pool_shapes, - ).eval() - vision_encoder_npu = Gemma4VisionEncoderAdapter( - model, - weights_dir=weights_dir, - native_image_soft_token_counts=native_image_soft_token_counts, - native_image_pool_shapes=None, - ).eval() - audio_encoder = Gemma4AudioEncoderAdapter(model, weights_dir=weights_dir).eval() - - image_features: torch.Tensor | None = None - audio_features: torch.Tensor | None = None - decoder_inputs: tuple[torch.Tensor, ...] | None = None - prefill_decoder_inputs: tuple[torch.Tensor, ...] | None = None - decoder_step_inputs: tuple[torch.Tensor, ...] | None = None - step_input_ids: torch.Tensor | None = None - step_position_ids: torch.Tensor | None = None - media_step_embeds: torch.Tensor | None = None - native_merge_plan: _Gemma4NativeMergePlan | None = None - - with torch.no_grad(): - if "vision_encoder" in expanded_components and ( - "lm_encoder" in expanded_components - or "decoder" in expanded_components - ): - image_features = vision_encoder(pixel_values, pixel_position_ids) - if "audio_encoder" in expanded_components and ( - "lm_encoder" in expanded_components - or "decoder" in expanded_components - ): - audio_features = audio_encoder(input_features, input_features_mask) - if "lm_encoder" in expanded_components or "decoder" in expanded_components: - if image_features is None: - image_features = vision_encoder(pixel_values, pixel_position_ids) - if audio_features is None: - audio_features = audio_encoder(input_features, input_features_mask) - native_merge_plan = _gemma4_build_native_merge_plan( - getattr(model, "model", model), - input_ids, - image_feature_count=_gemma4_feature_token_count(image_features), - audio_feature_count=_gemma4_feature_token_count(audio_features), - ) - lm_encoder = Gemma4LMEncoderAdapter( - model, - weights_dir=weights_dir, - native_merge_plan=native_merge_plan, - ).eval() - encoder_chunk_size = max(1, int(os.environ.get("CACTUS_GEMMA4_ENCODER_CHUNK", "128") or "128")) - encoder_chunk_size = min(encoder_chunk_size, int(input_ids.shape[1])) - lm_encoder_step = Gemma4LMEncoderStepAdapter(model, weights_dir=weights_dir).eval() - lm_encoder_text_chunk = Gemma4LMEncoderTextChunkAdapter(model, weights_dir=weights_dir).eval() - lm_encoder_media_step = Gemma4LMEncoderMediaStepAdapter(model, weights_dir=weights_dir).eval() - lm_encoder_media_chunk = Gemma4LMEncoderMediaChunkAdapter(model, weights_dir=weights_dir).eval() - decoder = Gemma4DecoderAdapter(model, weights_dir=weights_dir).eval() - decoder_prefill = Gemma4DecoderPrefillChunkAdapter(model, weights_dir=weights_dir).eval() - decoder_embed = Gemma4DecoderEmbedChunkAdapter(model, weights_dir=weights_dir).eval() - decoder_step = Gemma4DecoderStepAdapter(model, weights_dir=weights_dir).eval() - - with torch.no_grad(): - if "decoder" in expanded_components: - decoder_inputs = lm_encoder( - input_ids, - attention_mask, - token_type_ids, - image_features, - audio_features, - ) - if "lm_encoder_step" in expanded_components or "decoder_step" in expanded_components: - step_input_ids = input_ids[:, -1:].contiguous() - step_position_ids = torch.full( - (int(input_ids.shape[0]), 1), - max(0, int(input_ids.shape[1]) - 1), - dtype=torch.long, - device=input_ids.device, - ) - decoder_step_inputs = lm_encoder_step(step_input_ids, step_position_ids) - if "lm_encoder_media_step" in expanded_components: - if audio_features is not None and int(audio_features.shape[1]) > 0: - media_step_embeds = audio_features[:, :1, :].contiguous() - elif image_features is not None and int(image_features.shape[1]) > 0: - media_step_embeds = image_features[:, :1, :].contiguous() - else: - text_config = _gemma4_text_config(getattr(model, "model", model)) - hidden_size = int(getattr(text_config, "hidden_size", 0) or 0) - if hidden_size <= 0: - raise RuntimeError("Gemma4 lm_encoder_media_step spec could not infer hidden size") - dtype = _module_floating_dtype(model) or torch.float16 - media_step_embeds = torch.zeros( - (int(input_ids.shape[0]), 1, hidden_size), - dtype=dtype, - device=input_ids.device, - ) - chunk_input_ids = input_ids[:, :encoder_chunk_size].contiguous() - chunk_position_ids = torch.arange( - encoder_chunk_size, - dtype=torch.long, - device=input_ids.device, - ).unsqueeze(0).expand(int(input_ids.shape[0]), -1).contiguous() - if "decoder_prefill_chunk" in expanded_components and decoder_inputs is None: - prefill_decoder_inputs = lm_encoder_text_chunk(chunk_input_ids, chunk_position_ids) - if media_step_embeds is not None: - hidden_size = int(media_step_embeds.shape[-1]) - media_chunk_embeds = torch.zeros( - (int(input_ids.shape[0]), encoder_chunk_size, hidden_size), - dtype=media_step_embeds.dtype, - device=media_step_embeds.device, - ) - media_chunk_embeds[:, :1, :] = media_step_embeds - else: - text_config = _gemma4_text_config(getattr(model, "model", model)) - hidden_size = int(getattr(text_config, "hidden_size", 0) or 0) - if hidden_size <= 0: - hidden_size = int(chunk_input_ids.shape[-1]) - dtype = _module_floating_dtype(model) or torch.float16 - media_chunk_embeds = torch.zeros( - (int(input_ids.shape[0]), encoder_chunk_size, hidden_size), - dtype=dtype, - device=input_ids.device, - ) - - common_graph_meta = { - "weights_dir": weights_dir, - "task": "multimodal_causal_lm_logits", - "adapter_family": "gemma4", - } - specs: list[ComponentModuleSpec] = [] - if "vision_encoder" in expanded_components: - specs.append(ComponentModuleSpec( - component="vision_encoder", - module=vision_encoder, - example_inputs=(pixel_values, pixel_position_ids), - input_keys=("pixel_values", "pixel_position_ids"), - output_keys=("image_features",), - graph_meta={**common_graph_meta, "component": "vision_encoder"}, - metadata={"family": "gemma4", "task": "multimodal_causal_lm_logits"}, - npu_module=vision_encoder_npu, - npu_runtime_input_count=2, - npu_reparam=_gemma4_npu_projection_reparam, - )) - if "audio_encoder" in expanded_components: - specs.append(ComponentModuleSpec( - component="audio_encoder", - module=audio_encoder, - example_inputs=(input_features, input_features_mask), - input_keys=("input_features", "input_features_mask"), - output_keys=("audio_features",), - graph_meta={**common_graph_meta, "component": "audio_encoder"}, - metadata={"family": "gemma4", "task": "multimodal_causal_lm_logits"}, - npu_reparam=_gemma4_npu_projection_reparam, - )) - if "lm_encoder" in expanded_components: - if image_features is None or audio_features is None: - raise RuntimeError("Gemma4 lm_encoder spec requires precomputed image/audio features") - specs.append(ComponentModuleSpec( - component="lm_encoder", - module=lm_encoder, - example_inputs=(input_ids, attention_mask, token_type_ids, image_features, audio_features), - input_keys=("input_ids", "attention_mask", "token_type_ids", "image_features", "audio_features"), - output_keys=_GEMMA4_DECODER_PIPELINE_IO_KEYS, - graph_meta={**common_graph_meta, "component": "lm_encoder"}, - metadata={"family": "gemma4", "task": "multimodal_causal_lm_logits"}, - )) - cache_seq_len = _max_cache_seq_len(model, input_ids, cache_context_length, fallback_extra_tokens=256) - prefill_chunk_size = max(1, int(os.environ.get("CACTUS_GEMMA4_PREFILL_CHUNK", "128") or "128")) - prefill_chunk_size = min(prefill_chunk_size, int(input_ids.shape[1])) - if "decoder" in expanded_components: - if decoder_inputs is None: - raise RuntimeError("Gemma4 decoder spec requires precomputed decoder inputs") - specs.append(ComponentModuleSpec( - component="decoder", - module=decoder, - example_inputs=tuple(decoder_inputs), - input_keys=_GEMMA4_DECODER_PIPELINE_IO_KEYS, - output_keys=("logits",), - graph_meta={**common_graph_meta, "component": "decoder"}, - metadata={"family": "gemma4", "task": "multimodal_causal_lm_logits"}, - )) - if "decoder_prefill_chunk" in expanded_components: - if decoder_inputs is not None: - chunk_inputs = tuple( - value[:, :prefill_chunk_size, ...].contiguous() - if value.ndim >= 2 - else value - for value in decoder_inputs - ) - elif prefill_decoder_inputs is not None: - prefill_chunk_size = int(prefill_decoder_inputs[0].shape[1]) - chunk_inputs = tuple(prefill_decoder_inputs) - else: - raise RuntimeError("Gemma4 decoder_prefill_chunk spec requires precomputed decoder inputs") - specs.append(ComponentModuleSpec( - component="decoder_prefill_chunk", - module=decoder_prefill, - example_inputs=chunk_inputs, - input_keys=_GEMMA4_DECODER_PIPELINE_IO_KEYS, - output_keys=("logits",), - graph_meta={ - **common_graph_meta, - "component": "decoder_prefill_chunk", - "use_internal_kv_cache": True, - "max_cache_seq_len": cache_seq_len, - "cache_sink_size": 4, - "prefill_chunk_size": prefill_chunk_size, - }, - metadata={"family": "gemma4", "task": "multimodal_causal_lm_logits"}, - )) - if "decoder_embed_chunk" in expanded_components: - if decoder_inputs is not None: - embed_chunk_inputs = tuple( - value[:, :prefill_chunk_size, ...].contiguous() - if value.ndim >= 2 - else value - for value in decoder_inputs - ) - elif prefill_decoder_inputs is not None: - embed_chunk_inputs = tuple(prefill_decoder_inputs) - else: - raise RuntimeError("Gemma4 decoder_embed_chunk spec requires precomputed decoder inputs") - specs.append(ComponentModuleSpec( - component="decoder_embed_chunk", - module=decoder_embed, - example_inputs=embed_chunk_inputs, - input_keys=_GEMMA4_DECODER_PIPELINE_IO_KEYS, - output_keys=("last_hidden_state",), - graph_meta={ - **common_graph_meta, - "component": "decoder_embed_chunk", - "use_internal_kv_cache": True, - "max_cache_seq_len": cache_seq_len, - "cache_sink_size": 4, - "prefill_chunk_size": prefill_chunk_size, - }, - metadata={"family": "gemma4", "task": "multimodal_causal_lm_logits"}, - )) - if "lm_encoder_step" in expanded_components: - if step_input_ids is None or step_position_ids is None: - raise RuntimeError("Gemma4 lm_encoder_step spec requires step token inputs") - specs.append(ComponentModuleSpec( - component="lm_encoder_step", - module=lm_encoder_step, - example_inputs=(step_input_ids, step_position_ids), - input_keys=("input_ids", "position_ids"), - output_keys=_GEMMA4_DECODER_PIPELINE_IO_KEYS, - graph_meta={**common_graph_meta, "component": "lm_encoder_step"}, - metadata={"family": "gemma4", "task": "multimodal_causal_lm_logits"}, - dynamic_batch_axis=0 if dynamic_batch else None, - )) - if "lm_encoder_text_chunk" in expanded_components: - specs.append(ComponentModuleSpec( - component="lm_encoder_text_chunk", - module=lm_encoder_text_chunk, - example_inputs=(chunk_input_ids, chunk_position_ids), - input_keys=("input_ids", "position_ids"), - output_keys=_GEMMA4_DECODER_PIPELINE_IO_KEYS, - graph_meta={ - **common_graph_meta, - "component": "lm_encoder_text_chunk", - "encoder_chunk_size": encoder_chunk_size, - }, - metadata={"family": "gemma4", "task": "multimodal_causal_lm_logits"}, - )) - if "lm_encoder_media_step" in expanded_components: - if media_step_embeds is None or step_input_ids is None or step_position_ids is None: - raise RuntimeError("Gemma4 lm_encoder_media_step spec requires step token inputs") - specs.append(ComponentModuleSpec( - component="lm_encoder_media_step", - module=lm_encoder_media_step, - example_inputs=(media_step_embeds, torch.zeros_like(step_input_ids), step_position_ids), - input_keys=("inputs_embeds", "input_ids", "position_ids"), - output_keys=_GEMMA4_DECODER_PIPELINE_IO_KEYS, - graph_meta={**common_graph_meta, "component": "lm_encoder_media_step"}, - metadata={"family": "gemma4", "task": "multimodal_causal_lm_logits"}, - )) - if "lm_encoder_media_chunk" in expanded_components: - specs.append(ComponentModuleSpec( - component="lm_encoder_media_chunk", - module=lm_encoder_media_chunk, - example_inputs=(media_chunk_embeds, torch.zeros_like(chunk_input_ids), chunk_position_ids), - input_keys=("inputs_embeds", "input_ids", "position_ids"), - output_keys=_GEMMA4_DECODER_PIPELINE_IO_KEYS, - graph_meta={ - **common_graph_meta, - "component": "lm_encoder_media_chunk", - "encoder_chunk_size": encoder_chunk_size, - }, - metadata={"family": "gemma4", "task": "multimodal_causal_lm_logits"}, - )) - if "decoder_step" in expanded_components: - if decoder_step_inputs is None: - raise RuntimeError("Gemma4 decoder_step spec requires precomputed step decoder inputs") - specs.append(ComponentModuleSpec( - component="decoder_step", - module=decoder_step, - example_inputs=tuple(decoder_step_inputs), - input_keys=_GEMMA4_DECODER_PIPELINE_IO_KEYS, - output_keys=("logits", "probe_hidden"), - graph_meta={ - **common_graph_meta, - "component": "decoder_step", - "use_internal_kv_cache": True, - "max_cache_seq_len": cache_seq_len, - "cache_sink_size": 4, - "cache_num_slots": (max_slots if dynamic_batch else 1), - }, - metadata={"family": "gemma4", "task": "multimodal_causal_lm_logits"}, - dynamic_batch_axis=0 if dynamic_batch else None, - )) - return specs - - -class Qwen35CausalLMLogitsAdapter(torch.nn.Module): - def __init__(self, model: torch.nn.Module, *, pad_token_id: int | None = None): - super().__init__() - self.model = model - self.backbone = _resolve_qwen35_text_backbone(model) - self.pad_token_id = pad_token_id if pad_token_id is not None else _resolve_model_pad_token_id(model) - from transformers.models.qwen3_5.modeling_qwen3_5 import create_causal_mask # type: ignore - - self._create_causal_mask = create_causal_mask - - def forward(self, input_ids: torch.Tensor): - return self.debug_forward(input_ids)[0] - - def debug_forward(self, input_ids: torch.Tensor) -> tuple[torch.Tensor, list[torch.Tensor]]: - inputs_embeds = self.backbone.embed_tokens(input_ids) - attention_mask = ( - (input_ids != int(self.pad_token_id)).to(dtype=torch.int64) - if self.pad_token_id is not None - else None - ) - base_position_ids = torch.arange(inputs_embeds.shape[1], device=inputs_embeds.device).view(1, 1, -1) - position_ids = torch.cat( - (base_position_ids, base_position_ids, base_position_ids, base_position_ids), - dim=0, - ) - text_position_ids = position_ids[0] - multimodal_position_ids = position_ids[1:] - - causal_mask = self._create_causal_mask( - config=self.backbone.config, - inputs_embeds=inputs_embeds, - attention_mask=attention_mask, - past_key_values=None, - position_ids=text_position_ids, - ) - linear_attn_mask = attention_mask - - hidden_states = inputs_embeds - checkpoints: list[torch.Tensor] = [] - position_embeddings = self.backbone.rotary_emb(hidden_states, multimodal_position_ids) - - for i, decoder_layer in enumerate(self.backbone.layers[: self.backbone.config.num_hidden_layers]): - hidden_states = decoder_layer( - hidden_states, - position_embeddings=position_embeddings, - attention_mask=_qwen35_attention_mask_for_layer( - self.backbone.config, - i, - causal_mask=causal_mask, - linear_attention_mask=linear_attn_mask, - ), - position_ids=text_position_ids, - past_key_values=None, - use_cache=False, - ) - checkpoints.append(hidden_states) - - hidden_states = self.backbone.norm(hidden_states) - hidden_states = _select_last_non_pad_token( - hidden_states, - input_ids, - pad_token_id=self.pad_token_id, - ) - checkpoints.append(hidden_states) - return self.model.lm_head(hidden_states), checkpoints - - def get_transpile_metadata(self): - layer_types = _qwen35_layer_types(self.backbone.config) - return { - "graph": { - **_transpile_graph_meta( - self.model, - adapter_family="qwen3_5", - adapter_type=type(self).__name__, - input_names=("input_ids",), - ), - "num_hidden_layers": int(self.backbone.config.num_hidden_layers), - "layer_types": layer_types, - } - } - - -class Qwen35CausalLMStepAdapter(torch.nn.Module): - def __init__(self, model: torch.nn.Module, *, pad_token_id: int | None = None): - super().__init__() - self.model = model - self.backbone = _resolve_qwen35_text_backbone(model) - self.pad_token_id = pad_token_id if pad_token_id is not None else _resolve_model_pad_token_id(model) - from transformers.models.qwen3_5.modeling_qwen3_5 import create_causal_mask # type: ignore - - self._create_causal_mask = create_causal_mask - - def forward(self, input_ids: torch.Tensor, position_ids: torch.Tensor) -> torch.Tensor: - inputs_embeds = self.backbone.embed_tokens(input_ids) - text_position_ids = position_ids.to(dtype=torch.int64) - base_position_ids = text_position_ids.view(1, text_position_ids.shape[0], -1) - multimodal_position_ids = torch.cat( - (base_position_ids, base_position_ids, base_position_ids), - dim=0, - ) - causal_mask = self._create_causal_mask( - config=self.backbone.config, - inputs_embeds=inputs_embeds, - attention_mask=None, - past_key_values=None, - position_ids=text_position_ids, - ) - linear_attn_mask = None - - hidden_states = inputs_embeds - position_embeddings = self.backbone.rotary_emb(hidden_states, multimodal_position_ids) - for i, decoder_layer in enumerate(self.backbone.layers[: self.backbone.config.num_hidden_layers]): - hidden_states = decoder_layer( - hidden_states, - position_embeddings=position_embeddings, - attention_mask=_qwen35_attention_mask_for_layer( - self.backbone.config, - i, - causal_mask=causal_mask, - linear_attention_mask=linear_attn_mask, - ), - position_ids=text_position_ids, - past_key_values=None, - use_cache=False, - ) - - hidden_states = self.backbone.norm(hidden_states) - return self.model.lm_head(hidden_states[:, -1:, :]) - - def get_transpile_metadata(self): - layer_types = _qwen35_layer_types(self.backbone.config) - return { - "graph": { - **_transpile_graph_meta( - self.model, - adapter_family="qwen3_5", - adapter_type=type(self).__name__, - input_names=("input_ids", "position_ids"), - ), - "num_hidden_layers": int(self.backbone.config.num_hidden_layers), - "layer_types": layer_types, - } - } - - -def _resolve_qwen35_text_backbone(model: torch.nn.Module) -> torch.nn.Module: - backbone = getattr(model, "model", None) - language_model = getattr(backbone, "language_model", None) - if isinstance(language_model, torch.nn.Module): - return language_model - if isinstance(backbone, torch.nn.Module): - return backbone - raise AttributeError(f"{type(model).__name__} does not expose a Qwen3.5 text backbone") - - -def _resolve_qwen35_multimodal_backbone(model: torch.nn.Module) -> torch.nn.Module: - backbone = getattr(model, "model", None) - if isinstance(backbone, torch.nn.Module) and hasattr(backbone, "visual"): - return backbone - raise AttributeError(f"{type(model).__name__} does not expose a Qwen3.5 multimodal backbone") - - -def _qwen35_layer_types(config) -> tuple[str, ...]: - layer_types = tuple(str(value) for value in (getattr(config, "layer_types", ()) or ())) - if layer_types: - return layer_types - num_layers = int(getattr(config, "num_hidden_layers", 0) or 0) - return tuple("full_attention" for _ in range(num_layers)) - - -def _qwen35_attention_mask_for_layer( - config, - layer_index: int, - *, - causal_mask, - linear_attention_mask, -): - layer_types = _qwen35_layer_types(config) - layer_type = layer_types[layer_index] if layer_index < len(layer_types) else "full_attention" - return linear_attention_mask if layer_type == "linear_attention" else causal_mask - - -class Qwen35VisionEncoderAdapter(torch.nn.Module): - def __init__(self, model: torch.nn.Module, *, image_grid_thw: torch.Tensor): - super().__init__() - self.backbone = _resolve_qwen35_multimodal_backbone(model) - self.visual = self.backbone.visual - self.register_buffer("_image_grid_thw", image_grid_thw.detach().clone().to(dtype=torch.long), persistent=False) - with torch.no_grad(): - rotary_pos_emb = self.visual.rot_pos_emb(self._image_grid_thw) - emb = torch.cat((rotary_pos_emb.reshape(rotary_pos_emb.shape[0], -1),) * 2, dim=-1) - self.register_buffer("_vision_cos", emb.cos(), persistent=False) - self.register_buffer("_vision_sin", emb.sin(), persistent=False) - self.register_buffer( - "_vision_pos_embed", - self.visual.fast_pos_embed_interpolate(self._image_grid_thw), - persistent=False, - ) - - def _vision_attention(self, attn: torch.nn.Module, hidden_states: torch.Tensor) -> torch.Tensor: - seq_length = hidden_states.shape[0] - query_states, key_states, value_states = ( - attn.qkv(hidden_states).reshape(seq_length, 3, attn.num_heads, -1).permute(1, 0, 2, 3).unbind(0) - ) - from transformers.models.qwen3_5.modeling_qwen3_5 import apply_rotary_pos_emb_vision # type: ignore - - query_states, key_states = apply_rotary_pos_emb_vision( - query_states, - key_states, - self._vision_cos, - self._vision_sin, - ) - query_states = query_states.transpose(0, 1).unsqueeze(0) - key_states = key_states.transpose(0, 1).unsqueeze(0) - value_states = value_states.transpose(0, 1).unsqueeze(0) - attn_weights = torch.matmul(query_states, key_states.transpose(2, 3)) * float(attn.scaling) - attn_weights = torch.softmax(attn_weights, dim=-1, dtype=torch.float32).to(query_states.dtype) - attn_output = torch.matmul(attn_weights, value_states) - attn_output = attn_output.transpose(1, 2).reshape(seq_length, -1).contiguous() - return attn.proj(attn_output) - - def forward(self, pixel_values: torch.Tensor) -> torch.Tensor: - proj = self.visual.patch_embed.proj - patch_weight = proj.weight.reshape(int(proj.weight.shape[0]), -1) - hidden_states = torch.nn.functional.linear( - pixel_values.to(dtype=patch_weight.dtype), - patch_weight, - proj.bias, - ) - hidden_states = hidden_states + self._vision_pos_embed.to(dtype=hidden_states.dtype) - for block in self.visual.blocks: - hidden_states = hidden_states + self._vision_attention(block.attn, block.norm1(hidden_states)) - hidden_states = hidden_states + block.mlp(block.norm2(hidden_states)) - return self.visual.merger(hidden_states) - - def get_transpile_metadata(self): - return { - "graph": { - **_transpile_graph_meta( - self.backbone, - adapter_family="qwen3_5", - adapter_type=type(self).__name__, - input_names=("pixel_values",), - ), - } - } - - -class Qwen35LMEncoderAdapter(torch.nn.Module): - def __init__(self, model: torch.nn.Module, *, input_ids: torch.Tensor): - super().__init__() - self.backbone = _resolve_qwen35_multimodal_backbone(model) - self.language_model = self.backbone.language_model - self.image_token_id = int(self.backbone.config.image_token_id) - image_positions = (input_ids[0] == self.image_token_id).nonzero(as_tuple=False).flatten() - if int(image_positions.numel()) <= 0: - raise ValueError("Qwen3.5 multimodal LM encoder requires image placeholder tokens") - self.image_token_start = int(image_positions[0].item()) - self.image_token_count = int(image_positions.numel()) - - def forward( - self, - input_ids: torch.Tensor, - attention_mask: torch.Tensor, - position_ids: torch.Tensor, - image_features: torch.Tensor, - ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: - inputs_embeds = self.language_model.embed_tokens(input_ids) - image_embeds = image_features.to(device=inputs_embeds.device, dtype=inputs_embeds.dtype).reshape( - 1, - self.image_token_count, - int(inputs_embeds.shape[-1]), - ) - prefix = inputs_embeds[:, : self.image_token_start, :] - suffix = inputs_embeds[:, self.image_token_start + self.image_token_count :, :] - inputs_embeds = torch.cat((prefix, image_embeds, suffix), dim=1) - return inputs_embeds, attention_mask, position_ids - - def get_transpile_metadata(self): - return { - "graph": { - **_transpile_graph_meta( - self.backbone, - adapter_family="qwen3_5", - adapter_type=type(self).__name__, - input_names=("input_ids", "attention_mask", "position_ids", "image_features"), - ), - } - } - - -class Qwen35EmbedsCausalLMLogitsAdapter(torch.nn.Module): - def __init__(self, model: torch.nn.Module): - super().__init__() - self.model = model - self.backbone = _resolve_qwen35_text_backbone(model) - from transformers.models.qwen3_5.modeling_qwen3_5 import create_causal_mask # type: ignore - - self._create_causal_mask = create_causal_mask - - def forward( - self, - inputs_embeds: torch.Tensor, - attention_mask: torch.Tensor, - position_ids: torch.Tensor, - ) -> torch.Tensor: - if position_ids.ndim == 2: - multimodal_position_ids = position_ids[None, ...].expand(3, position_ids.shape[0], -1) - text_position_ids = position_ids - elif position_ids.ndim == 3 and position_ids.shape[0] == 4: - text_position_ids = position_ids[0] - multimodal_position_ids = position_ids[1:] - else: - text_position_ids = None - multimodal_position_ids = position_ids - causal_mask = self._create_causal_mask( - config=self.backbone.config, - inputs_embeds=inputs_embeds, - attention_mask=attention_mask, - past_key_values=None, - position_ids=text_position_ids, - ) - hidden_states = inputs_embeds - position_embeddings = self.backbone.rotary_emb(hidden_states, multimodal_position_ids) - for i, decoder_layer in enumerate(self.backbone.layers[: self.backbone.config.num_hidden_layers]): - hidden_states = decoder_layer( - hidden_states, - position_embeddings=position_embeddings, - attention_mask=_qwen35_attention_mask_for_layer( - self.backbone.config, - i, - causal_mask=causal_mask, - linear_attention_mask=attention_mask, - ), - position_ids=text_position_ids, - past_key_values=None, - use_cache=False, - ) - hidden_states = self.backbone.norm(hidden_states) - return self.model.lm_head(hidden_states) - - def get_transpile_metadata(self): - layer_types = _qwen35_layer_types(self.backbone.config) - return { - "graph": { - **_transpile_graph_meta( - self.model, - adapter_family="qwen3_5", - adapter_type=type(self).__name__, - input_names=("inputs_embeds", "attention_mask", "position_ids"), - ), - "num_hidden_layers": int(self.backbone.config.num_hidden_layers), - "layer_types": layer_types, - } - } - - -class Qwen35LMEncoderStepAdapter(torch.nn.Module): - def __init__(self, model: torch.nn.Module): - super().__init__() - self.backbone = _resolve_qwen35_text_backbone(model) - - def forward(self, input_ids: torch.Tensor, position_ids: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: - inputs_embeds = self.backbone.embed_tokens(input_ids) - base_position_ids = position_ids.to(dtype=torch.int64).view(1, position_ids.shape[0], -1) - multimodal_position_ids = torch.cat( - (base_position_ids, base_position_ids, base_position_ids), - dim=0, - ) - return inputs_embeds, multimodal_position_ids - - def get_transpile_metadata(self): - return { - "graph": { - **_transpile_graph_meta( - self.backbone, - adapter_family="qwen3_5", - adapter_type=type(self).__name__, - input_names=("input_ids", "position_ids"), - ), - } - } - - -class Qwen35LMEncoderTextChunkAdapter(Qwen35LMEncoderStepAdapter): - pass - - -class Qwen35EmbedsCausalLMStepAdapter(torch.nn.Module): - def __init__(self, model: torch.nn.Module): - super().__init__() - self.model = model - self.backbone = _resolve_qwen35_text_backbone(model) - from transformers.models.qwen3_5.modeling_qwen3_5 import create_causal_mask # type: ignore - - self._create_causal_mask = create_causal_mask - - def forward(self, inputs_embeds: torch.Tensor, position_ids: torch.Tensor) -> torch.Tensor: - if position_ids.ndim == 2: - text_position_ids = position_ids.to(dtype=torch.int64) - base_position_ids = text_position_ids.view(1, text_position_ids.shape[0], -1) - multimodal_position_ids = torch.cat( - (base_position_ids, base_position_ids, base_position_ids), - dim=0, - ) - else: - text_position_ids = None - multimodal_position_ids = position_ids.to(dtype=torch.int64) - causal_mask = self._create_causal_mask( - config=self.backbone.config, - inputs_embeds=inputs_embeds, - attention_mask=None, - past_key_values=None, - position_ids=text_position_ids, - ) - hidden_states = inputs_embeds - position_embeddings = self.backbone.rotary_emb(hidden_states, multimodal_position_ids) - for i, decoder_layer in enumerate(self.backbone.layers[: self.backbone.config.num_hidden_layers]): - hidden_states = decoder_layer( - hidden_states, - position_embeddings=position_embeddings, - attention_mask=_qwen35_attention_mask_for_layer( - self.backbone.config, - i, - causal_mask=causal_mask, - linear_attention_mask=None, - ), - position_ids=text_position_ids, - past_key_values=None, - use_cache=False, - ) - hidden_states = self.backbone.norm(hidden_states) - return self.model.lm_head(hidden_states[:, -1:, :]) - - def get_transpile_metadata(self): - layer_types = _qwen35_layer_types(self.backbone.config) - return { - "graph": { - **_transpile_graph_meta( - self.model, - adapter_family="qwen3_5", - adapter_type=type(self).__name__, - input_names=("inputs_embeds", "position_ids"), - ), - "num_hidden_layers": int(self.backbone.config.num_hidden_layers), - "layer_types": layer_types, - } - } - - -class Qwen35EmbedsCausalLMPrefillChunkAdapter(Qwen35EmbedsCausalLMStepAdapter): - pass - - -class Qwen3CausalLMLogitsAdapter(torch.nn.Module): - def __init__(self, model: torch.nn.Module, *, pad_token_id: int | None = None): - super().__init__() - self.model = model - self.backbone = model.model - self.pad_token_id = pad_token_id if pad_token_id is not None else _resolve_model_pad_token_id(model) - - def forward(self, input_ids: torch.Tensor): - return self.debug_forward(input_ids)[0] - - def debug_forward(self, input_ids: torch.Tensor) -> tuple[torch.Tensor, list[torch.Tensor]]: - inputs_embeds = self.backbone.embed_tokens(input_ids) - position_ids = torch.arange(inputs_embeds.shape[1], device=inputs_embeds.device).unsqueeze(0) - seq_len = int(inputs_embeds.shape[1]) - allowed_positions = torch.tril( - torch.ones((seq_len, seq_len), dtype=torch.bool, device=inputs_embeds.device), - ).view(1, 1, seq_len, seq_len) - if self.pad_token_id is not None: - key_mask = (input_ids != int(self.pad_token_id)).view(input_ids.shape[0], 1, 1, seq_len) - allowed_positions = torch.logical_and(allowed_positions, key_mask) - allowed_values = torch.ones( - (1, 1, seq_len, seq_len), - dtype=inputs_embeds.dtype, - device=inputs_embeds.device, - ) * 0.0 - blocked_values = torch.ones( - (1, 1, seq_len, seq_len), - dtype=inputs_embeds.dtype, - device=inputs_embeds.device, - ) * torch.finfo(inputs_embeds.dtype).min - causal_mask = torch.where( - allowed_positions, - allowed_values, - blocked_values, - ) - - hidden_states = inputs_embeds - checkpoints: list[torch.Tensor] = [] - position_embeddings = self.backbone.rotary_emb(hidden_states, position_ids) - - for i, decoder_layer in enumerate(self.backbone.layers[: self.backbone.config.num_hidden_layers]): - hidden_states = decoder_layer( - hidden_states, - position_embeddings=position_embeddings, - attention_mask=causal_mask, - position_ids=position_ids, - past_key_values=None, - use_cache=False, - ) - checkpoints.append(hidden_states) - - hidden_states = self.backbone.norm(hidden_states) - hidden_states = _select_last_non_pad_token( - hidden_states, - input_ids, - pad_token_id=self.pad_token_id, - ) - checkpoints.append(hidden_states) - return self.model.lm_head(hidden_states), checkpoints - - def get_transpile_metadata(self): - sliding_window = getattr(self.backbone.config, "sliding_window", None) - layer_types = list(getattr(self.backbone.config, "layer_types", [])) - return { - "graph": { - **_transpile_graph_meta( - self.model, - adapter_family="qwen3", - adapter_type=type(self).__name__, - input_names=("input_ids",), - ), - "num_hidden_layers": int(self.backbone.config.num_hidden_layers), - "layer_types": tuple(layer_types), - "sliding_window": None if sliding_window is None else int(sliding_window), - } - } - - -class Qwen3CausalLMStepAdapter(torch.nn.Module): - def __init__(self, model: torch.nn.Module, *, pad_token_id: int | None = None): - super().__init__() - self.model = model - self.backbone = model.model - self.pad_token_id = pad_token_id if pad_token_id is not None else _resolve_model_pad_token_id(model) - - def forward(self, input_ids: torch.Tensor, position_ids: torch.Tensor) -> torch.Tensor: - inputs_embeds = self.backbone.embed_tokens(input_ids) - seq_len = int(inputs_embeds.shape[1]) - allowed_positions = torch.ones( - (1, 1, seq_len, seq_len), - dtype=torch.bool, - device=inputs_embeds.device, - ) - allowed_values = torch.ones( - (1, 1, seq_len, seq_len), - dtype=inputs_embeds.dtype, - device=inputs_embeds.device, - ) * 0.0 - blocked_values = torch.ones( - (1, 1, seq_len, seq_len), - dtype=inputs_embeds.dtype, - device=inputs_embeds.device, - ) * torch.finfo(inputs_embeds.dtype).min - causal_mask = torch.where(allowed_positions, allowed_values, blocked_values) - - hidden_states = inputs_embeds - text_position_ids = position_ids.to(dtype=torch.int64) - position_embeddings = self.backbone.rotary_emb(hidden_states, text_position_ids) - for decoder_layer in self.backbone.layers[: self.backbone.config.num_hidden_layers]: - hidden_states = decoder_layer( - hidden_states, - position_embeddings=position_embeddings, - attention_mask=causal_mask, - position_ids=text_position_ids, - past_key_values=None, - use_cache=False, - ) - - hidden_states = self.backbone.norm(hidden_states) - return self.model.lm_head(hidden_states[:, -1:, :]) - - def get_transpile_metadata(self): - sliding_window = getattr(self.backbone.config, "sliding_window", None) - layer_types = list(getattr(self.backbone.config, "layer_types", [])) - return { - "graph": { - **_transpile_graph_meta( - self.model, - adapter_family="qwen3", - adapter_type=type(self).__name__, - input_names=("input_ids", "position_ids"), - ), - "num_hidden_layers": int(self.backbone.config.num_hidden_layers), - "layer_types": tuple(layer_types), - "sliding_window": None if sliding_window is None else int(sliding_window), - } - } - - -class Qwen3LMEncoderStepAdapter(torch.nn.Module): - def __init__(self, model: torch.nn.Module): - super().__init__() - self.backbone = model.model - - def forward(self, input_ids: torch.Tensor, position_ids: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: - inputs_embeds = self.backbone.embed_tokens(input_ids) - return inputs_embeds, position_ids.to(dtype=torch.int64) - - def get_transpile_metadata(self): - return { - "graph": { - **_transpile_graph_meta( - self.backbone, - adapter_family="qwen3", - adapter_type=type(self).__name__, - input_names=("input_ids", "position_ids"), - ), - } - } - - -class Qwen3LMEncoderTextChunkAdapter(Qwen3LMEncoderStepAdapter): - pass - - -class Qwen3EmbedsCausalLMStepAdapter(torch.nn.Module): - def __init__(self, model: torch.nn.Module): - super().__init__() - self.model = model - self.backbone = model.model - - def _encode_hidden_states(self, inputs_embeds: torch.Tensor, position_ids: torch.Tensor) -> torch.Tensor: - seq_len = int(inputs_embeds.shape[1]) - allowed_positions = torch.tril( - torch.ones((seq_len, seq_len), dtype=torch.bool, device=inputs_embeds.device), - ).view(1, 1, seq_len, seq_len) - allowed_values = torch.zeros( - (1, 1, seq_len, seq_len), - dtype=inputs_embeds.dtype, - device=inputs_embeds.device, - ) - blocked_values = torch.ones( - (1, 1, seq_len, seq_len), - dtype=inputs_embeds.dtype, - device=inputs_embeds.device, - ) * torch.finfo(inputs_embeds.dtype).min - causal_mask = torch.where(allowed_positions, allowed_values, blocked_values) - - hidden_states = inputs_embeds - text_position_ids = position_ids.to(dtype=torch.int64) - position_embeddings = self.backbone.rotary_emb(hidden_states, text_position_ids) - for decoder_layer in self.backbone.layers[: self.backbone.config.num_hidden_layers]: - hidden_states = decoder_layer( - hidden_states, - position_embeddings=position_embeddings, - attention_mask=causal_mask, - position_ids=text_position_ids, - past_key_values=None, - use_cache=False, - ) - return self.backbone.norm(hidden_states) - - def forward(self, inputs_embeds: torch.Tensor, position_ids: torch.Tensor) -> torch.Tensor: - hidden_states = self._encode_hidden_states(inputs_embeds, position_ids) - return self.model.lm_head(hidden_states[:, -1:, :]) - - def get_transpile_metadata(self): - sliding_window = getattr(self.backbone.config, "sliding_window", None) - layer_types = list(getattr(self.backbone.config, "layer_types", [])) - return { - "graph": { - **_transpile_graph_meta( - self.model, - adapter_family="qwen3", - adapter_type=type(self).__name__, - input_names=("inputs_embeds", "position_ids"), - ), - "num_hidden_layers": int(self.backbone.config.num_hidden_layers), - "layer_types": tuple(layer_types), - "sliding_window": None if sliding_window is None else int(sliding_window), - } - } - - -class Qwen3EmbedsCausalLMPrefillChunkAdapter(Qwen3EmbedsCausalLMStepAdapter): - pass - - -class Qwen3EmbedsCausalLMEmbedChunkAdapter(Qwen3EmbedsCausalLMStepAdapter): - """Like the prefill chunk, but emits the full-sequence last hidden state - (no lm_head, no last-token slice) for embedding readout.""" - - def forward(self, inputs_embeds: torch.Tensor, position_ids: torch.Tensor) -> torch.Tensor: - return self._encode_hidden_states(inputs_embeds, position_ids) - - -def _build_qwen_causal_lm_component_specs( - model: torch.nn.Module, - *, - named_tensors: dict[str, torch.Tensor], - weights_dir: str | None = None, - components: tuple[str, ...] | None = None, - cache_context_length: str | int | None = None, -) -> list[ComponentModuleSpec] | None: - input_ids = named_tensors.get("input_ids") - if input_ids is None: - return None - family = _family_key(model) - pad_token_id = _resolve_model_pad_token_id(model) - if family == "qwen3_5": - decoder = Qwen35CausalLMLogitsAdapter(model, pad_token_id=pad_token_id).eval() - decoder_step = Qwen35CausalLMStepAdapter(model, pad_token_id=pad_token_id).eval() - elif family == "qwen3": - decoder = Qwen3CausalLMLogitsAdapter(model, pad_token_id=pad_token_id).eval() - decoder_step = Qwen3CausalLMStepAdapter(model, pad_token_id=pad_token_id).eval() - else: - return None - - requested = tuple(components or ("decoder", "decoder_step")) - requested_set = set(requested) - common_graph_meta = { - "weights_dir": weights_dir, - "task": "causal_lm_logits", - "adapter_family": family, - } - - chunk_components = {"lm_encoder_step", "lm_encoder_text_chunk", "decoder_media_step", "decoder_prefill_chunk", "decoder_embed_chunk"} - wants_chunked = bool(chunk_components & requested_set) - if wants_chunked and family != "qwen3": - raise UnsupportedComponentsError( - f"text chunked-prefill components are only supported for the qwen3 family, got {family}" - ) - - specs: list[ComponentModuleSpec] = [] - if "decoder" in requested_set: - specs.append(ComponentModuleSpec( - component="decoder", - module=decoder, - example_inputs=(input_ids,), - input_keys=("input_ids",), - output_keys=("logits",), - graph_meta={**common_graph_meta, "component": "decoder"}, - metadata={"family": family, "task": "causal_lm_logits"}, - )) - if "decoder_step" in requested_set: - step_input_ids = input_ids[:, :1] - step_position_ids = torch.zeros_like(step_input_ids, dtype=torch.int64) - max_cache_seq_len = _max_cache_seq_len(model, input_ids, cache_context_length, fallback_extra_tokens=512) - specs.append(ComponentModuleSpec( - component="decoder_step", - module=decoder_step, - example_inputs=(step_input_ids, step_position_ids), - input_keys=("input_ids", "position_ids"), - output_keys=("logits",), - graph_meta={ - **common_graph_meta, - "component": "decoder_step", - "use_internal_kv_cache": True, - "use_internal_conv_cache": True, - "use_internal_gated_deltanet_cache": True, - "max_cache_seq_len": max_cache_seq_len, - "cache_sink_size": 4, - }, - metadata={"family": family, "task": "causal_lm_logits"}, - )) - - if wants_chunked: - lm_encoder_step = Qwen3LMEncoderStepAdapter(model).eval() - lm_encoder_text_chunk = Qwen3LMEncoderTextChunkAdapter(model).eval() - decoder_media_step = Qwen3EmbedsCausalLMStepAdapter(model).eval() - decoder_prefill_chunk = Qwen3EmbedsCausalLMPrefillChunkAdapter(model).eval() - decoder_embed_chunk = Qwen3EmbedsCausalLMEmbedChunkAdapter(model).eval() - max_cache_seq_len = _max_cache_seq_len(model, input_ids, cache_context_length, fallback_extra_tokens=512) - prefill_chunk_size = max(1, int(os.environ.get("CACTUS_QWEN_PREFILL_CHUNK", "128") or "128")) - chunk_input_ids = _tile_to_length(input_ids, prefill_chunk_size) - chunk_position_ids = torch.arange( - prefill_chunk_size, - dtype=torch.long, - device=input_ids.device, - ).unsqueeze(0).expand(int(input_ids.shape[0]), -1).contiguous() - with torch.no_grad(): - chunk_embeds, chunk_pos_out = lm_encoder_text_chunk(chunk_input_ids, chunk_position_ids) - step_input_ids = input_ids[:, :1] - step_position_ids = torch.zeros_like(step_input_ids, dtype=torch.int64) - with torch.no_grad(): - step_embeds, step_pos_out = lm_encoder_step(step_input_ids, step_position_ids) - - if "lm_encoder_step" in requested_set: - specs.append(ComponentModuleSpec( - component="lm_encoder_step", - module=lm_encoder_step, - example_inputs=(step_input_ids, step_position_ids), - input_keys=("input_ids", "position_ids"), - output_keys=("inputs_embeds", "position_ids"), - graph_meta={**common_graph_meta, "component": "lm_encoder_step"}, - metadata={"family": family, "task": "causal_lm_logits"}, - )) - if "lm_encoder_text_chunk" in requested_set: - specs.append(ComponentModuleSpec( - component="lm_encoder_text_chunk", - module=lm_encoder_text_chunk, - example_inputs=(chunk_input_ids, chunk_position_ids), - input_keys=("input_ids", "position_ids"), - output_keys=("inputs_embeds", "position_ids"), - graph_meta={**common_graph_meta, "component": "lm_encoder_text_chunk"}, - metadata={"family": family, "task": "causal_lm_logits"}, - )) - if "decoder_prefill_chunk" in requested_set: - specs.append(ComponentModuleSpec( - component="decoder_prefill_chunk", - module=decoder_prefill_chunk, - example_inputs=(chunk_embeds, chunk_pos_out), - input_keys=("inputs_embeds", "position_ids"), - output_keys=("logits",), - graph_meta={ - **common_graph_meta, - "component": "decoder_prefill_chunk", - "use_internal_kv_cache": True, - "use_internal_conv_cache": True, - "use_internal_gated_deltanet_cache": True, - "max_cache_seq_len": max_cache_seq_len, - "cache_sink_size": 4, - "prefill_chunk_size": prefill_chunk_size, - }, - metadata={"family": family, "task": "causal_lm_logits"}, - )) - if "decoder_embed_chunk" in requested_set: - specs.append(ComponentModuleSpec( - component="decoder_embed_chunk", - module=decoder_embed_chunk, - example_inputs=(chunk_embeds, chunk_pos_out), - input_keys=("inputs_embeds", "position_ids"), - output_keys=("last_hidden_state",), - graph_meta={ - **common_graph_meta, - "component": "decoder_embed_chunk", - "use_internal_kv_cache": True, - "use_internal_conv_cache": True, - "use_internal_gated_deltanet_cache": True, - "max_cache_seq_len": max_cache_seq_len, - "cache_sink_size": 4, - "prefill_chunk_size": prefill_chunk_size, - }, - metadata={"family": family, "task": "causal_lm_logits"}, - )) - if "decoder_media_step" in requested_set: - specs.append(ComponentModuleSpec( - component="decoder_media_step", - module=decoder_media_step, - example_inputs=(step_embeds, step_pos_out), - input_keys=("inputs_embeds", "position_ids"), - output_keys=("logits",), - graph_meta={ - **common_graph_meta, - "component": "decoder_media_step", - "use_internal_kv_cache": True, - "use_internal_conv_cache": True, - "use_internal_gated_deltanet_cache": True, - "max_cache_seq_len": max_cache_seq_len, - "cache_sink_size": 4, - }, - metadata={"family": family, "task": "causal_lm_logits"}, - )) - return specs - - -def _build_qwen3_5_multimodal_component_specs( - model: torch.nn.Module, - *, - named_tensors: dict[str, torch.Tensor], - weights_dir: str | None, - components: tuple[str, ...] | None = None, - cache_context_length: str | int | None = None, -) -> list[ComponentModuleSpec] | None: - input_ids = named_tensors.get("input_ids") - attention_mask = named_tensors.get("attention_mask") - position_ids = named_tensors.get("position_ids") - pixel_values = named_tensors.get("pixel_values") - image_grid_thw = named_tensors.get("image_grid_thw") - if not all(isinstance(value, torch.Tensor) for value in (input_ids, attention_mask, position_ids, pixel_values, image_grid_thw)): - return None - - requested = tuple(components or ( - "vision_encoder", - "lm_encoder", - "decoder", - "lm_encoder_text_chunk", - "decoder_prefill_chunk", - "lm_encoder_step", - "decoder_media_step", - "decoder_step", - )) - requested_set = set(requested) - expanded_components: list[str] = [] - - def _require(component: str) -> None: - if component not in expanded_components: - expanded_components.append(component) - - if "decoder" in requested_set: - _require("vision_encoder") - _require("lm_encoder") - _require("decoder") - _require("lm_encoder_text_chunk") - _require("decoder_prefill_chunk") - if "lm_encoder" in requested_set: - _require("vision_encoder") - _require("lm_encoder") - if "lm_encoder_text_chunk" in requested_set: - _require("lm_encoder_text_chunk") - if "decoder_prefill_chunk" in requested_set: - _require("lm_encoder_text_chunk") - _require("decoder_prefill_chunk") - if "vision_encoder" in requested_set: - _require("vision_encoder") - if "decoder_step" in requested_set: - _require("decoder_step") - if "decoder_media_step" in requested_set: - _require("decoder_media_step") - if "lm_encoder_step" in requested_set or "decoder_media_step" in requested_set: - _require("lm_encoder_step") - - vision_encoder = Qwen35VisionEncoderAdapter(model, image_grid_thw=image_grid_thw).eval() - lm_encoder = Qwen35LMEncoderAdapter(model, input_ids=input_ids).eval() - decoder = Qwen35EmbedsCausalLMLogitsAdapter(model).eval() - lm_encoder_step = Qwen35LMEncoderStepAdapter(model).eval() - lm_encoder_text_chunk = Qwen35LMEncoderTextChunkAdapter(model).eval() - decoder_media_step = Qwen35EmbedsCausalLMStepAdapter(model).eval() - decoder_prefill_chunk = Qwen35EmbedsCausalLMPrefillChunkAdapter(model).eval() - decoder_step = Qwen35CausalLMStepAdapter(model, pad_token_id=_resolve_model_pad_token_id(model)).eval() - - image_features: torch.Tensor | None = None - decoder_inputs: tuple[torch.Tensor, ...] | None = None - prefill_decoder_inputs: tuple[torch.Tensor, ...] | None = None - with torch.no_grad(): - if "vision_encoder" in expanded_components and ("lm_encoder" in expanded_components or "decoder" in expanded_components): - image_features = vision_encoder(pixel_values) - if "lm_encoder" in expanded_components or "decoder" in expanded_components: - if image_features is None: - image_features = vision_encoder(pixel_values) - decoder_inputs = lm_encoder(input_ids, attention_mask, position_ids, image_features) - prefill_chunk_size = max(1, int(os.environ.get("CACTUS_QWEN_PREFILL_CHUNK", "128") or "128")) - chunk_input_ids = _tile_to_length(input_ids, prefill_chunk_size) - chunk_position_ids = torch.arange( - prefill_chunk_size, - dtype=torch.long, - device=input_ids.device, - ).unsqueeze(0).expand(int(input_ids.shape[0]), -1).contiguous() - if "decoder_prefill_chunk" in expanded_components: - prefill_decoder_inputs = lm_encoder_text_chunk(chunk_input_ids, chunk_position_ids) - - common_graph_meta = { - "weights_dir": weights_dir, - "task": "multimodal_causal_lm_logits", - "adapter_family": "qwen3_5", - } - max_cache_seq_len = _max_cache_seq_len(model, input_ids, cache_context_length, fallback_extra_tokens=512) - specs: list[ComponentModuleSpec] = [] - if "vision_encoder" in expanded_components: - specs.append(ComponentModuleSpec( - component="vision_encoder", - module=vision_encoder, - example_inputs=(pixel_values,), - input_keys=("pixel_values",), - output_keys=("image_features",), - graph_meta={**common_graph_meta, "component": "vision_encoder"}, - metadata={"family": "qwen3_5", "task": "multimodal_causal_lm_logits"}, - )) - if "lm_encoder" in expanded_components: - if image_features is None: - raise RuntimeError("Qwen3.5 lm_encoder spec requires precomputed image features") - specs.append(ComponentModuleSpec( - component="lm_encoder", - module=lm_encoder, - example_inputs=(input_ids, attention_mask, position_ids, image_features), - input_keys=("input_ids", "attention_mask", "position_ids", "image_features"), - output_keys=("inputs_embeds", "attention_mask", "position_ids"), - graph_meta={**common_graph_meta, "component": "lm_encoder"}, - metadata={"family": "qwen3_5", "task": "multimodal_causal_lm_logits"}, - )) - if "decoder" in expanded_components: - if decoder_inputs is None: - raise RuntimeError("Qwen3.5 decoder spec requires precomputed decoder inputs") - specs.append(ComponentModuleSpec( - component="decoder", - module=decoder, - example_inputs=decoder_inputs, - input_keys=("inputs_embeds", "attention_mask", "position_ids"), - output_keys=("logits",), - graph_meta={**common_graph_meta, "component": "decoder"}, - metadata={"family": "qwen3_5", "task": "multimodal_causal_lm_logits"}, - )) - if "lm_encoder_step" in expanded_components: - step_input_ids = input_ids[:, :1] - step_position_ids = torch.zeros_like(step_input_ids, dtype=torch.int64) - specs.append(ComponentModuleSpec( - component="lm_encoder_step", - module=lm_encoder_step, - example_inputs=(step_input_ids, step_position_ids), - input_keys=("input_ids", "position_ids"), - output_keys=("inputs_embeds", "position_ids"), - graph_meta={**common_graph_meta, "component": "lm_encoder_step"}, - metadata={"family": "qwen3_5", "task": "multimodal_causal_lm_logits"}, - )) - if "lm_encoder_text_chunk" in expanded_components: - specs.append(ComponentModuleSpec( - component="lm_encoder_text_chunk", - module=lm_encoder_text_chunk, - example_inputs=(chunk_input_ids, chunk_position_ids), - input_keys=("input_ids", "position_ids"), - output_keys=("inputs_embeds", "position_ids"), - graph_meta={**common_graph_meta, "component": "lm_encoder_text_chunk"}, - metadata={"family": "qwen3_5", "task": "multimodal_causal_lm_logits"}, - )) - if "decoder_prefill_chunk" in expanded_components: - if prefill_decoder_inputs is None: - raise RuntimeError("Qwen3.5 decoder_prefill_chunk spec requires precomputed text chunk decoder inputs") - specs.append(ComponentModuleSpec( - component="decoder_prefill_chunk", - module=decoder_prefill_chunk, - example_inputs=prefill_decoder_inputs, - input_keys=("inputs_embeds", "position_ids"), - output_keys=("logits",), - graph_meta={ - **common_graph_meta, - "component": "decoder_prefill_chunk", - "use_internal_kv_cache": True, - "use_internal_conv_cache": True, - "use_internal_gated_deltanet_cache": True, - "max_cache_seq_len": max_cache_seq_len, - "cache_sink_size": 4, - "prefill_chunk_size": prefill_chunk_size, - }, - metadata={"family": "qwen3_5", "task": "multimodal_causal_lm_logits"}, - )) - if "decoder_media_step" in expanded_components: - if decoder_inputs is None: - raise RuntimeError("Qwen3.5 decoder_media_step spec requires precomputed decoder inputs") - step_inputs_embeds = decoder_inputs[0][:, :1, :] - step_position_ids = decoder_inputs[2][:, :, :1] - specs.append(ComponentModuleSpec( - component="decoder_media_step", - module=decoder_media_step, - example_inputs=(step_inputs_embeds, step_position_ids), - input_keys=("inputs_embeds", "position_ids"), - output_keys=("logits",), - graph_meta={ - **common_graph_meta, - "component": "decoder_media_step", - "use_internal_kv_cache": True, - "use_internal_conv_cache": True, - "use_internal_gated_deltanet_cache": True, - "max_cache_seq_len": max_cache_seq_len, - "cache_sink_size": 4, - }, - metadata={"family": "qwen3_5", "task": "multimodal_causal_lm_logits"}, - )) - if "decoder_step" in expanded_components: - step_input_ids = input_ids[:, :1] - step_position_ids = torch.zeros_like(step_input_ids, dtype=torch.int64) - specs.append(ComponentModuleSpec( - component="decoder_step", - module=decoder_step, - example_inputs=(step_input_ids, step_position_ids), - input_keys=("input_ids", "position_ids"), - output_keys=("logits",), - graph_meta={ - **common_graph_meta, - "component": "decoder_step", - "use_internal_kv_cache": True, - "use_internal_conv_cache": True, - "use_internal_gated_deltanet_cache": True, - "max_cache_seq_len": max_cache_seq_len, - "cache_sink_size": 4, - }, - metadata={"family": "qwen3_5", "task": "multimodal_causal_lm_logits"}, - )) - return specs - - -class CTCLogitsAdapter(BoundInputAdapter): - def __init__(self, model: torch.nn.Module, *, input_names: tuple[str, ...], family: str): - super().__init__(model, input_names=input_names, family=family, metadata_task="ctc_logits") - - def forward(self, *bound_inputs: torch.Tensor | None) -> torch.Tensor: - outputs = self.model(return_dict=True, **self._kwargs_from_bound_inputs(*bound_inputs)) - return _extract_tensor_output(outputs, preferred_field="logits") - - -class EncoderHiddenStatesAdapter(BoundInputAdapter): - def __init__(self, model: torch.nn.Module, *, input_names: tuple[str, ...], family: str): - encoder = None - get_encoder = getattr(model, "get_encoder", None) - if callable(get_encoder): - encoder = get_encoder() - if encoder is None: - encoder = getattr(model, "encoder", None) - if encoder is None: - model_attr = getattr(model, "model", None) - if model_attr is not None: - encoder = getattr(model_attr, "encoder", None) - if not isinstance(encoder, torch.nn.Module): - raise NotImplementedError(f"{type(model).__name__} does not expose an encoder module") - super().__init__(model, input_names=input_names, family=family, metadata_task="encoder_hidden_states") - self.encoder = encoder - - def forward(self, *bound_inputs: torch.Tensor | None) -> torch.Tensor: - outputs = self.encoder(return_dict=True, **self._kwargs_from_bound_inputs(*bound_inputs)) - return _extract_tensor_output(outputs, preferred_field="last_hidden_state") - - -def _nomic_rotate_half(x: torch.Tensor) -> torch.Tensor: - x1, x2 = x.chunk(2, dim=-1) - return torch.cat([-x2, x1], dim=-1) - - -class NomicTextEmbeddingAdapter(torch.nn.Module): - """Export-friendly nomic-bert encoder that returns last_hidden_state. - - Reuses the HF submodules (embeddings, layernorms, attention/MLP weights) but - reimplements the forward to avoid the HF rotary cache and the data-dependent - MoE dispatch (`torch.where`/`index_add_`) which do not survive torch.export. - Pooling and L2 normalization are applied downstream in the C++ engine. - """ - - def __init__(self, model: torch.nn.Module) -> None: - super().__init__() - for name in ("embeddings", "emb_ln", "encoder"): - if not isinstance(getattr(model, name, None), torch.nn.Module): - raise NotImplementedError(f"{type(model).__name__} is not a nomic-bert model") - self.embeddings = model.embeddings - self.emb_ln = model.emb_ln - self.encoder = model.encoder - config = model.config - self.n_heads = int(config.n_head) - self.head_dim = int(config.n_embd) // int(config.n_head) - self.hidden_dim = int(config.n_embd) - self.n_experts = int(getattr(config, "num_experts", 0) or 0) - self.top_k = int(getattr(config, "moe_top_k", 0) or 0) - self.ffn_dim = int(config.n_inner) - self.rope_base = float(getattr(config, "rotary_emb_base", 10000.0) or 10000.0) - # Pre-transpose each MoE expert w2 to [hidden, n_experts*ffn] so the second - # expert matmul consumes it as a direct linear weight — transposing a packed - # quantized weight inside the graph is not lowerable. Idempotent by shape so - # repeated adapter construction on the same model is safe. - packed_rows = self.n_experts * self.ffn_dim - for layer in self.encoder.layers: - if getattr(layer, "moe", False): - w2 = layer.mlp.experts.mlp.w2 - if w2.shape[0] == packed_rows: - layer.mlp.experts.mlp.w2 = torch.nn.Parameter( - w2.detach().t().contiguous(), requires_grad=False - ) - - def _rope_cos_sin(self, seq_len: int, device, dtype): - inv_freq = 1.0 / (self.rope_base ** (torch.arange(0, self.head_dim, 2, device=device, dtype=torch.float32) / self.head_dim)) - positions = torch.arange(seq_len, device=device, dtype=torch.float32) - freqs = positions.unsqueeze(1) * inv_freq.unsqueeze(0) - cos = torch.cat([freqs.cos(), freqs.cos()], dim=-1).to(dtype) - sin = torch.cat([freqs.sin(), freqs.sin()], dim=-1).to(dtype) - return cos, sin - - def _attention(self, layer, hidden_states, additive_mask, cos, sin): - batch, seq, _ = hidden_states.shape - qkv = layer.attn.Wqkv(hidden_states).view(batch, seq, 3, self.n_heads, self.head_dim) - query, key, value = qkv[:, :, 0], qkv[:, :, 1], qkv[:, :, 2] - cos_b = cos[None, :, None, :] - sin_b = sin[None, :, None, :] - query = query * cos_b + _nomic_rotate_half(query) * sin_b - key = key * cos_b + _nomic_rotate_half(key) * sin_b - query = query.permute(0, 2, 1, 3) - key = key.permute(0, 2, 1, 3) - value = value.permute(0, 2, 1, 3) - scores = torch.matmul(query, key.transpose(-1, -2)) / (self.head_dim ** 0.5) - scores = scores + additive_mask - context = torch.matmul(F.softmax(scores, dim=-1), value) - context = context.permute(0, 2, 1, 3).reshape(batch, seq, self.hidden_dim) - return layer.attn.out_proj(context) - - def _topk_gate(self, weights): - # torch.topk does not lower; build the top-k softmax gate via iterative - # max + masking (top_k is small). gate[n, e] = softmax weight when expert e - # is among the top-k for token n, else 0. - neg = torch.finfo(weights.dtype).min - remaining = weights - gate = weights * 0.0 - for _ in range(self.top_k): - current_max = torch.amax(remaining, dim=-1, keepdim=True) - selected = (remaining >= current_max).to(weights.dtype) - gate = gate + selected * weights - remaining = remaining + selected * neg - return gate - - def _moe_mlp(self, moe, hidden_states): - # Dense MoE: run all experts via two fused matmuls against the packed - # expert weights (no per-expert indexing — slicing a quantized weight is - # not lowerable), then gate by scaling each expert's activations. The - # gated sum is identical to top-k routing because non-selected experts - # have gate 0. - batch, seq, _ = hidden_states.shape - flat = hidden_states.reshape(-1, self.hidden_dim) - weights = moe.router.layer(flat).softmax(dim=-1, dtype=torch.float32) - gate = self._topk_gate(weights).to(hidden_states.dtype) - all_hidden = F.linear(flat, moe.experts.mlp.w1) # [tokens, n_experts * ffn] - activated = F.gelu(all_hidden).reshape(-1, self.n_experts, self.ffn_dim) - activated = activated * gate.unsqueeze(-1) - activated = activated.reshape(-1, self.n_experts * self.ffn_dim) - # w2 is pre-transposed to [hidden, n_experts*ffn] in __init__. - out = F.linear(activated, moe.experts.mlp.w2) - out = out + moe.experts.bias - return out.reshape(batch, seq, self.hidden_dim) - - def forward(self, input_ids: torch.Tensor, attention_mask: torch.Tensor) -> torch.Tensor: - hidden_states = self.emb_ln(self.embeddings.word_embeddings(input_ids)) - seq_len = input_ids.shape[1] - cos, sin = self._rope_cos_sin(seq_len, input_ids.device, hidden_states.dtype) - mask_float = attention_mask[:, None, None, :].to(hidden_states.dtype) - additive_mask = (mask_float - 1.0) * (-torch.finfo(hidden_states.dtype).min) - for layer in self.encoder.layers: - hidden_states = layer.norm1(self._attention(layer, hidden_states, additive_mask, cos, sin) + hidden_states) - mlp_out = self._moe_mlp(layer.mlp, hidden_states) if layer.moe else layer.mlp(hidden_states) - hidden_states = layer.norm2(mlp_out + hidden_states) - return hidden_states - - -class WhisperEncoderComponentAdapter(torch.nn.Module): - def __init__(self, encoder: torch.nn.Module): - super().__init__() - self.encoder = encoder - - def forward(self, input_features: torch.Tensor) -> torch.Tensor: - hidden_states = F.gelu(self.encoder.conv1(input_features)) - hidden_states = F.gelu(self.encoder.conv2(hidden_states)) - hidden_states = hidden_states.permute(0, 2, 1) - - position_ids = torch.arange(hidden_states.shape[1], device=hidden_states.device) - positions = self.encoder.embed_positions(position_ids) - hidden_states = hidden_states + positions.to( - device=hidden_states.device, - dtype=hidden_states.dtype, - ).unsqueeze(0) - - for layer in self.encoder.layers: - residual = hidden_states - hidden_states = layer.self_attn_layer_norm(hidden_states) - hidden_states, _ = layer.self_attn( - hidden_states=hidden_states, - attention_mask=None, - ) - hidden_states = residual + hidden_states - - residual = hidden_states - hidden_states = layer.final_layer_norm(hidden_states) - hidden_states = layer.activation_fn(layer.fc1(hidden_states)) - hidden_states = layer.fc2(hidden_states) - hidden_states = residual + hidden_states - - return self.encoder.layer_norm(hidden_states) - - -class WhisperDecoderCrossKVComponentAdapter(torch.nn.Module): - def __init__(self, decoder: torch.nn.Module): - super().__init__() - self.decoder = decoder - - def forward(self, encoder_hidden_states: torch.Tensor) -> tuple[torch.Tensor, ...]: - outputs: list[torch.Tensor] = [] - batch = int(encoder_hidden_states.shape[0]) - for layer in self.decoder.layers: - attn = layer.encoder_attn - kv_shape = (batch, -1, int(attn.num_heads), int(attn.head_dim)) - outputs.append(attn.k_proj(encoder_hidden_states).view(kv_shape).contiguous()) - outputs.append(attn.v_proj(encoder_hidden_states).view(kv_shape).contiguous()) - return tuple(outputs) - - -class WhisperDecoderStepWithCrossKVComponentAdapter(torch.nn.Module): - def __init__(self, decoder: torch.nn.Module, proj_out: torch.nn.Module, *, pad_token_id: int | None): - super().__init__() - self.decoder = decoder - self.proj_out = proj_out - self.pad_token_id = pad_token_id - - def forward( - self, - decoder_input_ids: torch.Tensor, - position_ids: torch.Tensor, - *cross_kv: torch.Tensor, - ) -> torch.Tensor: - hidden_states = self.decoder.embed_tokens(decoder_input_ids) - positions = self.decoder.embed_positions( - decoder_input_ids, - past_key_values_length=0, - position_ids=position_ids, - ) - hidden_states = hidden_states + positions.to( - device=hidden_states.device, - dtype=hidden_states.dtype, - ) - - for layer_index, layer in enumerate(self.decoder.layers): - residual = hidden_states - hidden_states = layer.self_attn_layer_norm(hidden_states) - hidden_states, _ = layer.self_attn( - hidden_states, - attention_mask=None, - past_key_values=None, - ) - hidden_states = residual + hidden_states - - residual = hidden_states - hidden_states = layer.encoder_attn_layer_norm(hidden_states) - attn = layer.encoder_attn - query_shape = (*hidden_states.shape[:-1], int(attn.num_heads), int(attn.head_dim)) - query_states = (attn.q_proj(hidden_states) * attn.scaling).view(query_shape).transpose(1, 2).contiguous() - key_states = cross_kv[layer_index * 2].transpose(1, 2).contiguous() - value_states = cross_kv[layer_index * 2 + 1].transpose(1, 2).contiguous() - attn_output = torch.nn.functional.scaled_dot_product_attention( - query_states, - key_states, - value_states, - attn_mask=None, - dropout_p=0.0, - scale=1.0, - ) - attn_output = attn_output.transpose(1, 2).reshape(*hidden_states.shape[:-1], -1).contiguous() - hidden_states = residual + attn.out_proj(attn_output) - - residual = hidden_states - hidden_states = layer.final_layer_norm(hidden_states) - hidden_states = layer.activation_fn(layer.fc1(hidden_states)) - hidden_states = layer.fc2(hidden_states) - hidden_states = residual + hidden_states - - hidden_states = self.decoder.layer_norm(hidden_states) - hidden_states = _select_last_non_pad_token( - hidden_states, - decoder_input_ids, - pad_token_id=self.pad_token_id, - ) - return self.proj_out(hidden_states) - - -class NeedleSourceEncoderComponentAdapter(torch.nn.Module): - def __init__(self, model: torch.nn.Module): - super().__init__() - self.model = model - - def forward(self, input_ids: torch.Tensor, attention_mask: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: - encode = getattr(self.model, "cactus_source_encode", None) - if not callable(encode): - raise TypeError(f"{type(self.model).__name__} does not expose cactus_source_encode") - return encode(input_ids, attention_mask) - - def get_transpile_metadata(self): - return { - "graph": { - **_transpile_graph_meta( - self.model, - adapter_family="needle", - adapter_type=type(self).__name__, - input_names=("input_ids", "attention_mask"), - ), - "task": "causal_lm_logits", - } - } - - -class NeedleDecoderCrossKVComponentAdapter(torch.nn.Module): - def __init__(self, model: torch.nn.Module): - super().__init__() - self.model = model - - def forward( - self, - encoder_hidden_states: torch.Tensor, - encoder_attention_mask: torch.Tensor, - ) -> tuple[torch.Tensor, ...]: - cross_kv = getattr(self.model, "cactus_decoder_cross_kv", None) - if not callable(cross_kv): - raise TypeError(f"{type(self.model).__name__} does not expose cactus_decoder_cross_kv") - return cross_kv(encoder_hidden_states, encoder_attention_mask) - - def get_transpile_metadata(self): - return { - "graph": { - **_transpile_graph_meta( - self.model, - adapter_family="needle", - adapter_type=type(self).__name__, - input_names=("encoder_hidden_states", "encoder_attention_mask"), - ), - "task": "causal_lm_logits", - } - } - - -class NeedleDecoderStepComponentAdapter(torch.nn.Module): - def __init__(self, model: torch.nn.Module): - super().__init__() - self.model = model - - def forward( - self, - decoder_input_ids: torch.Tensor, - position_ids: torch.Tensor, - encoder_attention_mask: torch.Tensor, - *cross_kv: torch.Tensor, - ) -> torch.Tensor: - step = getattr(self.model, "cactus_decoder_step", None) - if not callable(step): - raise TypeError(f"{type(self.model).__name__} does not expose cactus_decoder_step") - return step(decoder_input_ids, position_ids, encoder_attention_mask, *cross_kv) - - def get_transpile_metadata(self): - return { - "graph": { - **_transpile_graph_meta( - self.model, - adapter_family="needle", - adapter_type=type(self).__name__, - input_names=("decoder_input_ids", "position_ids", "encoder_attention_mask"), - ), - "task": "causal_lm_logits", - } - } - - -def _family_key(model: torch.nn.Module) -> str: - explicit_family = getattr(model, "family", None) - if isinstance(explicit_family, str) and explicit_family: - return explicit_family.lower() - module_name = type(model).__module__ - if module_name.startswith("transformers.models.whisper."): - return "whisper" - if module_name.startswith("transformers.models.gemma4."): - return "gemma4" - if module_name.startswith("transformers.models.gemma3."): - return "gemma3" - if module_name.startswith("transformers.models.gemma."): - return "gemma" - if module_name.startswith("transformers.models.qwen3_5."): - return "qwen3_5" - if module_name.startswith("transformers.models.qwen3_vl."): - return "qwen3_5" - if module_name.startswith("transformers.models.qwen3."): - return "qwen3" - if module_name.startswith("transformers.models.lfm2_vl."): - return "lfm2_vl" - if module_name.startswith("transformers.models.lfm2_moe."): - return "lfm2_moe" - if module_name.startswith("transformers.models.lfm2."): - return "lfm2" - model_type = str(getattr(getattr(model, "config", None), "model_type", "") or "").lower() - if model_type == "needle": - return "needle" - if "nomic" in module_name.lower() or "nomic" in model_type: - return "nomic" - return "generic" - - -_ENCODER_CROSS_KV_ROUTE = "encoder_cross_kv_decoder_step" - - -def _cross_kv_output_keys(num_layers: int) -> tuple[str, ...]: - return tuple( - name - for layer_index in range(int(num_layers)) - for name in (f"cross_k_{layer_index}", f"cross_v_{layer_index}") - ) - - -def _component_spec( - *, - component: str, - module: torch.nn.Module, - example_inputs: tuple[torch.Tensor, ...], - input_keys: tuple[str, ...], - output_keys: tuple[str, ...], - common_graph_meta: dict[str, object], - family: str, - task: str, - graph_meta: dict[str, object] | None = None, - runtime_role: str | None = None, - source_kind: str | None = None, -) -> ComponentModuleSpec: - component_metadata: dict[str, object] = {"family": family, "task": task} - if runtime_role is not None: - component_metadata["runtime_route"] = _ENCODER_CROSS_KV_ROUTE - component_metadata["runtime_role"] = runtime_role - if source_kind is not None: - component_metadata["source_kind"] = source_kind - return ComponentModuleSpec( - component=component, - module=module, - example_inputs=example_inputs, - input_keys=input_keys, - output_keys=output_keys, - graph_meta={**common_graph_meta, "component": component, **(graph_meta or {})}, - metadata=component_metadata, - ) - - -def _encoder_cross_kv_step_specs( - *, - source_component: str, - source_module: torch.nn.Module, - source_inputs: tuple[torch.Tensor, ...], - source_input_keys: tuple[str, ...], - source_output_keys: tuple[str, ...], - source_kind: str, - cross_kv_module: torch.nn.Module, - cross_kv_inputs: tuple[torch.Tensor, ...], - cross_kv_input_keys: tuple[str, ...], - cross_kv_output_keys: tuple[str, ...], - decoder_step_module: torch.nn.Module, - decoder_step_inputs: tuple[torch.Tensor, ...], - decoder_step_input_keys: tuple[str, ...], - common_graph_meta: dict[str, object], - family: str, - task: str, - max_cache_seq_len: int, - decoder_step_graph_meta: dict[str, object] | None = None, -) -> list[ComponentModuleSpec]: - return [ - _component_spec( - component=source_component, - module=source_module, - example_inputs=source_inputs, - input_keys=source_input_keys, - output_keys=source_output_keys, - common_graph_meta=common_graph_meta, - family=family, - task=task, - runtime_role="source_encoder", - source_kind=source_kind, - ), - _component_spec( - component="decoder_cross_kv", - module=cross_kv_module, - example_inputs=cross_kv_inputs, - input_keys=cross_kv_input_keys, - output_keys=cross_kv_output_keys, - common_graph_meta=common_graph_meta, - family=family, - task=task, - runtime_role="decoder_cross_kv", - ), - _component_spec( - component="decoder_step", - module=decoder_step_module, - example_inputs=decoder_step_inputs, - input_keys=decoder_step_input_keys, - output_keys=("logits",), - common_graph_meta=common_graph_meta, - family=family, - task=task, - graph_meta={ - "use_internal_kv_cache": True, - "max_cache_seq_len": max_cache_seq_len, - "cache_sink_size": 0, - **(decoder_step_graph_meta or {}), - }, - runtime_role="decoder_step", - ), - ] - - -def _build_whisper_seq2seq_component_specs( - model: torch.nn.Module, - *, - named_tensors: dict[str, torch.Tensor], - inputs_metadata: dict[str, object] | None = None, - weights_dir: str | None = None, -) -> list[ComponentModuleSpec]: - input_features = named_tensors.get("input_features") - if not isinstance(input_features, torch.Tensor): - raise RuntimeError("Whisper component transpile requires input_features") - - get_encoder = getattr(model, "get_encoder", None) - encoder = get_encoder() if callable(get_encoder) else None - if not isinstance(encoder, torch.nn.Module): - encoder = getattr(getattr(model, "model", None), "encoder", None) - decoder = getattr(getattr(model, "model", None), "decoder", None) - proj_out = getattr(model, "proj_out", None) - if not isinstance(encoder, torch.nn.Module) or not isinstance(decoder, torch.nn.Module): - raise RuntimeError(f"{type(model).__name__} does not expose Whisper encoder/decoder modules") - if not isinstance(proj_out, torch.nn.Module): - raise RuntimeError(f"{type(model).__name__} does not expose a Whisper projection head") - - metadata = dict(inputs_metadata or {}) - decoder_prompt_ids = metadata.get("decoder_input_ids") - if not isinstance(decoder_prompt_ids, list) or not decoder_prompt_ids: - decoder_start_token_id = getattr(getattr(model, "config", None), "decoder_start_token_id", None) - if not isinstance(decoder_start_token_id, int): - raise RuntimeError("Whisper component transpile requires decoder_input_ids metadata") - decoder_prompt_ids = [int(decoder_start_token_id)] - else: - decoder_prompt_ids = [int(value) for value in decoder_prompt_ids] - - target_token_count = int(metadata.get("target_token_count", len(decoder_prompt_ids)) or len(decoder_prompt_ids)) - target_token_count = max(target_token_count, len(decoder_prompt_ids)) - pad_token_id = int(metadata.get("pad_token_id", getattr(getattr(model, "config", None), "pad_token_id", 0)) or 0) - encoder_adapter = WhisperEncoderComponentAdapter(encoder).eval() - decoder_cross_kv_adapter = WhisperDecoderCrossKVComponentAdapter(decoder).eval() - decoder_step_adapter = WhisperDecoderStepWithCrossKVComponentAdapter( - decoder, - proj_out, - pad_token_id=pad_token_id, - ).eval() - - with torch.no_grad(): - encoder_hidden_states = encoder_adapter(input_features) - decoder_cross_kv = decoder_cross_kv_adapter(encoder_hidden_states) - - decoder_input_ids = torch.tensor([[decoder_prompt_ids[0]]], dtype=torch.int64, device=input_features.device) - step_position_ids = torch.zeros_like(decoder_input_ids, dtype=torch.int64) - - common_graph_meta = { - **_transpile_graph_meta( - model, - adapter_family="whisper", - adapter_type="component_pipeline", - input_names=("input_features",), - ), - "task": "seq2seq_transcription", - "adapter_family": "whisper", - } - if weights_dir: - common_graph_meta["weights_dir"] = weights_dir - cross_kv_output_keys = _cross_kv_output_keys(len(decoder.layers)) - return _encoder_cross_kv_step_specs( - source_component="audio_encoder", - source_module=encoder_adapter, - source_inputs=(input_features,), - source_input_keys=("input_features",), - source_output_keys=("encoder_hidden_states",), - source_kind="audio_features", - cross_kv_module=decoder_cross_kv_adapter, - cross_kv_inputs=(encoder_hidden_states,), - cross_kv_input_keys=("encoder_hidden_states",), - cross_kv_output_keys=cross_kv_output_keys, - decoder_step_module=decoder_step_adapter, - decoder_step_inputs=(decoder_input_ids, step_position_ids, *decoder_cross_kv), - decoder_step_input_keys=("decoder_input_ids", "position_ids", *cross_kv_output_keys), - common_graph_meta=common_graph_meta, - family="whisper", - task="seq2seq_transcription", - max_cache_seq_len=max(16, int(target_token_count)), - ) - - -def _resolve_decoder_start_token_id(model: torch.nn.Module, *, fallback: int = 0) -> int: - config = getattr(model, "config", None) - generation_config = getattr(model, "generation_config", None) - for owner in (config, generation_config): - for attr_name in ("decoder_start_token_id", "bos_token_id", "eos_token_id", "pad_token_id"): - value = getattr(owner, attr_name, None) if owner is not None else None - if value is not None: - return int(value) - return int(fallback) - -NEEDLE_CONTEXT_LENGTH = 1024 - -def _build_needle_causal_lm_component_specs( - model: torch.nn.Module, - *, - named_tensors: dict[str, torch.Tensor], - weights_dir: str | None = None, -) -> list[ComponentModuleSpec] | None: - input_ids = named_tensors.get("input_ids") - if input_ids is None: - return None - for method_name in ("cactus_source_encode", "cactus_decoder_cross_kv", "cactus_decoder_step"): - if not callable(getattr(model, method_name, None)): - raise RuntimeError(f"{type(model).__name__} does not expose {method_name}") - - pad_token_id = _resolve_model_pad_token_id(model) - current_len = int(input_ids.shape[1]) - if current_len < NEEDLE_CONTEXT_LENGTH: - pad_value = int(pad_token_id) if pad_token_id is not None else 0 - padding = torch.full( - (int(input_ids.shape[0]), NEEDLE_CONTEXT_LENGTH - current_len), - pad_value, - dtype=input_ids.dtype, - device=input_ids.device, - ) - input_ids = torch.cat([input_ids, padding], dim=1) - if pad_token_id is None: - attention_mask = torch.ones_like(input_ids, dtype=torch.int64) - else: - attention_mask = (input_ids != int(pad_token_id)).to(dtype=torch.int64) - source_encoder = NeedleSourceEncoderComponentAdapter(model).eval() - decoder_cross_kv = NeedleDecoderCrossKVComponentAdapter(model).eval() - decoder_step = NeedleDecoderStepComponentAdapter(model).eval() - - with torch.no_grad(): - encoder_hidden_states, encoder_attention_mask = source_encoder(input_ids, attention_mask) - cross_kv = decoder_cross_kv(encoder_hidden_states, encoder_attention_mask) - - decoder_start_token_id = _resolve_decoder_start_token_id(model, fallback=int(input_ids[0, 0].item())) - decoder_input_ids = torch.full( - (int(input_ids.shape[0]), 1), - decoder_start_token_id, - dtype=torch.int64, - device=input_ids.device, - ) - step_position_ids = torch.zeros_like(decoder_input_ids, dtype=torch.int64) - cross_kv_output_keys = _cross_kv_output_keys(len(cross_kv) // 2) - - common_graph_meta = { - **_transpile_graph_meta( - model, - adapter_family="needle", - adapter_type="component_pipeline", - input_names=("input_ids",), - ), - "task": "causal_lm_logits", - "adapter_family": "needle", - } - if weights_dir: - common_graph_meta["weights_dir"] = weights_dir - - return _encoder_cross_kv_step_specs( - source_component="source_encoder", - source_module=source_encoder, - source_inputs=(input_ids, attention_mask), - source_input_keys=("input_ids", "attention_mask"), - source_output_keys=("encoder_hidden_states", "encoder_attention_mask"), - source_kind="text_tokens", - cross_kv_module=decoder_cross_kv, - cross_kv_inputs=(encoder_hidden_states, encoder_attention_mask), - cross_kv_input_keys=("encoder_hidden_states", "encoder_attention_mask"), - cross_kv_output_keys=cross_kv_output_keys, - decoder_step_module=decoder_step, - decoder_step_inputs=(decoder_input_ids, step_position_ids, encoder_attention_mask, *cross_kv), - decoder_step_input_keys=("decoder_input_ids", "position_ids", "encoder_attention_mask", *cross_kv_output_keys), - common_graph_meta=common_graph_meta, - family="needle", - task="causal_lm_logits", - max_cache_seq_len=max(16, int(input_ids.shape[1])), - decoder_step_graph_meta={ - "external_cross_kv_cache_inputs": True, - "cross_kv_input_start_index": 3, - "cross_kv_input_count": len(cross_kv_output_keys), - "cross_kv_input_layout": "bthd", - }, - ) - - -def _build_nomic_text_embedding_component_specs( - model: torch.nn.Module, - *, - named_tensors: dict[str, torch.Tensor], - weights_dir: str | None = None, -) -> list[ComponentModuleSpec]: - input_ids = named_tensors.get("input_ids") - if not isinstance(input_ids, torch.Tensor): - raise RuntimeError("nomic text_embedding transpile requires input_ids") - attention_mask = named_tensors.get("attention_mask") - if not isinstance(attention_mask, torch.Tensor): - attention_mask = torch.ones_like(input_ids) - - adapter = NomicTextEmbeddingAdapter(model).eval() - - common_graph_meta = { - **_transpile_graph_meta( - model, - adapter_family="nomic", - adapter_type="component_pipeline", - input_names=("input_ids", "attention_mask"), - ), - "task": "text_embedding", - "adapter_family": "nomic", - } - if weights_dir: - common_graph_meta["weights_dir"] = weights_dir - - return [ - ComponentModuleSpec( - component="text_embedding", - module=adapter, - example_inputs=(input_ids, attention_mask), - input_keys=("input_ids", "attention_mask"), - output_keys=("last_hidden_state",), - graph_meta={**common_graph_meta, "component": "text_embedding"}, - metadata={"family": "nomic", "task": "text_embedding"}, - ) - ] - - -def _build_lfm2_causal_lm_component_specs( - model: torch.nn.Module, - *, - named_tensors: dict[str, torch.Tensor], - weights_dir: str | None = None, - components: tuple[str, ...] | None = None, - cache_context_length: str | int | None = None, -) -> list[ComponentModuleSpec] | None: - input_ids = named_tensors.get("input_ids") - if input_ids is None: - return None - requested = tuple(components or ("decoder", "decoder_step")) - requested_set = set(requested) - common_graph_meta = { - "weights_dir": weights_dir, - "task": "causal_lm_logits", - "adapter_family": "lfm2", - } - metadata = {"family": "lfm2", "task": "causal_lm_logits"} - chunk_components = {"lm_encoder_step", "lm_encoder_text_chunk", "decoder_prefill_chunk"} - if chunk_components & requested_set: - return _lfm2_chunked_pipeline_specs( - model, - requested_set=requested_set, - input_ids=input_ids, - weights_dir=weights_dir, - cache_context_length=cache_context_length, - common_graph_meta=common_graph_meta, - metadata=metadata, - ) - - pad_token_id = _resolve_model_pad_token_id(model) - decoder = Lfm2CausalLMLogitsAdapter(model, pad_token_id=pad_token_id).eval() - decoder_step = Lfm2CausalLMStepAdapter(model, pad_token_id=pad_token_id).eval() - specs: list[ComponentModuleSpec] = [] - if "decoder" in requested_set: - specs.append(ComponentModuleSpec( - component="decoder", - module=decoder, - example_inputs=(input_ids,), - input_keys=("input_ids",), - output_keys=("logits",), - graph_meta={**common_graph_meta, "component": "decoder"}, - metadata=metadata, - )) - if "decoder_step" in requested_set: - step_input_ids = input_ids[:, :1] - step_position_ids = torch.zeros_like(step_input_ids, dtype=torch.int64) - max_cache_seq_len = _max_cache_seq_len(model, input_ids, cache_context_length, fallback_extra_tokens=512) - specs.append(ComponentModuleSpec( - component="decoder_step", - module=decoder_step, - example_inputs=(step_input_ids, step_position_ids), - input_keys=("input_ids", "position_ids"), - output_keys=("logits",), - graph_meta={ - **common_graph_meta, - "component": "decoder_step", - "use_internal_kv_cache": True, - "use_internal_conv_cache": True, - "use_internal_gated_deltanet_cache": True, - "max_cache_seq_len": max_cache_seq_len, - "cache_sink_size": 4, - }, - metadata=metadata, - )) - return specs - - -def _lfm2_chunked_pipeline_specs( - model: torch.nn.Module, - *, - requested_set: set[str], - input_ids: torch.Tensor, - weights_dir: str | None, - cache_context_length: str | int | None, - common_graph_meta: dict[str, object], - metadata: dict[str, str], - decoder_inputs_fn: Callable[[], tuple[torch.Tensor, ...]] | None = None, -) -> list[ComponentModuleSpec]: - lm_encoder_step = Lfm2VlLMEncoderStepAdapter(model, weights_dir=weights_dir).eval() - lm_encoder_text_chunk = Lfm2VlLMEncoderTextChunkAdapter(model, weights_dir=weights_dir).eval() - decoder_last_token = Lfm2VlDecoderAdapter(model, weights_dir=weights_dir, last_token_only=True).eval() - decoder_embed = Lfm2VlDecoderAdapter(model, weights_dir=weights_dir, last_token_only=False, return_hidden=True).eval() - - prefill_chunk = max(2, int(os.environ.get("CACTUS_LFM2_PREFILL_CHUNK", "128") or "128")) - chunk_input_ids = _tile_to_length(input_ids, prefill_chunk) - chunk_position_ids = torch.arange( - prefill_chunk, dtype=torch.long, device=input_ids.device, - ).unsqueeze(0).expand(int(input_ids.shape[0]), -1).contiguous() - step_input_ids = input_ids[:, :1] - step_position_ids = torch.zeros_like(step_input_ids, dtype=torch.int64) - - decoder_inputs: tuple[torch.Tensor, ...] | None = None - if {"decoder_prefill_chunk", "decoder_embed_chunk", "decoder_step"} & requested_set: - if decoder_inputs_fn is not None: - decoder_inputs = decoder_inputs_fn() - else: - with torch.no_grad(): - decoder_inputs = lm_encoder_text_chunk(chunk_input_ids, chunk_position_ids) - - cache_graph_meta = { - "use_internal_kv_cache": True, - "use_internal_conv_cache": True, - "max_cache_seq_len": _max_cache_seq_len(model, input_ids, cache_context_length, fallback_extra_tokens=512), - "cache_sink_size": 4, - } - - specs: list[ComponentModuleSpec] = [] - if "lm_encoder_step" in requested_set: - specs.append(ComponentModuleSpec( - component="lm_encoder_step", - module=lm_encoder_step, - example_inputs=(step_input_ids, step_position_ids), - input_keys=("input_ids", "position_ids"), - output_keys=("inputs_embeds", "attention_mask", "position_ids"), - graph_meta={**common_graph_meta, "component": "lm_encoder_step"}, - metadata=metadata, - )) - if "lm_encoder_text_chunk" in requested_set: - specs.append(ComponentModuleSpec( - component="lm_encoder_text_chunk", - module=lm_encoder_text_chunk, - example_inputs=(chunk_input_ids, chunk_position_ids), - input_keys=("input_ids", "position_ids"), - output_keys=("inputs_embeds", "attention_mask", "position_ids"), - graph_meta={ - **common_graph_meta, - "component": "lm_encoder_text_chunk", - "encoder_chunk_size": prefill_chunk, - }, - metadata=metadata, - )) - if "decoder_prefill_chunk" in requested_set: - specs.append(ComponentModuleSpec( - component="decoder_prefill_chunk", - module=decoder_last_token, - example_inputs=tuple(tensor[:, :prefill_chunk, ...] for tensor in decoder_inputs), - input_keys=("inputs_embeds", "attention_mask", "position_ids"), - output_keys=("logits",), - graph_meta={**common_graph_meta, "component": "decoder_prefill_chunk", **cache_graph_meta}, - metadata=metadata, - )) - if "decoder_embed_chunk" in requested_set: - specs.append(ComponentModuleSpec( - component="decoder_embed_chunk", - module=decoder_embed, - example_inputs=tuple(tensor[:, :prefill_chunk, ...] for tensor in decoder_inputs), - input_keys=("inputs_embeds", "attention_mask", "position_ids"), - output_keys=("last_hidden_state",), - graph_meta={**common_graph_meta, "component": "decoder_embed_chunk", **cache_graph_meta}, - metadata=metadata, - )) - if "decoder_step" in requested_set: - specs.append(ComponentModuleSpec( - component="decoder_step", - module=decoder_last_token, - example_inputs=tuple(tensor[:, :1, ...] for tensor in decoder_inputs), - input_keys=("inputs_embeds", "attention_mask", "position_ids"), - output_keys=("logits",), - graph_meta={**common_graph_meta, "component": "decoder_step", **cache_graph_meta}, - metadata=metadata, - )) - return specs - - -def _build_lfm2_vl_multimodal_component_specs( - model: torch.nn.Module, - *, - named_tensors: dict[str, torch.Tensor], - weights_dir: str | None, - components: tuple[str, ...] | None = None, - cache_context_length: str | int | None = None, -) -> list[ComponentModuleSpec]: - input_ids = named_tensors["input_ids"] - attention_mask = named_tensors["attention_mask"] - pixel_values = named_tensors["pixel_values"] - spatial_shapes = named_tensors["spatial_shapes"] - pixel_attention_mask = named_tensors["pixel_attention_mask"] - text_context = max(8, int(os.environ.get("CACTUS_LFM2_TEXT_CONTEXT", "128") or "128")) - text_context = min(text_context, int(input_ids.shape[1])) - text_input_ids = torch.zeros( - (int(input_ids.shape[0]), text_context), - dtype=input_ids.dtype, - device=input_ids.device, - ) - text_attention_mask = torch.ones( - (int(attention_mask.shape[0]), text_context), - dtype=attention_mask.dtype, - device=attention_mask.device, - ) - - requested_components = tuple(components or ("vision_encoder", "lm_encoder", "text_lm_encoder", "decoder")) - requested_set = set(requested_components) - if not requested_set: - return [] - - expanded_components: list[str] = [] - - def _require(component: str) -> None: - if component not in expanded_components: - expanded_components.append(component) - - if "decoder" in requested_set: - _require("vision_encoder") - _require("vision_projector") - _require("lm_encoder") - _require("text_lm_encoder") - _require("decoder") - _require("text_decoder") - _require("lm_encoder_step") - _require("lm_encoder_text_chunk") - _require("decoder_prefill_chunk") - _require("decoder_embed_chunk") - _require("decoder_step") - if "lm_encoder" in requested_set: - _require("vision_encoder") - _require("vision_projector") - _require("lm_encoder") - if "lm_encoder_step" in requested_set: - _require("lm_encoder_step") - if "lm_encoder_text_chunk" in requested_set: - _require("lm_encoder_text_chunk") - if "decoder_step" in requested_set: - _require("lm_encoder_step") - _require("decoder_step") - if "decoder_prefill_chunk" in requested_set: - _require("lm_encoder_text_chunk") - _require("decoder_prefill_chunk") - if "text_lm_encoder" in requested_set: - _require("text_lm_encoder") - if "text_decoder" in requested_set: - _require("text_lm_encoder") - _require("text_decoder") - if "vision_encoder" in requested_set: - _require("vision_encoder") - _require("vision_projector") - - vision_encoder = Lfm2VlVisionEncoderAdapter(model, weights_dir=weights_dir).eval() - vision_projector = Lfm2VlVisionProjectorAdapter(model, weights_dir=weights_dir).eval() - - _vemb = vision_encoder.embeddings - with torch.no_grad(): - _pos_grid = _vemb.position_embedding.weight.reshape( - _vemb.position_embedding_size, _vemb.position_embedding_size, -1 - ) - example_positional_embeddings = _vemb.resize_positional_embeddings( - _pos_grid, spatial_shapes.detach().cpu(), max_length=int(pixel_attention_mask.shape[1]) - ).to(pixel_values.dtype) - _root = _lfm2_vl_model_root(model) - _factor = int(_root.multi_modal_projector.factor) - _proj_in = int(_vemb.embed_dim) * _factor * _factor - _max_subimage_tokens = int(getattr(_root.config, "max_image_tokens", 256) or 256) - example_vision_features = torch.zeros((_max_subimage_tokens, _proj_in), dtype=pixel_values.dtype) - - lm_encoder = Lfm2VlLMEncoderAdapter(model, weights_dir=weights_dir).eval() - text_lm_encoder = Lfm2VlLMEncoderAdapter(model, weights_dir=weights_dir).eval() - decoder = Lfm2VlDecoderAdapter(model, weights_dir=weights_dir, last_token_only=False).eval() - - decoder_inputs: tuple[torch.Tensor, ...] | None = None - text_decoder_inputs: tuple[torch.Tensor, ...] | None = None - with torch.no_grad(): - if "lm_encoder" in expanded_components or "decoder" in expanded_components: - decoder_inputs = lm_encoder(input_ids, attention_mask) - if "text_lm_encoder" in expanded_components: - text_decoder_inputs = text_lm_encoder(text_input_ids, text_attention_mask) - - def _ensure_decoder_inputs() -> tuple[torch.Tensor, ...]: - nonlocal decoder_inputs - if decoder_inputs is None: - decoder_inputs = lm_encoder(input_ids, attention_mask) - return decoder_inputs - - common_graph_meta = { - "weights_dir": weights_dir, - "task": "multimodal_causal_lm_logits", - "adapter_family": "lfm2_vl", - } - metadata = {"family": "lfm2_vl", "task": "multimodal_causal_lm_logits"} - max_cache_seq_len = _max_cache_seq_len(model, input_ids, cache_context_length, fallback_extra_tokens=512) - lm_seq = max(int(input_ids.shape[1]), int(max_cache_seq_len)) - lm_example_input_ids = torch.zeros((int(input_ids.shape[0]), lm_seq), dtype=input_ids.dtype, device=input_ids.device) - lm_example_attention_mask = torch.ones((int(attention_mask.shape[0]), lm_seq), dtype=attention_mask.dtype, device=attention_mask.device) - specs: list[ComponentModuleSpec] = [] - if "vision_encoder" in expanded_components: - specs.append(ComponentModuleSpec( - component="vision_encoder", - module=vision_encoder, - example_inputs=(pixel_values, pixel_attention_mask, example_positional_embeddings), - input_keys=("pixel_values", "pixel_attention_mask", "positional_embeddings"), - output_keys=("last_hidden_state",), - graph_meta={**common_graph_meta, "component": "vision_encoder"}, - metadata=metadata, - npu_runtime_input_count=3, - )) - if "vision_projector" in expanded_components: - specs.append(ComponentModuleSpec( - component="vision_projector", - module=vision_projector, - example_inputs=(example_vision_features,), - input_keys=("vision_features",), - output_keys=("image_features",), - graph_meta={**common_graph_meta, "component": "vision_projector"}, - metadata=metadata, - )) - if "lm_encoder" in expanded_components: - specs.append(ComponentModuleSpec( - component="lm_encoder", - module=lm_encoder, - example_inputs=(lm_example_input_ids, lm_example_attention_mask), - input_keys=("input_ids", "attention_mask"), - output_keys=("inputs_embeds", "attention_mask", "position_ids"), - graph_meta={**common_graph_meta, "component": "lm_encoder"}, - metadata=metadata, - )) - if "text_lm_encoder" in expanded_components: - if text_decoder_inputs is None: - raise RuntimeError("LFM2-VL text_lm_encoder spec requires precomputed text decoder inputs") - specs.append(ComponentModuleSpec( - component="text_lm_encoder", - module=text_lm_encoder, - example_inputs=(text_input_ids, text_attention_mask), - input_keys=("input_ids", "attention_mask"), - output_keys=("inputs_embeds", "attention_mask", "position_ids"), - graph_meta={**common_graph_meta, "component": "text_lm_encoder"}, - metadata=metadata, - )) - if "decoder" in expanded_components: - if decoder_inputs is None: - decoder_inputs = text_decoder_inputs - if decoder_inputs is None: - raise RuntimeError("LFM2-VL decoder spec requires precomputed decoder inputs") - specs.append(ComponentModuleSpec( - component="decoder", - module=decoder, - example_inputs=decoder_inputs, - input_keys=("inputs_embeds", "attention_mask", "position_ids"), - output_keys=("logits",), - graph_meta={**common_graph_meta, "component": "decoder"}, - metadata=metadata, - )) - if "text_decoder" in expanded_components: - if text_decoder_inputs is None: - raise RuntimeError("LFM2-VL text_decoder spec requires precomputed text decoder inputs") - specs.append(ComponentModuleSpec( - component="text_decoder", - module=decoder, - example_inputs=text_decoder_inputs, - input_keys=("inputs_embeds", "attention_mask", "position_ids"), - output_keys=("logits",), - graph_meta={**common_graph_meta, "component": "text_decoder"}, - metadata=metadata, - )) - specs.extend(_lfm2_chunked_pipeline_specs( - model, - requested_set=set(expanded_components), - input_ids=input_ids, - weights_dir=weights_dir, - cache_context_length=cache_context_length, - common_graph_meta=common_graph_meta, - metadata=metadata, - decoder_inputs_fn=_ensure_decoder_inputs, - )) - return specs - - -def build_component_module_specs( - model: torch.nn.Module, - *, - task: str, - named_tensors: dict[str, torch.Tensor], - weights_dir: str | None = None, - inputs_metadata: dict[str, object] | None = None, - components: tuple[str, ...] | None = None, - cache_context_length: str | int | None = None, - dynamic_batch: bool = False, - max_slots: int = 1, -) -> list[ComponentModuleSpec] | None: - family = _family_key(model) - if family == "qwen3_5" and task == "multimodal_causal_lm_logits": - return _build_qwen3_5_multimodal_component_specs( - model, - named_tensors=named_tensors, - weights_dir=weights_dir, - components=components, - cache_context_length=cache_context_length, - ) - if family in {"qwen3", "qwen3_5"} and task == "causal_lm_logits": - return _build_qwen_causal_lm_component_specs( - model, - named_tensors=named_tensors, - weights_dir=weights_dir, - components=components, - cache_context_length=cache_context_length, - ) - if family == "gemma3" and task == "causal_lm_logits": - return _build_gemma3_causal_lm_component_specs( - model, - named_tensors=named_tensors, - weights_dir=weights_dir, - components=components, - cache_context_length=cache_context_length, - ) - if family == "gemma4" and task == "multimodal_causal_lm_logits": - return _build_gemma4_multimodal_component_specs( - model, - named_tensors=named_tensors, - weights_dir=weights_dir, - components=components, - cache_context_length=cache_context_length, - dynamic_batch=dynamic_batch, - max_slots=max_slots, - ) - if family == "lfm2_vl" and task == "multimodal_causal_lm_logits": - return _build_lfm2_vl_multimodal_component_specs( - model, - named_tensors=named_tensors, - weights_dir=weights_dir, - components=components, - cache_context_length=cache_context_length, - ) - if family == "lfm2" and task == "causal_lm_logits": - return _build_lfm2_causal_lm_component_specs( - model, - named_tensors=named_tensors, - weights_dir=weights_dir, - components=components, - cache_context_length=cache_context_length, - ) - if family == "needle" and task == "causal_lm_logits": - return _build_needle_causal_lm_component_specs( - model, - named_tensors=named_tensors, - weights_dir=weights_dir, - ) - if family == "parakeet_tdt" and task == "tdt_transcription": - from cactus.transpile.tdt_runtime import build_parakeet_tdt_component_specs - - return build_parakeet_tdt_component_specs( - model, - named_tensors=named_tensors, - weights_dir=weights_dir, - ) - if family == "whisper" and task == "seq2seq_transcription": - return _build_whisper_seq2seq_component_specs( - model, - named_tensors=named_tensors, - inputs_metadata=inputs_metadata, - weights_dir=weights_dir, - ) - if family == "nomic" and task == "text_embedding": - return _build_nomic_text_embedding_component_specs( - model, - named_tensors=named_tensors, - weights_dir=weights_dir, - ) - return None - - -def canonicalize_model_interface( - model: torch.nn.Module, - task: str = "causal_lm_logits", - *, - input_names: tuple[str, ...] | None = None, - weights_dir: str | None = None, - inputs_metadata: dict[str, object] | None = None, -) -> CanonicalizedModel: - family = _family_key(model) - adapter_factory: Callable[[torch.nn.Module], torch.nn.Module] - resolved_input_names = tuple(input_names or ()) - padding_token_id_value = None - if isinstance(inputs_metadata, dict): - raw_padding_token_id = inputs_metadata.get("padding_token_id") - if isinstance(raw_padding_token_id, int): - padding_token_id_value = int(raw_padding_token_id) - - if task == "causal_lm_logits": - if family == "gemma": - adapter_factory = lambda inner_model: GemmaCausalLMLogitsAdapter( # type: ignore[assignment] - inner_model, - pad_token_id=padding_token_id_value, - ) - elif family == "gemma4": - adapter_factory = lambda inner_model: Gemma4CausalLMLogitsAdapter( # type: ignore[assignment] - inner_model, - pad_token_id=padding_token_id_value, - ) - elif family == "gemma3": - adapter_factory = lambda inner_model: Gemma3CausalLMLogitsAdapter( # type: ignore[assignment] - inner_model, - pad_token_id=padding_token_id_value, - ) - elif family == "qwen3_5": - adapter_factory = lambda inner_model: Qwen35CausalLMLogitsAdapter( # type: ignore[assignment] - inner_model, - pad_token_id=padding_token_id_value, - ) - elif family == "qwen3": - adapter_factory = lambda inner_model: Qwen3CausalLMLogitsAdapter( # type: ignore[assignment] - inner_model, - pad_token_id=padding_token_id_value, - ) - elif family in {"lfm2", "lfm2_vl", "lfm2_moe"}: - adapter_factory = lambda inner_model: Lfm2CausalLMLogitsAdapter( # type: ignore[assignment] - inner_model, - pad_token_id=padding_token_id_value, - ) - else: - adapter_factory = lambda inner_model: CausalLMLogitsAdapter( # type: ignore[assignment] - inner_model, - pad_token_id=padding_token_id_value, - ) - resolved_input_names = ("input_ids",) - elif task == "multimodal_causal_lm_logits": - if family not in {"gemma4", "lfm2_vl", "qwen3_5"}: - raise NotImplementedError(f"{type(model).__name__} does not support task={task}") - if family == "qwen3_5": - if not resolved_input_names: - resolved_input_names = ( - "input_ids", - "attention_mask", - "position_ids", - "pixel_values", - "image_grid_thw", - ) - adapter_factory = lambda inner_model: BoundInputAdapter( # type: ignore[assignment] - inner_model, - input_names=resolved_input_names, - family="qwen3_5", - metadata_task="multimodal_causal_lm_logits", - ) - elif family == "lfm2_vl": - if not resolved_input_names: - resolved_input_names = ( - "input_ids", - "attention_mask", - "pixel_values", - "spatial_shapes", - "pixel_attention_mask", - ) - adapter_factory = lambda inner_model: Lfm2VlMultimodalCausalLMLogitsAdapter( # type: ignore[assignment] - inner_model, - input_names=resolved_input_names, - ) - else: - if not resolved_input_names: - resolved_input_names = ( - "input_ids", - "attention_mask", - "token_type_ids", - "pixel_values", - "pixel_position_ids", - "input_features", - "input_features_mask", - ) - adapter_factory = lambda inner_model: Gemma4MultimodalCausalLMLogitsAdapter( # type: ignore[assignment] - inner_model, - input_names=resolved_input_names, - weights_dir=weights_dir, - ) - elif task == "ctc_logits": - if not resolved_input_names: - resolved_input_names = _infer_input_names( - model, - preferred=("input_values", "input_features", "attention_mask"), - ) - adapter_factory = lambda inner_model: CTCLogitsAdapter( # type: ignore[assignment] - inner_model, - input_names=resolved_input_names, - family=family, - ) - elif task == "encoder_hidden_states": - if not resolved_input_names: - resolved_input_names = _infer_input_names( - model, - preferred=("input_features", "input_values", "attention_mask"), - ) - adapter_factory = lambda inner_model: EncoderHiddenStatesAdapter( # type: ignore[assignment] - inner_model, - input_names=resolved_input_names, - family=family, - ) - elif task == "text_embedding": - if not resolved_input_names: - resolved_input_names = ("input_ids", "attention_mask") - adapter_factory = NomicTextEmbeddingAdapter - else: - raise NotImplementedError(f"unsupported task={task}") - - return CanonicalizedModel( - module=adapter_factory(model).eval(), - task=task, - family=family, - input_names=resolved_input_names, - ) diff --git a/python/cactus/transpile/model_patterns.py b/python/cactus/transpile/model_patterns.py deleted file mode 100644 index fa226e9b0..000000000 --- a/python/cactus/transpile/model_patterns.py +++ /dev/null @@ -1,64 +0,0 @@ -from __future__ import annotations - -from dataclasses import dataclass - - -@dataclass(frozen=True) -class GoldPattern: - name: str - semantic_ops: tuple[str, ...] - description: str - - -GOLD_PATTERNS: tuple[GoldPattern, ...] = ( - GoldPattern( - name="decoder_attention_gqa", - semantic_ops=("rms_norm", "rope", "attention", "attention_int8_hybrid"), - description=( - "Q/K/V projections from the same normalized hidden state; Q and K optionally get per-head RMSNorm, " - "Q/K are rotary-embedded, then Cactus attention consumes native Q-head and KV-head counts." - ), - ), - GoldPattern( - name="gated_mlp_gelu", - semantic_ops=("matmul", "gelu", "multiply"), - description="Two projection branches from the same normalized input: GELU(gate) * up, then a down projection.", - ), - GoldPattern( - name="gated_mlp_silu", - semantic_ops=("matmul", "silu", "multiply"), - description="Two projection branches from the same normalized input: SiLU(gate) * up, then a down projection.", - ), - GoldPattern( - name="decoder_block_post_attn_norm", - semantic_ops=("rms_norm", "attention", "add_clipped"), - description=( - "Pre-norm attention block with post-attention RMSNorm before the residual add, followed by a second " - "pre-FFN RMSNorm and a post-FFN RMSNorm before the final residual add." - ), - ), - GoldPattern( - name="decoder_block_simple_residual", - semantic_ops=("rms_norm", "attention", "add"), - description=( - "Pre-norm attention block with a direct residual add after attention, followed by a second RMSNorm and " - "a gated MLP branch with another direct residual add." - ), - ), - GoldPattern( - name="gemma4_partial_rope_attention", - semantic_ops=("rms_norm", "rope", "concat", "attention"), - description=( - "Gemma4 attention where only a prefix of the head dimension is rotary-embedded, optionally with shared " - "K/V heads and sliding-window attention." - ), - ), - GoldPattern( - name="sliding_window_attention_mask", - semantic_ops=("attention",), - description=( - "Exported PyTorch may build a boolean mask with diff/cumsum/comparisons, but the handwritten Cactus " - "models encode the same behavior directly as an attention `window_size`." - ), - ), -) diff --git a/python/cactus/transpile/model_profiles.py b/python/cactus/transpile/model_profiles.py deleted file mode 100644 index 75580eb47..000000000 --- a/python/cactus/transpile/model_profiles.py +++ /dev/null @@ -1,305 +0,0 @@ -from __future__ import annotations - -from dataclasses import dataclass -import re -from typing import Callable - -import torch - - -@dataclass(frozen=True) -class ModelProfile: - family: str - model_types: tuple[str, ...] = () - transformer_module: str | None = None - multimodal_context_tokens: int | None = None - input_combinations: tuple[tuple[str, ...], ...] = () - model_id_aliases: tuple[tuple[str, str], ...] = () - model_id_markers: tuple[str, ...] = () - family_aliases: tuple[str, ...] = () - stop_tokens: tuple[str, ...] = () - avoid_native_loader: bool = False - cached_step_components: tuple[str, ...] = () - cached_step_skip_components: tuple[str, ...] = () - dynamic_media_step_component: str | None = None - fp16_kv_cache_components: tuple[str, ...] = () - ir_replay_components: tuple[str, ...] = () - prompt_style: str = "auto" - multimodal_preprocessor: str = "auto" - text_only_component: str | None = None - default_task: str | None = None - default_components: tuple[str, ...] = () - text_component_plans: tuple[tuple[str, tuple[str, ...]], ...] = () - default_max_new_tokens: int | None = None - needs_image: bool = False - needs_audio: bool = False - force_component_pipeline: bool = False - aliases: tuple[tuple[str, str], ...] = () - regex_aliases: tuple[tuple[str, str], ...] = () - - -GEMMA4_PROFILE = ModelProfile( - family="gemma4", - model_types=("gemma4",), - transformer_module="transformers.models.gemma4.modeling_gemma4", - multimodal_context_tokens=2048, - input_combinations=((), ("image",), ("audio",), ("image", "audio")), - model_id_aliases=( - ("gemma4", "google/gemma-4-E2B-it"), - ("gemma4-e2b", "google/gemma-4-E2B-it"), - ), - model_id_markers=("gemma-4", "gemma4"), - family_aliases=("gemma",), - stop_tokens=("", ""), - avoid_native_loader=True, - cached_step_components=("lm_encoder_step", "decoder_step"), - cached_step_skip_components=("decoder",), - dynamic_media_step_component="lm_encoder_media_step", - fp16_kv_cache_components=("decoder_prefill_chunk", "decoder_step"), - ir_replay_components=( - "vision_encoder", - "audio_encoder", - "decoder_prefill_chunk", - "decoder_step", - ), - prompt_style="gemma4", - multimodal_preprocessor="gemma4", - default_task="multimodal_causal_lm_logits", - default_components=("vision_encoder", "audio_encoder", "lm_encoder", "decoder"), - needs_image=True, - needs_audio=True, - force_component_pipeline=True, -) - - -LFM2_PROFILE = ModelProfile( - family="lfm2", - model_types=("lfm2_vl", "lfm2"), - multimodal_context_tokens=256, - input_combinations=((), ("image",)), - model_id_aliases=(("lfm", "LiquidAI/LFM2-VL-450M"),), - model_id_markers=("lfm2-vl", "lfm-vl"), - family_aliases=("lfm2_vl", "lfm"), - stop_tokens=("<|im_end|>",), - avoid_native_loader=True, - cached_step_components=("lm_encoder_step", "decoder_step"), - cached_step_skip_components=("decoder",), - fp16_kv_cache_components=("decoder_prefill_chunk", "decoder_step"), - ir_replay_components=("decoder_prefill_chunk", "decoder_step"), - prompt_style="lfm_chat", - multimodal_preprocessor="lfm2_vl", - text_only_component="text_lm_encoder", - default_task="multimodal_causal_lm_logits", - default_components=("vision_encoder", "lm_encoder", "decoder"), - text_component_plans=( - ("lfm2forcausallm", ("decoder_step", "lm_encoder_step", "lm_encoder_text_chunk", - "decoder_prefill_chunk")), - ), - needs_image=True, - force_component_pipeline=True, -) - - -PARAKEET_TDT_PROFILE = ModelProfile( - family="parakeet_tdt", - model_types=("parakeet_tdt",), - model_id_aliases=( - ("parakeet", "nvidia/parakeet-tdt-0.6b-v3"), - ("parakeet-tdt", "nvidia/parakeet-tdt-0.6b-v3"), - ), - model_id_markers=("parakeet-tdt",), - avoid_native_loader=True, - default_task="tdt_transcription", - default_components=("audio_encoder", "decoder"), - needs_audio=True, - force_component_pipeline=True, - aliases=( - ("decoder.prediction.embed.weight", "decoder.embedding.weight"), - ("joint.enc.weight", "encoder_projector.weight"), - ("joint.enc.bias", "encoder_projector.bias"), - ("joint.pred.weight", "decoder.decoder_projector.weight"), - ("joint.pred.bias", "decoder.decoder_projector.bias"), - ("joint.joint_net.2.weight", "joint.head.weight"), - ("joint.joint_net.2.bias", "joint.head.bias"), - ("encoder.pre_encode.out.weight", "encoder.subsampling.linear.weight"), - ("encoder.pre_encode.out.bias", "encoder.subsampling.linear.bias"), - ), - regex_aliases=( - ( - r"encoder\.layers\.(\d+)\.self_attn\.(q|k|v)_proj\.weight", - r"encoder.layers.\1.self_attn.linear_\2.weight", - ), - ( - r"encoder\.layers\.(\d+)\.self_attn\.o_proj\.weight", - r"encoder.layers.\1.self_attn.linear_out.weight", - ), - ( - r"encoder\.layers\.(\d+)\.self_attn\.relative_k_proj\.weight", - r"encoder.layers.\1.self_attn.linear_pos.weight", - ), - ( - r"encoder\.layers\.(\d+)\.self_attn\.bias_u", - r"encoder.layers.\1.self_attn.pos_bias_u", - ), - ( - r"encoder\.layers\.(\d+)\.self_attn\.bias_v", - r"encoder.layers.\1.self_attn.pos_bias_v", - ), - ( - r"encoder\.layers\.(\d+)\.conv\.norm\.(.+)", - r"encoder.layers.\1.conv.batch_norm.\2", - ), - ), -) - - -WHISPER_PROFILE = ModelProfile( - family="whisper", - model_types=("whisper",), - model_id_aliases=(("whisper", "openai/whisper-small"),), - model_id_markers=("whisper",), - family_aliases=("whisperforconditionalgeneration",), - avoid_native_loader=True, - cached_step_components=("decoder_step",), - cached_step_skip_components=("decoder",), - fp16_kv_cache_components=("decoder_step",), - default_task="seq2seq_transcription", - default_components=("audio_encoder", "decoder"), - needs_audio=True, - force_component_pipeline=True, -) - - -QWEN_PROFILE = ModelProfile( - family="qwen", - model_types=("qwen", "qwen2", "qwen3", "qwen3_5", "qwen3.5", "qwen3_vl"), - multimodal_context_tokens=512, - input_combinations=((), ("image",)), - model_id_aliases=( - ("qwen", "Qwen/Qwen3.5-0.8B"), - ("qwen3.5", "Qwen/Qwen3.5-0.8B"), - ("qwen35", "Qwen/Qwen3.5-0.8B"), - ("qwen3.5-0.8b", "Qwen/Qwen3.5-0.8B"), - ("qwen3", "Qwen/Qwen3-1.7B"), - ), - model_id_markers=("qwen",), - family_aliases=("qwen2", "qwen3", "qwen3_5", "qwen3.5", "qwen3_vl"), - stop_tokens=("<|im_end|>",), - avoid_native_loader=True, - cached_step_components=("decoder_media_step",), - cached_step_skip_components=("decoder",), - fp16_kv_cache_components=("decoder_prefill_chunk", "decoder_media_step", "decoder_step"), - prompt_style="qwen_chat", - multimodal_preprocessor="qwen3_5", - default_task="multimodal_causal_lm_logits", - default_components=( - "vision_encoder", - "lm_encoder", - "decoder", - "lm_encoder_text_chunk", - "decoder_prefill_chunk", - "lm_encoder_step", - "decoder_media_step", - "decoder_step", - ), - text_component_plans=( - ("qwen3forcausallm", ("decoder_step", "lm_encoder_step", "lm_encoder_text_chunk", - "decoder_prefill_chunk", "decoder_embed_chunk", "decoder_media_step")), - ), - needs_image=True, - force_component_pipeline=True, -) - - -NEEDLE_PROFILE = ModelProfile( - family="needle", - model_types=("needle",), - model_id_markers=("needle",), - avoid_native_loader=True, - default_task="causal_lm_logits", - default_components=("source_encoder", "decoder_cross_kv", "decoder_step"), - default_max_new_tokens=1022, - force_component_pipeline=True, -) - - -PROFILES: tuple[ModelProfile, ...] = ( - GEMMA4_PROFILE, - LFM2_PROFILE, - PARAKEET_TDT_PROFILE, - WHISPER_PROFILE, - QWEN_PROFILE, - NEEDLE_PROFILE, -) - - -def profile_for_model_type(model_type: str) -> ModelProfile | None: - normalized = str(model_type or "").strip().lower() - for profile in PROFILES: - if normalized in profile.model_types: - return profile - return None - - -def profile_for_family(family: str) -> ModelProfile | None: - normalized = str(family or "").strip().lower() - for profile in PROFILES: - if normalized == profile.family or normalized in profile.family_aliases: - return profile - return None - - -def profile_for_model_id(model_id: str) -> ModelProfile | None: - normalized = str(model_id or "").strip().lower() - if not normalized: - return None - alias_map = model_id_alias_map() - normalized = alias_map.get(normalized, normalized).lower() - for profile in PROFILES: - for _, profile_model_id in profile.model_id_aliases: - if normalized == profile_model_id.lower(): - return profile - if any(marker and marker in normalized for marker in profile.model_id_markers): - return profile - return None - - -def model_id_alias_map() -> dict[str, str]: - aliases: dict[str, str] = {} - for profile in PROFILES: - for alias, model_id in profile.model_id_aliases: - aliases[str(alias).strip().lower()] = str(model_id).strip() - return aliases - - -def multimodal_context_tokens_for_model_type(model_type: str, default: int) -> int: - profile = profile_for_model_type(model_type) - if profile is not None and profile.multimodal_context_tokens is not None: - return max(0, int(profile.multimodal_context_tokens)) - return max(0, int(default)) - - -def add_tensor_aliases( - state_dict: dict[str, torch.Tensor], - profile: ModelProfile, - *, - derived_aliases: Callable[[dict[str, torch.Tensor]], None] | None = None, -) -> dict[str, torch.Tensor]: - def alias(target: str, source: str) -> None: - if target not in state_dict and source in state_dict: - state_dict[target] = state_dict[source] - - for target, source in profile.aliases: - alias(target, source) - - for source_key in list(state_dict): - for source_pattern, target_template in profile.regex_aliases: - match = re.fullmatch(source_pattern, source_key) - if match is None: - continue - alias(match.expand(target_template), source_key) - - if derived_aliases is not None: - derived_aliases(state_dict) - - return state_dict diff --git a/python/cactus/transpile/multimodal_runtime.py b/python/cactus/transpile/multimodal_runtime.py deleted file mode 100644 index 3747afcf6..000000000 --- a/python/cactus/transpile/multimodal_runtime.py +++ /dev/null @@ -1,722 +0,0 @@ -from __future__ import annotations - -from pathlib import Path - -import numpy as np -import torch - -from cactus.transpile.audio_preprocess import audio_duration_limit_seconds -from cactus.transpile.audio_preprocess import prepare_native_gemma4_audio_features -from cactus.transpile.media_limits import resize_static_image -from cactus.transpile.runtime_support import patch_torch_flex_attention_compat -from cactus.transpile.runtime_support import patch_transformers_torchvision_probe -from cactus.transpile.runtime_support import PreparedInputs - - -def _resolve_audio_sample_rate(processor: object) -> int: - for attr_name in ("feature_extractor", "tokenizer"): - child = getattr(processor, attr_name, None) - sample_rate = getattr(child, "sampling_rate", None) - if isinstance(sample_rate, int) and sample_rate > 0: - return sample_rate - sample_rate = getattr(processor, "sampling_rate", None) - if isinstance(sample_rate, int) and sample_rate > 0: - return sample_rate - return 16000 - - -_GEMMA4_MULTIMODAL_INPUT_ORDER = ( - "input_ids", - "attention_mask", - "token_type_ids", - "pixel_values", - "pixel_position_ids", - "input_features", - "input_features_mask", -) - - -def _normalize_multimodal_prompt( - prompt: str, - *, - image_token: str | None, - num_images: int, - audio_token: str | None, - num_audio_segments: int, -) -> str: - normalized = prompt.strip() - prefixes: list[str] = [] - - if image_token and num_images > 0: - image_count = normalized.count(image_token) - if image_count < num_images: - prefixes.append(" ".join(image_token for _ in range(num_images - image_count))) - if audio_token and num_audio_segments > 0: - audio_count = normalized.count(audio_token) - if audio_count < num_audio_segments: - prefixes.append(" ".join(audio_token for _ in range(num_audio_segments - audio_count))) - - if prefixes: - prefix = "\n".join(part for part in prefixes if part) - if normalized: - return f"{prefix}\n{normalized}" - return prefix - return normalized - - -def _build_gemma4_chat_prompt( - *, - prompt: str, - image_token: str | None, - num_images: int, - audio_token: str | None, - num_audio_segments: int, - system_prompt: str = "", - enable_thinking_if_supported: bool = False, -) -> str: - result = "" - normalized_system = system_prompt.strip() - if enable_thinking_if_supported or normalized_system: - result += "<|turn>system\n" - if enable_thinking_if_supported: - result += "<|think|>" - result += normalized_system - result += "\n" - - result += "<|turn>user\n" - if image_token and num_images > 0: - for _ in range(num_images): - result += f"\n\n{image_token}\n\n" - result += prompt - if audio_token and num_audio_segments > 0: - result += "".join(audio_token for _ in range(num_audio_segments)) - result += "\n" - result += "<|turn>model\n" - return result - - -def _gemma4_split_cactus_newline_token_merges(batch: object) -> None: - input_ids = batch.get("input_ids") if hasattr(batch, "get") else None - if not isinstance(input_ids, torch.Tensor) or input_ids.ndim != 2: - return - - expansions = { - 108: (107, 107), - 109: (107, 107, 107), - } - if not any(int(token) in expansions for token in input_ids.reshape(-1).tolist()): - return - - lengths: list[list[int]] = [] - max_len = 0 - for row in input_ids.detach().cpu().tolist(): - row_lengths = [len(expansions.get(int(token), (int(token),))) for token in row] - lengths.append(row_lengths) - max_len = max(max_len, sum(row_lengths)) - - for key in ("input_ids", "attention_mask", "token_type_ids"): - value = batch.get(key) if hasattr(batch, "get") else None - if not isinstance(value, torch.Tensor) or value.ndim != 2: - continue - expanded = torch.full( - (value.shape[0], max_len), - 0, - dtype=value.dtype, - device=value.device, - ) - for row_idx, row in enumerate(value.detach().cpu().tolist()): - out: list[int] = [] - for token_idx, item in enumerate(row): - if key == "input_ids": - out.extend(expansions.get(int(item), (int(item),))) - else: - out.extend([int(item)] * lengths[row_idx][token_idx]) - expanded[row_idx, : len(out)] = torch.tensor(out, dtype=value.dtype, device=value.device) - batch[key] = expanded - - -def _load_image_inputs(image_files: tuple[str, ...]) -> list[object]: - if not image_files: - return [] - try: - from PIL import Image # type: ignore - except Exception as exc: # pragma: no cover - raise RuntimeError(f"Pillow is required for --image-file: {exc}") from exc - - images: list[object] = [] - for image_file in image_files: - path = Path(image_file).resolve() - if not path.exists(): - raise RuntimeError(f"image_file does not exist: {path}") - with Image.open(path) as image: - images.append(resize_static_image(image.convert("RGB")).copy()) - return images - - -def _get_processor_image_attr(processor: object, name: str, default: object) -> object: - image_processor = getattr(processor, "image_processor", None) - if image_processor is not None and hasattr(image_processor, name): - return getattr(image_processor, name) - if isinstance(image_processor, dict) and name in image_processor: - return image_processor[name] - return default - - -def _image_channel_array(value: object, default: float) -> np.ndarray: - if isinstance(value, (list, tuple, np.ndarray)): - array = np.asarray(value, dtype=np.float32).reshape(-1) - if array.size >= 3: - return array[:3] - if array.size == 1: - return np.full((3,), float(array[0]), dtype=np.float32) - if isinstance(value, (int, float)): - return np.full((3,), float(value), dtype=np.float32) - return np.full((3,), float(default), dtype=np.float32) - - -def _prepare_gemma4_native_image_tensors( - processor: object, - image_files: tuple[str, ...], -) -> tuple[torch.Tensor, torch.Tensor] | None: - if not image_files: - return None - try: - from PIL import Image # type: ignore - except Exception as exc: # pragma: no cover - raise RuntimeError(f"Pillow is required for Gemma4 native image preprocessing: {exc}") from exc - - patch_size = int(_get_processor_image_attr(processor, "patch_size", 16)) - pooling_kernel_size = int(_get_processor_image_attr(processor, "pooling_kernel_size", 3)) - max_soft_tokens = int(_get_processor_image_attr(processor, "max_soft_tokens", 280)) - do_rescale = bool(_get_processor_image_attr(processor, "do_rescale", True)) - do_normalize = bool(_get_processor_image_attr(processor, "do_normalize", False)) - rescale_factor = float(_get_processor_image_attr(processor, "rescale_factor", 1.0 / 255.0)) if do_rescale else 1.0 - image_mean = ( - _image_channel_array(_get_processor_image_attr(processor, "image_mean", 0.0), 0.0) - if do_normalize - else np.zeros((3,), dtype=np.float32) - ) - image_std = ( - _image_channel_array(_get_processor_image_attr(processor, "image_std", 1.0), 1.0) - if do_normalize - else np.ones((3,), dtype=np.float32) - ) - max_patches = max_soft_tokens * pooling_kernel_size * pooling_kernel_size - side_multiple = pooling_kernel_size * patch_size - patch_dim = 3 * patch_size * patch_size - if patch_size <= 0 or pooling_kernel_size <= 0 or max_patches <= 0: - return None - - try: - resample_bilinear = Image.Resampling.BILINEAR - except AttributeError: # pragma: no cover - resample_bilinear = Image.BILINEAR - - pixel_batches: list[np.ndarray] = [] - position_batches: list[np.ndarray] = [] - for image_file in image_files: - path = Path(image_file).resolve() - with Image.open(path) as image: - rgb = image.convert("RGB") - width, height = rgb.size - target_pixels = float(max_patches * patch_size * patch_size) - factor = float(np.sqrt(target_pixels / max(1.0, float(width * height)))) - target_h = int(np.floor(factor * height / side_multiple)) * side_multiple - target_w = int(np.floor(factor * width / side_multiple)) * side_multiple - if target_h == 0: - target_h = side_multiple - if target_w == 0: - target_w = side_multiple - if (target_w, target_h) != rgb.size: - rgb = rgb.resize((target_w, target_h), resample=resample_bilinear) - array = np.asarray(rgb, dtype=np.float32) * rescale_factor - array = (array - image_mean.reshape(1, 1, 3)) / image_std.reshape(1, 1, 3) - - patch_h = target_h // patch_size - patch_w = target_w // patch_size - num_patches = patch_h * patch_w - if num_patches > max_patches: - raise RuntimeError( - f"Gemma4 native image preprocessing produced {num_patches} patches, " - f"but max_patches={max_patches}" - ) - chw = np.transpose(array, (2, 0, 1)) - patches = ( - chw.reshape(3, patch_h, patch_size, patch_w, patch_size) - .transpose(1, 3, 2, 4, 0) - .reshape(num_patches, patch_dim) - ) - - padded_patches = np.zeros((max_patches, patch_dim), dtype=np.float32) - padded_patches[:num_patches] = patches - positions = np.full((max_patches, 2), -1, dtype=np.int64) - valid_positions = np.zeros((num_patches, 2), dtype=np.int64) - for patch_y in range(patch_h): - row_start = patch_y * patch_w - valid_positions[row_start : row_start + patch_w, 0] = np.arange(patch_w, dtype=np.int64) - valid_positions[row_start : row_start + patch_w, 1] = patch_y - positions[:num_patches] = valid_positions - - pixel_batches.append(padded_patches) - position_batches.append(positions) - - return ( - torch.from_numpy(np.stack(pixel_batches, axis=0)), - torch.from_numpy(np.stack(position_batches, axis=0)), - ) - - -def _gemma4_image_soft_token_counts( - pixel_position_ids: torch.Tensor, - *, - pooling_kernel_size: int, -) -> list[int]: - counts: list[int] = [] - for positions in pixel_position_ids: - valid = positions[(positions != -1).any(dim=-1)] - if valid.numel() == 0: - counts.append(0) - continue - max_x = int(valid[:, 0].max().item()) + 1 - max_y = int(valid[:, 1].max().item()) + 1 - counts.append((max_x // pooling_kernel_size) * (max_y // pooling_kernel_size)) - return counts - - -def _gemma4_audio_soft_token_count(num_frames: int) -> int: - after_stage1 = (int(num_frames) + 1) // 2 - return (after_stage1 + 1) // 2 - - -def _resolve_gemma4_audio_mels(processor: object, default: int = 128) -> int: - for child_name in ("feature_extractor", "audio_processor", "audio_feature_extractor"): - child = getattr(processor, child_name, None) - if child is None: - continue - for attr_name in ("feature_size", "num_mel_bins", "n_mels", "input_feat_size"): - value = getattr(child, attr_name, None) - if isinstance(value, int) and value > 0: - return int(value) - for attr_name in ("feature_size", "num_mel_bins", "n_mels", "input_feat_size"): - value = getattr(processor, attr_name, None) - if isinstance(value, int) and value > 0: - return int(value) - return default - - -def _prepare_gemma4_processor_audio_tensors( - processor: object, - audio_waveforms: list[np.ndarray], - *, - torch_dtype: torch.dtype, - max_seconds: float | None, -) -> tuple[torch.Tensor, torch.Tensor, int] | None: - if not audio_waveforms: - return None - extractor = ( - getattr(processor, "audio_processor", None) - or getattr(processor, "feature_extractor", None) - or getattr(processor, "audio_feature_extractor", None) - ) - if not callable(extractor): - return None - batch = extractor(audio_waveforms, return_tensors="pt") - input_features = batch.get("input_features") if hasattr(batch, "get") else None - input_features_mask = batch.get("input_features_mask") if hasattr(batch, "get") else None - if not isinstance(input_features, torch.Tensor) or not isinstance(input_features_mask, torch.Tensor): - return None - if input_features.ndim != 3 or input_features_mask.ndim != 2: - return None - - active_frames = int(input_features_mask[0].sum().item()) - output_frames = int(input_features.shape[1]) - if max_seconds is not None and max_seconds > 0.0: - # Preserve the existing static Gemma4 graph capacity so a single bundle - # can accept any <=30s clip while the active mask still matches HF. - max_samples = int(round(16000.0 * float(max_seconds))) - output_frames = max(output_frames, int(np.ceil((max_samples + 640) / 160.0)) + 8) - - if output_frames > int(input_features.shape[1]): - padded_features = torch.zeros( - (input_features.shape[0], output_frames, input_features.shape[2]), - dtype=input_features.dtype, - ) - padded_mask = torch.zeros( - (input_features_mask.shape[0], output_frames), - dtype=input_features_mask.dtype, - ) - padded_features[:, : input_features.shape[1], :] = input_features - padded_mask[:, : input_features_mask.shape[1]] = input_features_mask - input_features = padded_features - input_features_mask = padded_mask - - return ( - input_features.to(dtype=torch_dtype), - input_features_mask.to(dtype=torch.bool), - active_frames, - ) - - -def _pad_gemma4_audio_batch_to_static_limit( - batch: object, - *, - max_seconds: float | None, -) -> int | None: - input_features = batch.get("input_features") if hasattr(batch, "get") else None - input_features_mask = batch.get("input_features_mask") if hasattr(batch, "get") else None - if not isinstance(input_features, torch.Tensor) or not isinstance(input_features_mask, torch.Tensor): - return None - active_frames = int(input_features_mask[0].sum().item()) - if max_seconds is None or max_seconds <= 0.0: - return active_frames - - max_samples = int(round(16000.0 * float(max_seconds))) - output_frames = max( - int(input_features.shape[1]), - int(np.ceil((max_samples + 640) / 160.0)) + 8, - ) - if output_frames > int(input_features.shape[1]): - padded_features = torch.zeros( - (input_features.shape[0], output_frames, input_features.shape[2]), - dtype=input_features.dtype, - device=input_features.device, - ) - padded_mask = torch.zeros( - (input_features_mask.shape[0], output_frames), - dtype=input_features_mask.dtype, - device=input_features_mask.device, - ) - padded_features[:, : input_features.shape[1], :] = input_features - padded_mask[:, : input_features_mask.shape[1]] = input_features_mask - batch["input_features"] = padded_features - batch["input_features_mask"] = padded_mask - return active_frames - - -def _build_gemma4_processor_chat_prompt( - processor: object, - *, - prompt: str, - image_files: tuple[str, ...], - audio_file: str | None, - system_prompt: str = "", -) -> str | None: - apply_chat_template = getattr(processor, "apply_chat_template", None) - if not callable(apply_chat_template): - return None - - content: list[dict[str, object]] = [] - for image_file in image_files: - content.append({"type": "image", "image": str(Path(image_file).expanduser().resolve())}) - if prompt.strip(): - content.append({"type": "text", "text": prompt.strip()}) - if audio_file: - content.append({"type": "audio", "audio": str(Path(audio_file).expanduser().resolve())}) - - messages: list[dict[str, object]] = [] - if system_prompt.strip(): - messages.append({"role": "system", "content": [{"type": "text", "text": system_prompt.strip()}]}) - messages.append({"role": "user", "content": content or [{"type": "text", "text": prompt.strip()}]}) - try: - return str(apply_chat_template(messages, tokenize=False, add_generation_prompt=True)) - except Exception: - return None - - -def _build_gemma4_native_chat_prompt( - *, - prompt: str, - image_soft_token_counts: tuple[int, ...], - audio_soft_token_count: int, - system_prompt: str = "", - enable_thinking_if_supported: bool = False, -) -> str: - result = "" - normalized_system = system_prompt.strip() - if enable_thinking_if_supported or normalized_system: - result += "<|turn>system\n" - if enable_thinking_if_supported: - result += "<|think|>" - result += normalized_system - result += "\n" - - result += "<|turn>user\n" - for count in image_soft_token_counts: - if count > 0: - result += "\n\n<|image>" - result += "<|image|>" * int(count) - result += "\n\n" - result += prompt - if audio_soft_token_count > 0: - result += "<|audio>" - result += "<|audio|>" * int(audio_soft_token_count) - result += "" - result += "\n" - result += "<|turn>model\n" - return result - - -def _tokenize_gemma4_native_prompt( - processor: object, - prompt: str, - *, - image_token: str | None, - audio_token: str | None, -) -> dict[str, torch.Tensor]: - tokenizer = getattr(processor, "tokenizer", processor) - if not callable(tokenizer): - raise RuntimeError("Gemma4 native prompt path requires a tokenizer") - batch = tokenizer(prompt, return_tensors="pt", add_special_tokens=False) - input_ids = batch["input_ids"] - if "attention_mask" not in batch: - batch["attention_mask"] = torch.ones_like(input_ids) - - token_type_ids = torch.zeros_like(input_ids) - convert = getattr(tokenizer, "convert_tokens_to_ids", None) - if callable(convert): - if image_token: - image_id = convert(image_token) - if isinstance(image_id, int) and image_id >= 0: - token_type_ids = torch.where(input_ids == image_id, torch.ones_like(token_type_ids), token_type_ids) - if audio_token: - audio_id = convert(audio_token) - if isinstance(audio_id, int) and audio_id >= 0: - token_type_ids = torch.where(input_ids == audio_id, torch.full_like(token_type_ids, 2), token_type_ids) - batch["token_type_ids"] = token_type_ids - return {key: value for key, value in batch.items() if isinstance(value, torch.Tensor)} - - -def prepare_gemma4_multimodal_inputs( - processor: object | None, - *, - prompt: str, - image_files: tuple[str, ...], - audio_file: str | None, - torch_dtype: torch.dtype, - system_prompt: str = "", - enable_thinking_if_supported: bool = False, - use_gemma4_chat_template: bool = False, - static_image_soft_token_counts: tuple[int, ...] | None = None, - static_audio_soft_token_count: int | None = None, - force_static_multimodal_placeholders: bool = False, -) -> PreparedInputs: - if processor is None: - raise RuntimeError("multimodal Gemma4 transpile requires an AutoProcessor with image and audio support") - - images = _load_image_inputs(image_files) - audio_waveforms: list[np.ndarray] = [] - sample_rate: int | None = None - if audio_file: - sample_rate = _resolve_audio_sample_rate(processor) - from cactus.transpile.audio_preprocess import load_audio_waveform - - audio_waveforms.append( - load_audio_waveform( - audio_file, - target_sample_rate=sample_rate, - max_seconds=audio_duration_limit_seconds(), - ) - ) - - image_token = getattr(processor, "image_token", None) - audio_token = getattr(processor, "audio_token", None) - processor_prompt = _normalize_multimodal_prompt( - prompt, - image_token=image_token if isinstance(image_token, str) else None, - num_images=len(images), - audio_token=audio_token if isinstance(audio_token, str) else None, - num_audio_segments=len(audio_waveforms), - ) - normalized_prompt = processor_prompt - if use_gemma4_chat_template: - normalized_prompt = prompt.strip() - native_image_tensors = _prepare_gemma4_native_image_tensors(processor, image_files) - native_audio: torch.Tensor | None = None - native_audio_mask: torch.Tensor | None = None - native_audio_frames: int | None = None - if audio_file: - expected_mels = _resolve_gemma4_audio_mels(processor) - try: - native_audio, native_audio_mask, native_audio_frames = prepare_native_gemma4_audio_features( - audio_file, - expected_mels=expected_mels, - torch_dtype=torch_dtype, - max_seconds=audio_duration_limit_seconds(), - pad_to_max_seconds=True, - ) - except Exception as exc: - try: - processor_audio = _prepare_gemma4_processor_audio_tensors( - processor, - audio_waveforms, - torch_dtype=torch_dtype, - max_seconds=audio_duration_limit_seconds(), - ) - except Exception: - processor_audio = None - if processor_audio is None: - print(f"note=falling back to processor gemma4 audio prompt packing: {exc}") - native_audio = None - native_audio_mask = None - native_audio_frames = None - else: - native_audio, native_audio_mask, native_audio_frames = processor_audio - - has_native_media = ( - native_image_tensors is not None - or native_audio is not None - or ( - force_static_multimodal_placeholders - and ( - bool(static_image_soft_token_counts) - or (static_audio_soft_token_count is not None and int(static_audio_soft_token_count) > 0) - ) - ) - ) - if has_native_media and (not audio_file or native_audio is not None): - image_soft_counts: tuple[int, ...] = () - if native_image_tensors is not None: - pixel_values, pixel_position_ids = native_image_tensors - pooling_kernel_size = int(_get_processor_image_attr(processor, "pooling_kernel_size", 3)) - computed_image_counts = tuple( - _gemma4_image_soft_token_counts( - pixel_position_ids, - pooling_kernel_size=pooling_kernel_size, - ) - ) - if ( - static_image_soft_token_counts is not None - and len(static_image_soft_token_counts) == len(computed_image_counts) - ): - image_soft_counts = tuple(int(count) for count in static_image_soft_token_counts) - else: - image_soft_counts = computed_image_counts - elif force_static_multimodal_placeholders and static_image_soft_token_counts is not None: - image_soft_counts = tuple(int(count) for count in static_image_soft_token_counts) - audio_soft_count = 0 - if audio_file: - audio_soft_count = ( - int(static_audio_soft_token_count) - if static_audio_soft_token_count is not None and int(static_audio_soft_token_count) > 0 - else _gemma4_audio_soft_token_count(native_audio_frames or 0) - ) - elif ( - force_static_multimodal_placeholders - and static_audio_soft_token_count is not None - and int(static_audio_soft_token_count) > 0 - ): - audio_soft_count = int(static_audio_soft_token_count) - processor_prompt = _build_gemma4_native_chat_prompt( - prompt=normalized_prompt, - image_soft_token_counts=image_soft_counts, - audio_soft_token_count=audio_soft_count, - system_prompt=system_prompt, - enable_thinking_if_supported=enable_thinking_if_supported, - ) - batch = _tokenize_gemma4_native_prompt( - processor, - processor_prompt, - image_token=image_token if isinstance(image_token, str) else None, - audio_token=audio_token if isinstance(audio_token, str) else None, - ) - if native_image_tensors is not None: - batch["pixel_values"] = pixel_values - batch["pixel_position_ids"] = pixel_position_ids - if native_audio is not None and native_audio_mask is not None: - batch["input_features"] = native_audio - batch["input_features_mask"] = native_audio_mask - batch["native_audio_frames"] = int(native_audio_frames or native_audio.shape[1]) - else: - processor_chat_prompt = None - if image_files or audio_file: - processor_chat_prompt = _build_gemma4_processor_chat_prompt( - processor, - prompt=normalized_prompt, - image_files=image_files, - audio_file=audio_file, - system_prompt=system_prompt, - ) - processor_prompt = ( - processor_chat_prompt - if processor_chat_prompt is not None - else _build_gemma4_chat_prompt( - prompt=normalized_prompt, - image_token=image_token if isinstance(image_token, str) else None, - num_images=len(images), - audio_token=audio_token if isinstance(audio_token, str) else None, - num_audio_segments=len(audio_waveforms), - system_prompt=system_prompt, - enable_thinking_if_supported=enable_thinking_if_supported, - ) - ) - batch = processor( - text=processor_prompt, - images=images or None, - audio=audio_waveforms or None, - return_tensors="pt", - ) - if "mm_token_type_ids" in batch: - batch["token_type_ids"] = batch["mm_token_type_ids"] - if "image_position_ids" in batch: - batch["pixel_position_ids"] = batch["image_position_ids"] - native_audio_frames = _pad_gemma4_audio_batch_to_static_limit( - batch, - max_seconds=audio_duration_limit_seconds(), - ) - if isinstance(native_audio_frames, int): - batch["native_audio_frames"] = native_audio_frames - _gemma4_split_cactus_newline_token_merges(batch) - if native_image_tensors is not None: - batch["pixel_values"], batch["pixel_position_ids"] = native_image_tensors - else: - batch = processor( - text=processor_prompt, - images=images or None, - audio=audio_waveforms or None, - return_tensors="pt", - ) - - ordered_keys = [ - key - for key in _GEMMA4_MULTIMODAL_INPUT_ORDER - if isinstance(batch.get(key), torch.Tensor) - ] - if not ordered_keys: - raise RuntimeError("Gemma4 multimodal processor did not return any tensor inputs") - - tensors: list[torch.Tensor] = [] - for key in ordered_keys: - value = batch[key] - if not isinstance(value, torch.Tensor): - continue - if torch.is_floating_point(value): - value = value.to(dtype=torch_dtype) - tensors.append(value) - - metadata: dict[str, object] = { - "prompt": normalized_prompt, - "processor_prompt": processor_prompt, - "image_files": [str(Path(path).resolve()) for path in image_files], - "input_shapes": { - name: list(tensor.shape) - for name, tensor in zip(ordered_keys, tensors) - }, - } - if audio_file: - metadata["audio_file"] = str(Path(audio_file).resolve()) - if sample_rate is not None: - metadata["sample_rate"] = sample_rate - native_audio_frames = batch.get("native_audio_frames") - if isinstance(native_audio_frames, int): - metadata["native_audio_frames"] = native_audio_frames - - return PreparedInputs( - names=tuple(ordered_keys[: len(tensors)]), - tensors=tuple(tensors), - metadata=metadata, - ) - - -_patch_transformers_torchvision_probe = patch_transformers_torchvision_probe -_patch_torch_flex_attention_compat = patch_torch_flex_attention_compat -_prepare_gemma4_multimodal_inputs = prepare_gemma4_multimodal_inputs diff --git a/python/cactus/transpile/normalize.py b/python/cactus/transpile/normalize.py deleted file mode 100644 index c7a0382a8..000000000 --- a/python/cactus/transpile/normalize.py +++ /dev/null @@ -1,29 +0,0 @@ -from __future__ import annotations - -from typing import Any - -import torch - -from cactus.transpile.aten_ops import canonical_ir_op - - -def normalize_target(target: Any) -> str: - return canonical_ir_op(target) - - -def dtype_to_ir(dtype: Any | None) -> str | None: - if dtype is None: - return None - - dtype_map = { - torch.float16: "fp16", - torch.float32: "fp32", - torch.float64: "fp64", - torch.bfloat16: "bf16", - torch.int8: "int8", - torch.int16: "int16", - torch.int32: "int32", - torch.int64: "int64", - torch.bool: "bool", - } - return dtype_map.get(dtype, str(dtype)) diff --git a/python/cactus/transpile/npu/README.md b/python/cactus/transpile/npu/README.md deleted file mode 100644 index 8a51e1013..000000000 --- a/python/cactus/transpile/npu/README.md +++ /dev/null @@ -1,86 +0,0 @@ -# NPU Transpiler Pipeline - -Emits CoreML `.mlpackage`s so the runtime engine can dispatch through the -Apple Neural Engine. Triggered by `--npu` on `cactus convert` (forwarded to -`python -m cactus.transpile.hf_model --npu …`), or programmatically via -`TranspileOptions(npu=True)` in `cli/model.py`. Runtime side: -`cactus-engine/src/model_npu.cpp`. - -**Scope: audio + vision encoders only.** Text-decoder prefill is -intentionally *not* on NPU — CPU prefill is the supported path. - -| Entry point | Emits | manifest key | Runtime dispatch | -|---|---|---|---| -| `run_encoder_pipeline` → `emit_audio_encoder_mlpackage` (audio.py) | audio encoder | `npu_audio_encoder` | `audio_encode_via_npu` | -| `run_encoder_pipeline` → `emit_vision_encoder_mlpackage` (vision.py) | vision encoder | `npu_vision_encoder` | `vision_encode_via_npu` | - -`run_encoder_pipeline` reuses the exact `ComponentModuleSpec` adapter -modules + example inputs the graph transpiler captures, so the NPU encoder -stays in lockstep with the CPU component graph. Auxiliary adapter inputs -(masks, position ids) are baked into the exported model as constants — -the ANE requires a static single-input signature. - -## Status - -| Target | Works | Notes | -|---|---|---| -| Parakeet TDT audio encoder | **yes — verified end-to-end** | 4.5× faster warm than CPU on M-series (~155ms vs ~693ms TTFT). Use `CACTUS_ANE_AUDIO_COMPUTE_UNITS=ALL`. | -| LFM-VL vision encoder | **yes — verified end-to-end** | ~2× faster on image-specific TTFT overhead, ~4× faster on the vision encoder graph alone (the rest of TTFT is preprocessing + LM prefill, CPU-bound). | -| Gemma 4 audio + vision encoder emit | **emit verified (int8); mlpackage output matches HF reference** | Vision int8: cos_sim 0.9995 vs HF reference (MAE 0.018). Audio int8: cos_sim 0.9963 (MAE 0.0055 after applying the cactus `1/16` audio rescale, see `_gemma4_compute_native_like_audio_features`). Vision: 162 MB; audio: 282 MB. C++ runtime dispatch path is the same code that ships Parakeet ASR + LFM-VL chat — exercised end-to-end on those models in this session. Full Gemma 4 cactus runtime end-to-end requires a v2 bundle (`cactus convert`, ~15–20 GB peak), which doesn't fit on the 16 GB host where these mlpackages were produced. | -| `coremltools.optimize.linear_quantize_weights` | works | Per-component defaults: **audio int8, vision fp16**. ANE compute is fp16 internally either way; quant only affects weight storage + memory bandwidth. Audio encoders (Conformer-style) absorb int8 quant noise well — verified on Parakeet + Gemma 4. Vision encoders (ViT-style: SigLIP2, Gemma 4 vision) are more sensitive — int4 visibly corrupts Gemma 4 vision output; fp16 is the rock-solid default given the towers are small (one-shot per request, bandwidth not critical). Override via `--npu-quantize` (forces both) or `--npu-audio-quantize` / `--npu-vision-quantize` (per-component). | - -## Why not NPU text prefill? - -Earlier iterations also emitted a text-decoder prefill `.mlpackage` -(via the now-removed `prefill.py`). Removed because: - -1. Wrapping the full text decoder (~2B params for Gemma 4 E2B) needs - ~15–20 GB peak during `coremltools.convert`, which OOMs 16 GB hosts — - regardless of whether the encoder emit succeeds. -2. CPU chunked prefill (post-PR #659) is fast enough that ANE prefill - isn't a clear win for the v2 runtime. -3. `origin/main`'s NPU prefill mlpackages were pre-built externally, - not regenerated by a transpiler — there's no precedent for v2 to - regenerate them reliably. - -The C++ runtime side (`NPUPrefill` class, `prefill_via_npu`, `load_npu_prefill`, -the `ANEPrefill` Obj-C wrapper) has been removed in lockstep. - -## Known coremltools 9.0 patches - -`coremltools_patches.py` carries op patches for gaps the v2 encoders hit -during conversion: - -- `__and__` / `__or__` / `__xor__` — bitwise ops on bool tensors -- `new_ones` — `tensor.new_ones(shape)` -- `one_hot` — cast indices to int32 (Gemma 4 vision feeds fp32 ids) -- `unfold` — translates to `mb.sliding_windows` + a transpose so the - window-size dim lands at the *end* (matching PyTorch semantics; MIL - inserts it at `axis+1` by default, which breaks Gemma 4 audio's - chunked-local attention permute chain) -- `layer_norm` epsilon dtype unification (iOS17+ enforces gamma/beta/eps - same dtype; HF fp16-traced models hit this) -- `add_fp16_cast` MIL pass skipping `layer_norm` / `batch_norm` / etc. -- mb binop scalar-cast wrappers (sub/mul/add/div SDPA mask path). - *Only Var operands get a `mb.cast` injected* — numpy arrays/scalars - are retyped via `numpy.astype` so we don't accidentally materialize a - Var in the wrong block scope when downstream passes (e.g. - `divide_to_multiply`) call `mb.mul(..., before_op=op)` from inside a - nested block. Mis-scoping trips the "var not visible in the block" - validator and aborts conversion. - -Add more here as new architectures expose new gaps. - -## Where future NPU perf gains can come from - -1. **Weight quantization (int4 / int8)** — biggest single perf knob on - ANE. Smaller weights ⇒ less memory bandwidth ⇒ faster matmuls. - Wired via `--npu-quantize`. -2. **SDPA preservation** — `torch.nn.functional.scaled_dot_product_attention` - maps to a single MIL op when coremltools sees it intact. HF defaults - to `attn_implementation="sdpa"` for compatible models, so we benefit - automatically. -3. **Per-family adapter optimizations** — e.g., the Gemma 4 vision tower - could probably be split (CPU patch embed + NPU transformer + CPU - projector, matching main's `model_siglip2.cpp` pattern) for cleaner - memory usage. Not done; whole-encoder mlpackage works today. diff --git a/python/cactus/transpile/npu/__init__.py b/python/cactus/transpile/npu/__init__.py deleted file mode 100644 index e2fc7da08..000000000 --- a/python/cactus/transpile/npu/__init__.py +++ /dev/null @@ -1,3 +0,0 @@ -from .pipeline import run_encoder_pipeline - -__all__ = ["run_encoder_pipeline"] diff --git a/python/cactus/transpile/npu/audio.py b/python/cactus/transpile/npu/audio.py deleted file mode 100644 index 6d72bedee..000000000 --- a/python/cactus/transpile/npu/audio.py +++ /dev/null @@ -1,115 +0,0 @@ -from __future__ import annotations - -from pathlib import Path -from typing import Any - -import torch - - -class AudioEncoderWrapper(torch.nn.Module): - def __init__( - self, - audio_module: torch.nn.Module, - baked_inputs: tuple[torch.Tensor, ...] = (), - ): - super().__init__() - self.audio = audio_module - self._n_baked = len(baked_inputs) - for idx, tensor in enumerate(baked_inputs): - self.register_buffer(f"_baked_{idx}", tensor, persistent=False) - - def forward(self, input_features: torch.Tensor) -> torch.Tensor: - extra = tuple(getattr(self, f"_baked_{i}") for i in range(self._n_baked)) - return self.audio(input_features, *extra) - - -def _import_coremltools() -> Any: - import coremltools as ct - from .coremltools_patches import apply_all_coremltools_patches - apply_all_coremltools_patches() - return ct - - -def _apply_weight_quantization(mlmodel: Any, bits: int) -> Any: - try: - from coremltools.optimize.coreml import ( - linear_quantize_weights, - OpLinearQuantizerConfig, - OptimizationConfig, - ) - op_config = OpLinearQuantizerConfig( - mode="linear_symmetric", - dtype=f"int{bits}", - granularity="per_channel", - ) - config = OptimizationConfig(global_config=op_config) - return linear_quantize_weights(mlmodel, config) - except Exception as exc: - print(f"npu.audio: weight quantization to int{bits} failed ({type(exc).__name__}: {exc}); keeping FP16 weights") - return mlmodel - - -def emit_audio_encoder_mlpackage( - audio_module: torch.nn.Module, - bundle_dir: Path, - *, - example_input: torch.Tensor, - baked_inputs: tuple[torch.Tensor, ...] = (), - filename: str = "audio_encoder.mlpackage", - input_name: str = "x", - output_name: str = "encoded", - minimum_deployment_target: str = "iOS17", - quantize_bits: int | None = None, -) -> str | None: - ct = _import_coremltools() - - wrapper = AudioEncoderWrapper(audio_module, baked_inputs) - wrapper.eval() - - try: - with torch.no_grad(): - exported = torch.export.export(wrapper, (example_input,)) - exported = exported.run_decompositions({}) - except Exception as exc: - print(f"npu.audio: torch.export failed ({type(exc).__name__}: {exc}); skipping mlpackage emit") - return None - - del wrapper - import gc as _gc - _gc.collect() - - if quantize_bits == 4: - minimum_deployment_target = "iOS18" - target_attr = getattr(ct.target, minimum_deployment_target, None) or ct.target.iOS17 - - from .coremltools_patches import build_cactus_pass_pipeline - try: - mlmodel = ct.convert( - exported, - inputs=[ct.TensorType(name=input_name, shape=tuple(example_input.shape))], - outputs=[ct.TensorType(name=output_name)], - compute_precision=ct.precision.FLOAT16, - convert_to="mlprogram", - minimum_deployment_target=target_attr, - pass_pipeline=build_cactus_pass_pipeline(), - ) - except Exception as exc: - print(f"npu.audio: coremltools.convert failed ({type(exc).__name__}: {exc})") - return None - - if quantize_bits is not None: - before_id = id(mlmodel) - mlmodel = _apply_weight_quantization(mlmodel, quantize_bits) - if id(mlmodel) != before_id: - print(f"npu.audio: applied int{quantize_bits} weight quantization") - - bundle_dir.mkdir(parents=True, exist_ok=True) - out_path = bundle_dir / filename - try: - mlmodel.save(str(out_path)) - except Exception as exc: - print(f"npu.audio: mlpackage save failed ({type(exc).__name__}: {exc})") - return None - - print(f"npu.audio: wrote {out_path} (input_shape={tuple(example_input.shape)})") - return filename diff --git a/python/cactus/transpile/npu/coremltools_patches.py b/python/cactus/transpile/npu/coremltools_patches.py deleted file mode 100644 index 9bb0de1cd..000000000 --- a/python/cactus/transpile/npu/coremltools_patches.py +++ /dev/null @@ -1,434 +0,0 @@ -from __future__ import annotations - -_APPLIED = False - - -def apply_all_coremltools_patches() -> None: - global _APPLIED - if _APPLIED: - return - try: - _patch_register_func_allow_dunder() - _register_new_ones_op() - _register_logical_and_op() - _override_layer_norm_translator() - _override_one_hot_translator() - _register_unfold_op() - _register_noop_assert_ops() - _patch_pipeline_remove_fuse_prelu() - _patch_mb_binops_scalar_cast() - _patch_fp16_cast_skip_layer_norm() - _APPLIED = True - print("npu.coremltools_patches: applied") - except Exception as exc: - print(f"npu.coremltools_patches: failed to apply ({type(exc).__name__}: {exc})") - - -def _patch_register_func_allow_dunder() -> None: - from coremltools.converters.mil.frontend.torch import torch_op_registry - - cls = torch_op_registry.TorchOpsRegistry - if getattr(cls.register_func, "_cactus_dunder_patched", False): - return - - - def register_func(self, func=None, torch_alias=None, override=False): - f_name = func.__name__ - all_f_names = [f_name] - if torch_alias is not None: - all_f_names.extend(torch_alias) - for name in all_f_names: - is_dunder = name.startswith("__") and name.endswith("__") - if name.endswith("_") and not is_dunder: - raise Exception( - f'Attempting to register "{name}" op. Do not register inplace ops.' - ) - if not override and name in self.name_to_func_mapping: - raise ValueError(f"Torch op {name} already registered.") - self.set_func_by_name(func, name) - - register_func._cactus_dunder_patched = True - cls.register_func = register_func - - -def _register_new_ones_op() -> None: - from coremltools.converters.mil.frontend.torch.torch_op_registry import ( - _TORCH_OPS_REGISTRY, - ) - from coremltools.converters.mil import Builder as mb - from coremltools.converters.mil.frontend.torch.ops import _get_inputs - - if "new_ones" in _TORCH_OPS_REGISTRY.name_to_func_mapping: - return - - def _to_int32(v): - if isinstance(v, list): - casted = [_to_int32(x) for x in v] - return mb.concat(values=casted, axis=0) - dtype = getattr(v, "dtype", None) - name = getattr(dtype, "__name__", str(dtype) if dtype else "") - if "int32" in name: - return v - return mb.cast(x=v, dtype="int32") - - def new_ones(context, node): - inputs = _get_inputs(context, node, min_expected=2) - size = _to_int32(inputs[1]) - res = mb.fill(shape=size, value=1.0, name=node.name) - context.add(res, node.name) - - _TORCH_OPS_REGISTRY.set_func_by_name(new_ones, "new_ones") - - -def _register_logical_and_op() -> None: - from coremltools.converters.mil.frontend.torch.torch_op_registry import ( - _TORCH_OPS_REGISTRY, - ) - from coremltools.converters.mil import Builder as mb - - def _force_bool(v): - dtype = getattr(v, "dtype", None) - name = getattr(dtype, "__name__", str(dtype) if dtype else "") - if "bool" in name: - return v - if "fp16" in name or "fp32" in name or "fp64" in name or "double" in name or "float" in name: - non_zero = mb.not_equal(x=v, y=0.0) - return non_zero - if "int" in name: - non_zero = mb.not_equal(x=v, y=0) - return non_zero - return mb.cast(x=v, dtype="bool") - - def make_logical(op_kind): - def _impl(context, node): - inputs = [context[i] for i in node.inputs] - x = _force_bool(inputs[0]) - y = _force_bool(inputs[1]) - if op_kind == "and": - out = mb.logical_and(x=x, y=y, name=node.name) - elif op_kind == "or": - out = mb.logical_or(x=x, y=y, name=node.name) - else: - out = mb.logical_xor(x=x, y=y, name=node.name) - context.add(out, node.name) - return _impl - - for tag, op_kind in [("__and_", "and"), ("__or_", "or"), ("__xor_", "xor")]: - if tag not in _TORCH_OPS_REGISTRY.name_to_func_mapping: - _TORCH_OPS_REGISTRY.set_func_by_name(make_logical(op_kind), tag) - for tag, op_kind in [("__and__", "and"), ("__or__", "or"), ("__xor__", "xor")]: - if tag not in _TORCH_OPS_REGISTRY.name_to_func_mapping: - _TORCH_OPS_REGISTRY.set_func_by_name(make_logical(op_kind), tag) - for tag in ("bitwise_and", "and"): - _TORCH_OPS_REGISTRY.set_func_by_name(make_logical("and"), tag) - for tag in ("bitwise_or", "or"): - _TORCH_OPS_REGISTRY.set_func_by_name(make_logical("or"), tag) - - -def _override_layer_norm_translator() -> None: - import numpy as np - from coremltools.converters.mil.frontend.torch.torch_op_registry import ( - _TORCH_OPS_REGISTRY, - ) - from coremltools.converters.mil.frontend.torch.ops import _get_inputs - from coremltools.converters.mil import Builder as mb - - def _dtype_name(v): - d = getattr(v, "dtype", None) - return getattr(d, "__name__", str(d) if d else "") - - def _to_typed_const(v, target): - if v is None: - return v - np_dtype = np.float16 if target == "fp16" else np.float32 - if isinstance(v, (int, float)) and not isinstance(v, bool): - return mb.const(val=np_dtype(v)) - if hasattr(v, "val") and v.val is not None and hasattr(v, "dtype"): - try: - return mb.const(val=np.asarray(v.val, dtype=np_dtype)) - except Exception: - pass - if hasattr(v, "dtype") and _dtype_name(v) != target: - return mb.cast(x=v, dtype=target) - return v - - def layer_norm(context, node): - inputs = _get_inputs(context, node, min_expected=2) - nargs = len(inputs) - x = inputs[0] - normalized_shape = inputs[1] - weight = inputs[2] if nargs > 2 else None - bias = inputs[3] if nargs > 3 else None - eps = inputs[4] if nargs > 4 else None - if eps is None: - eps = 1e-5 - - ref = weight if weight is not None else x - ref_name = _dtype_name(ref) - target = "fp16" if ("fp16" in ref_name or "float16" in ref_name) else ( - "fp32" if ("fp32" in ref_name or "float32" in ref_name) else "fp16" - ) - - if x is not None and hasattr(x, "dtype") and _dtype_name(x) != target: - x = mb.cast(x=x, dtype=target) - weight = _to_typed_const(weight, target) - bias = _to_typed_const(bias, target) - eps = _to_typed_const(eps, target) - - out = mb.layer_norm( - x=x, - axes=list(range(-len(normalized_shape.val), 0)), - gamma=weight, - beta=bias, - epsilon=eps, - name=node.name, - ) - - if node.kind == "native_layer_norm": - context.add((out, None, None), torch_name=node.name) - else: - context.add(out) - - _TORCH_OPS_REGISTRY.set_func_by_name(layer_norm, "layer_norm") - _TORCH_OPS_REGISTRY.set_func_by_name(layer_norm, "native_layer_norm") - - -def _patch_pipeline_remove_fuse_prelu() -> None: - return - - -_PASSES_TO_DROP = [ - "common::fuse_prelu", -] - - -def build_cactus_pass_pipeline(): - from coremltools.converters.mil.mil.passes.pass_pipeline import PassPipeline - pipeline = PassPipeline.DEFAULT - for pass_name in _PASSES_TO_DROP: - try: - pipeline.remove_passes([pass_name]) - except Exception: - pass - return pipeline - - -def _patch_mb_binops_scalar_cast() -> None: - import numpy as np - from coremltools.converters.mil import Builder as mb - - def _dtype_name(v): - d = getattr(v, "dtype", None) - return getattr(d, "__name__", str(d) if d else "") - - def _is_fp(name): - return "fp" in name or "float" in name or "double" in name - - def _np_for(dtype_name): - if "fp16" in dtype_name or "float16" in dtype_name: - return np.float16 - if "fp32" in dtype_name or "float32" in dtype_name: - return np.float32 - if "fp64" in dtype_name or "float64" in dtype_name or "double" in dtype_name: - return np.float64 - return np.float32 - - def _is_var(v): - return hasattr(v, "op") and hasattr(v, "sym_type") - - def _retype_in_place(v, target_dtype_name): - np_dtype = _np_for(target_dtype_name) - if isinstance(v, np.ndarray): - return v.astype(np_dtype) - if isinstance(v, (int, float)) and not isinstance(v, bool): - return np_dtype(v) - try: - return np.asarray(v, dtype=np_dtype) - except Exception: - return v - - def _scalar_const(s, dtype_name): - return mb.const(val=_np_for(dtype_name)(s)) - - for op_name in ("sub", "mul", "add", "div", "real_div"): - original = getattr(mb, op_name, None) - if original is None or getattr(original, "_cactus_scalar_cast_patched", False): - continue - - def make_wrapper(orig): - def wrapper(**kwargs): - x = kwargs.get("x") - y = kwargs.get("y") - xn = _dtype_name(x) - yn = _dtype_name(y) - if xn and not yn and _is_fp(xn) and isinstance(y, (int, float)) and not isinstance(y, bool): - kwargs["y"] = _scalar_const(y, xn) - elif yn and not xn and _is_fp(yn) and isinstance(x, (int, float)) and not isinstance(x, bool): - kwargs["x"] = _scalar_const(x, yn) - elif xn and yn and xn != yn: - target = "fp16" if (xn == "fp16" or yn == "fp16") else xn - if xn != target: - kwargs["x"] = mb.cast(x=x, dtype=target) if _is_var(x) else _retype_in_place(x, target) - if yn != target: - kwargs["y"] = mb.cast(x=y, dtype=target) if _is_var(y) else _retype_in_place(y, target) - return orig(**kwargs) - wrapper._cactus_scalar_cast_patched = True - wrapper.__name__ = orig.__name__ - return wrapper - - setattr(mb, op_name, make_wrapper(original)) - - select_op = getattr(mb, "select", None) - if select_op is not None and not getattr(select_op, "_cactus_scalar_cast_patched", False): - def _select_wrap(orig): - def wrapper(**kwargs): - a = kwargs.get("a") - b = kwargs.get("b") - an = _dtype_name(a) - bn = _dtype_name(b) - if an and not bn and _is_fp(an) and isinstance(b, (int, float)) and not isinstance(b, bool): - kwargs["b"] = _scalar_const(b, an) - elif bn and not an and _is_fp(bn) and isinstance(a, (int, float)) and not isinstance(a, bool): - kwargs["a"] = _scalar_const(a, bn) - elif an and bn and an != bn: - target = "fp16" if (an == "fp16" or bn == "fp16") else an - if an != target: - kwargs["a"] = mb.cast(x=a, dtype=target) if _is_var(a) else _retype_in_place(a, target) - if bn != target: - kwargs["b"] = mb.cast(x=b, dtype=target) if _is_var(b) else _retype_in_place(b, target) - return orig(**kwargs) - wrapper._cactus_scalar_cast_patched = True - wrapper.__name__ = orig.__name__ - return wrapper - setattr(mb, "select", _select_wrap(select_op)) - - orig_ln = getattr(mb, "layer_norm", None) - if orig_ln is not None and not getattr(orig_ln, "_cactus_scalar_cast_patched", False): - def _ln_wrap(orig): - def wrapper(**kwargs): - gamma = kwargs.get("gamma") - gn = _dtype_name(gamma) - if _is_fp(gn): - eps = kwargs.get("epsilon") - if isinstance(eps, (int, float)) and not isinstance(eps, bool): - kwargs["epsilon"] = _scalar_const(eps, gn) - elif hasattr(eps, "dtype") and _dtype_name(eps) != gn: - kwargs["epsilon"] = mb.cast(x=eps, dtype=gn) - return orig(**kwargs) - wrapper._cactus_scalar_cast_patched = True - wrapper.__name__ = orig.__name__ - return wrapper - setattr(mb, "layer_norm", _ln_wrap(orig_ln)) - - for op_name in ("batch_norm", "instance_norm", "rms_norm"): - orig_op = getattr(mb, op_name, None) - if orig_op is None or getattr(orig_op, "_cactus_scalar_cast_patched", False): - continue - def _bn_wrap(orig): - def wrapper(**kwargs): - ref = kwargs.get("gamma") or kwargs.get("x") - rn = _dtype_name(ref) - if _is_fp(rn): - eps = kwargs.get("epsilon") - if isinstance(eps, (int, float)) and not isinstance(eps, bool): - kwargs["epsilon"] = _scalar_const(eps, rn) - elif hasattr(eps, "dtype") and _dtype_name(eps) != rn: - kwargs["epsilon"] = mb.cast(x=eps, dtype=rn) - return orig(**kwargs) - wrapper._cactus_scalar_cast_patched = True - wrapper.__name__ = orig.__name__ - return wrapper - setattr(mb, op_name, _bn_wrap(orig_op)) - - -def _patch_fp16_cast_skip_layer_norm() -> None: - try: - from coremltools.converters.mil.mil.passes.defs.quantization import ( - FP16ComputePrecision, - ) - except Exception as exc: - print(f"npu.coremltools_patches: skip_layer_norm patch unavailable ({exc})") - return - base = set(FP16ComputePrecision._UNSUPPORTED_FP16_OPS) - for op in ("layer_norm", "batch_norm", "instance_norm", "rms_norm"): - base.add(op) - FP16ComputePrecision._UNSUPPORTED_FP16_OPS = base - - -def _override_one_hot_translator() -> None: - from coremltools.converters.mil.frontend.torch.torch_op_registry import ( - _TORCH_OPS_REGISTRY, - ) - from coremltools.converters.mil.frontend.torch.ops import _get_inputs, _get_kwinputs - from coremltools.converters.mil import Builder as mb - - def _dtype_name(v): - d = getattr(v, "dtype", None) - return getattr(d, "__name__", str(d) if d else "") - - def one_hot(context, node): - inputs = _get_inputs(context, node, expected=(1, 2)) - labels = inputs[0] - num_classes = inputs[1] if len(inputs) > 1 else -1 - num_classes = _get_kwinputs(context, node, "num_classes", default=[num_classes])[0] - if hasattr(num_classes, "val") and num_classes.val is not None: - num_classes = num_classes.val - - if hasattr(labels, "dtype") and _dtype_name(labels) != "int32": - labels = mb.cast(x=labels, dtype="int32") - - res = mb.one_hot(indices=labels, one_hot_vector_size=num_classes, name=node.name) - context.add(res) - - _TORCH_OPS_REGISTRY.set_func_by_name(one_hot, "one_hot") - - -def _register_noop_assert_ops() -> None: - from coremltools.converters.mil.frontend.torch.torch_op_registry import ( - _TORCH_OPS_REGISTRY, - ) - - def _noop(context, node): - return - - _TORCH_OPS_REGISTRY.set_func_by_name(_noop, "_assert_tensor_metadata") - - -def _register_unfold_op() -> None: - from coremltools.converters.mil.frontend.torch.torch_op_registry import ( - _TORCH_OPS_REGISTRY, - ) - from coremltools.converters.mil.frontend.torch.ops import _get_inputs - from coremltools.converters.mil import Builder as mb - - def unfold(context, node): - inputs = _get_inputs(context, node, min_expected=4) - x = inputs[0] - dimension = inputs[1].val if hasattr(inputs[1], "val") else int(inputs[1]) - size = inputs[2].val if hasattr(inputs[2], "val") else int(inputs[2]) - step = inputs[3].val if hasattr(inputs[3], "val") else int(inputs[3]) - - rank = len(x.shape) if hasattr(x, "shape") else x.rank - dim = int(dimension) - if dim < 0: - dim += rank - - windowed = mb.sliding_windows( - x=x, - axis=dim, - size=int(size), - stride=int(step), - ) - - out_rank = rank + 1 - perm = list(range(out_rank)) - size_axis = dim + 1 - perm.pop(size_axis) - perm.append(size_axis) - if perm == list(range(out_rank)): - out = mb.identity(x=windowed, name=node.name) - else: - out = mb.transpose(x=windowed, perm=perm, name=node.name) - context.add(out) - - _TORCH_OPS_REGISTRY.set_func_by_name(unfold, "unfold") diff --git a/python/cactus/transpile/npu/pipeline.py b/python/cactus/transpile/npu/pipeline.py deleted file mode 100644 index e15b25cce..000000000 --- a/python/cactus/transpile/npu/pipeline.py +++ /dev/null @@ -1,94 +0,0 @@ -from __future__ import annotations - -from contextlib import nullcontext -from pathlib import Path - -from .audio import emit_audio_encoder_mlpackage -from .source import emit_source_encoder_mlpackage -from .vision import emit_vision_encoder_mlpackage - - -_ENCODER_COMPONENTS = { - "audio_encoder": (emit_audio_encoder_mlpackage, "audio_encoder.mlpackage"), - "source_encoder": (emit_source_encoder_mlpackage, "source_encoder.mlpackage"), - "vision_encoder": (emit_vision_encoder_mlpackage, "vision_encoder.mlpackage"), -} - - -def run_encoder_pipeline( - component_specs, - artifact_dir: Path, - *, - enabled: bool = True, - quantize_bits: int | None = None, - audio_quantize_bits: int | None = None, - vision_quantize_bits: int | None = None, -) -> dict[str, str]: - results: dict[str, str] = {} - if not enabled or not component_specs: - return results - - bundle_root = artifact_dir / "components" - bundle_root.mkdir(parents=True, exist_ok=True) - - def _resolve(per_component: int | None, default: int | None) -> int | None: - chosen = per_component if per_component is not None else ( - quantize_bits if quantize_bits is not None else default - ) - return None if chosen == 0 else chosen - - component_quants = { - "audio_encoder": _resolve(audio_quantize_bits, default=8), - "source_encoder": _resolve(None, default=None), - "vision_encoder": _resolve(vision_quantize_bits, default=None), - } - - for spec in component_specs: - component = getattr(spec, "component", None) - if component not in _ENCODER_COMPONENTS: - continue - emit_fn, filename = _ENCODER_COMPONENTS[component] - example_inputs = tuple( - getattr(spec, "npu_example_inputs", None) or getattr(spec, "example_inputs", ()) or () - ) - if not example_inputs: - print(f"npu.pipeline: {component} spec has no example inputs; skipping") - continue - module = spec.npu_module or spec.module - n_runtime = min(max(1, spec.npu_runtime_input_count), len(example_inputs)) - qbits = component_quants[component] - qdesc = f"int{qbits}" if qbits else "fp16" - print(f"npu.pipeline: emitting {component} quant={qdesc}") - reparam = getattr(spec, "npu_reparam", None) - with (reparam(module) if reparam else nullcontext()): - if component == "source_encoder": - emitted = emit_fn( - module, - bundle_root, - example_inputs=example_inputs, - input_names=tuple(getattr(spec, "input_keys", ()) or ()), - filename=filename, - quantize_bits=qbits, - ) - elif component == "vision_encoder": - names = ("x",) + tuple(spec.input_keys[1:n_runtime]) - emitted = emit_fn( - module, - bundle_root, - runtime_inputs=tuple(zip(names, example_inputs[:n_runtime])), - baked_inputs=example_inputs[n_runtime:], - filename=filename, - quantize_bits=qbits, - ) - else: - emitted = emit_fn( - module, - bundle_root, - example_input=example_inputs[0], - baked_inputs=example_inputs[1:], - filename=filename, - quantize_bits=qbits, - ) - if emitted: - results[f"npu_{component}"] = f"components/{emitted}" - return results diff --git a/python/cactus/transpile/npu/source.py b/python/cactus/transpile/npu/source.py deleted file mode 100644 index 8dad63283..000000000 --- a/python/cactus/transpile/npu/source.py +++ /dev/null @@ -1,118 +0,0 @@ -from __future__ import annotations - -from pathlib import Path - -import numpy as np -import torch -import torch.nn.functional as F - -from .audio import _apply_weight_quantization, _import_coremltools - - -class SourceEncoderWrapper(torch.nn.Module): - def __init__( - self, - source_module: torch.nn.Module, - input_dtypes: tuple[torch.dtype, ...], - ): - super().__init__() - self.source = source_module - self.input_dtypes = input_dtypes - - def forward(self, *args: torch.Tensor) -> torch.Tensor: - typed_args = [] - for arg, dtype in zip(args, self.input_dtypes, strict=True): - typed_args.append(arg.to(dtype=dtype) if arg.dtype != dtype else arg) - out = self.source(*typed_args) - if isinstance(out, tuple): - return out[0] - return out - - -def emit_source_encoder_mlpackage( - source_module: torch.nn.Module, - bundle_dir: Path, - *, - example_inputs: tuple[torch.Tensor, ...], - input_names: tuple[str, ...], - filename: str = "source_encoder.mlpackage", - output_name: str = "encoder_hidden_states", - minimum_deployment_target: str = "iOS18", - quantize_bits: int | None = None, -) -> str | None: - ct = _import_coremltools() - if not example_inputs: - print("npu.source: source_encoder spec has no example inputs; skipping") - return None - - input_dtypes = tuple(t.dtype for t in example_inputs) - exported_inputs = tuple( - t.to(dtype=torch.int32) if t.dtype in {torch.int64, torch.long} else t - for t in example_inputs - ) - wrapper = SourceEncoderWrapper(source_module, input_dtypes) - wrapper.eval() - - original_sdpa = F.scaled_dot_product_attention - - def _sdpa_without_gqa(query, key, value, *args, **kwargs): - enable_gqa = bool(kwargs.pop("enable_gqa", False)) - if enable_gqa and query.dim() >= 3 and key.dim() >= 3 and query.shape[-3] != key.shape[-3]: - repeats = int(query.shape[-3]) // max(int(key.shape[-3]), 1) - key = key.repeat_interleave(repeats, dim=-3) - value = value.repeat_interleave(repeats, dim=-3) - return original_sdpa(query, key, value, *args, **kwargs) - - try: - F.scaled_dot_product_attention = _sdpa_without_gqa - with torch.no_grad(): - exported = torch.export.export(wrapper, exported_inputs) - exported = exported.run_decompositions({}) - except Exception as exc: - print(f"npu.source: torch.export failed ({type(exc).__name__}: {exc}); skipping mlpackage emit") - return None - finally: - F.scaled_dot_product_attention = original_sdpa - - del wrapper - import gc as _gc - _gc.collect() - - target_attr = getattr(ct.target, minimum_deployment_target, None) or ct.target.iOS17 - input_types = [] - for name, tensor in zip(input_names, exported_inputs, strict=True): - dtype = np.int32 if tensor.dtype in {torch.int32, torch.int64, torch.long} else None - input_types.append(ct.TensorType(name=name, shape=tuple(tensor.shape), dtype=dtype)) - - from .coremltools_patches import build_cactus_pass_pipeline - try: - mlmodel = ct.convert( - exported, - inputs=input_types, - outputs=[ct.TensorType(name=output_name)], - compute_precision=ct.precision.FLOAT16, - convert_to="mlprogram", - minimum_deployment_target=target_attr, - pass_pipeline=build_cactus_pass_pipeline(), - ) - except Exception as exc: - print(f"npu.source: coremltools.convert failed ({type(exc).__name__}: {exc})") - return None - - if quantize_bits is not None: - before_id = id(mlmodel) - mlmodel = _apply_weight_quantization(mlmodel, quantize_bits) - if id(mlmodel) != before_id: - print(f"npu.source: applied int{quantize_bits} weight quantization") - - bundle_dir.mkdir(parents=True, exist_ok=True) - out_path = bundle_dir / filename - try: - mlmodel.save(str(out_path)) - except Exception as exc: - print(f"npu.source: mlpackage save failed ({type(exc).__name__}: {exc})") - return None - - shapes = ",".join(f"{name}={tuple(t.shape)}" for name, t in zip(input_names, exported_inputs, strict=True)) - print(f"npu.source: wrote {out_path} ({shapes})") - return filename diff --git a/python/cactus/transpile/npu/vision.py b/python/cactus/transpile/npu/vision.py deleted file mode 100644 index 97087efdd..000000000 --- a/python/cactus/transpile/npu/vision.py +++ /dev/null @@ -1,126 +0,0 @@ -from __future__ import annotations - -from pathlib import Path -from typing import Any - -import torch - - -class VisionEncoderWrapper(torch.nn.Module): - def __init__( - self, - vision_module: torch.nn.Module, - baked_inputs: tuple[torch.Tensor, ...] = (), - ): - super().__init__() - self.vision = vision_module - self._n_baked = len(baked_inputs) - for idx, tensor in enumerate(baked_inputs): - self.register_buffer(f"_baked_{idx}", tensor, persistent=False) - - def forward(self, *runtime_inputs: torch.Tensor) -> torch.Tensor: - coerced = tuple( - t if torch.is_floating_point(t) else t.to(torch.long) - for t in runtime_inputs - ) - extra = tuple(getattr(self, f"_baked_{i}") for i in range(self._n_baked)) - return self.vision(*coerced, *extra) - - -def _import_coremltools() -> Any: - import coremltools as ct - from .coremltools_patches import apply_all_coremltools_patches - apply_all_coremltools_patches() - return ct - - -def _apply_weight_quantization(mlmodel: Any, bits: int) -> Any: - try: - from coremltools.optimize.coreml import ( - linear_quantize_weights, - OpLinearQuantizerConfig, - OptimizationConfig, - ) - op_config = OpLinearQuantizerConfig( - mode="linear_symmetric", - dtype=f"int{bits}", - granularity="per_channel", - ) - config = OptimizationConfig(global_config=op_config) - return linear_quantize_weights(mlmodel, config) - except Exception as exc: - print(f"npu.vision: weight quantization to int{bits} failed ({type(exc).__name__}: {exc}); keeping FP16 weights") - return mlmodel - - -def emit_vision_encoder_mlpackage( - vision_module: torch.nn.Module, - bundle_dir: Path, - *, - runtime_inputs: tuple[tuple[str, torch.Tensor], ...], - baked_inputs: tuple[torch.Tensor, ...] = (), - filename: str = "vision_encoder.mlpackage", - output_name: str = "encoded", - minimum_deployment_target: str = "iOS18", - quantize_bits: int | None = None, -) -> str | None: - ct = _import_coremltools() - - example_tensors = tuple(tensor for _, tensor in runtime_inputs) - wrapper = VisionEncoderWrapper(vision_module, baked_inputs=baked_inputs) - wrapper.eval() - - try: - with torch.no_grad(): - exported = torch.export.export(wrapper, example_tensors) - exported = exported.run_decompositions({}) - except Exception as exc: - print(f"npu.vision: torch.export failed ({type(exc).__name__}: {exc}); skipping mlpackage emit") - return None - - del wrapper - import gc as _gc - _gc.collect() - - target_attr = getattr(ct.target, minimum_deployment_target, None) or ct.target.iOS17 - - import numpy as np - ct_inputs = [ - ct.TensorType(name=name, shape=tuple(tensor.shape)) - if torch.is_floating_point(tensor) - else ct.TensorType(name=name, shape=tuple(tensor.shape), dtype=np.int32) - for name, tensor in runtime_inputs - ] - - from .coremltools_patches import build_cactus_pass_pipeline - try: - mlmodel = ct.convert( - exported, - inputs=ct_inputs, - outputs=[ct.TensorType(name=output_name)], - compute_precision=ct.precision.FLOAT16, - convert_to="mlprogram", - minimum_deployment_target=target_attr, - pass_pipeline=build_cactus_pass_pipeline(), - ) - except Exception as exc: - print(f"npu.vision: coremltools.convert failed ({type(exc).__name__}: {exc})") - return None - - if quantize_bits is not None: - before_id = id(mlmodel) - mlmodel = _apply_weight_quantization(mlmodel, quantize_bits) - if id(mlmodel) != before_id: - print(f"npu.vision: applied int{quantize_bits} weight quantization") - - bundle_dir.mkdir(parents=True, exist_ok=True) - out_path = bundle_dir / filename - try: - mlmodel.save(str(out_path)) - except Exception as exc: - print(f"npu.vision: mlpackage save failed ({type(exc).__name__}: {exc})") - return None - - shapes = ", ".join(f"{name}={tuple(t.shape)}" for name, t in runtime_inputs) - print(f"npu.vision: wrote {out_path} (inputs: {shapes})") - return filename diff --git a/python/cactus/transpile/ops.py b/python/cactus/transpile/ops.py deleted file mode 100644 index e1a870f8b..000000000 --- a/python/cactus/transpile/ops.py +++ /dev/null @@ -1,255 +0,0 @@ -from __future__ import annotations - -from dataclasses import dataclass - - -@dataclass(frozen=True) -class OpSchema: - name: str - num_inputs: int | None = None - attrs: tuple[str, ...] = () - backend_op: str | None = None - is_layout: bool = False - is_reduction: bool = False - is_fusible: bool = False - - -# Small canonical op registry for the transpiler. -# This is metadata for canonical IR ops, not another IR layer. - -OPS: tuple[OpSchema, ...] = ( - OpSchema("add", num_inputs=2, backend_op="add", is_fusible=True), - OpSchema("add_clipped", num_inputs=2, backend_op="add_clipped", is_fusible=True), - OpSchema("subtract", num_inputs=2, backend_op="subtract", is_fusible=True), - OpSchema("multiply", num_inputs=2, backend_op="multiply", is_fusible=True), - OpSchema("divide", num_inputs=2, backend_op="divide", is_fusible=True), - OpSchema("not_equal", num_inputs=2, backend_op="not_equal"), - OpSchema("equal", num_inputs=2), - OpSchema("greater", num_inputs=2), - OpSchema("greater_equal", num_inputs=2), - OpSchema("less", num_inputs=2), - OpSchema("less_equal", num_inputs=2), - OpSchema("logical_and", num_inputs=2), - OpSchema("logical_or", num_inputs=2), - OpSchema("logical_not", num_inputs=1), - OpSchema("abs", num_inputs=1, backend_op="abs", is_fusible=True), - OpSchema("clamp", num_inputs=1, attrs=("min", "max")), - OpSchema("pow", num_inputs=1, attrs=("exponent",), backend_op="pow"), - OpSchema("negate", num_inputs=1, is_fusible=True), - OpSchema("scalar_add", num_inputs=1, attrs=("value",), backend_op="scalar_add", is_fusible=True), - OpSchema("scalar_subtract", num_inputs=1, attrs=("value",), backend_op="scalar_subtract", is_fusible=True), - OpSchema("scalar_subtract_reverse", num_inputs=1, attrs=("value",)), - OpSchema("scalar_multiply", num_inputs=1, attrs=("value",), backend_op="scalar_multiply", is_fusible=True), - OpSchema("scalar_divide", num_inputs=1, attrs=("value",), backend_op="scalar_divide", is_fusible=True), - OpSchema("scalar_floor_divide", num_inputs=1, attrs=("value",), backend_op="scalar_floor_divide"), - OpSchema("scalar_divide_reverse", num_inputs=1, attrs=("value",)), - OpSchema("scalar_not_equal", num_inputs=1, attrs=("value",), backend_op="scalar_not_equal"), - OpSchema("scalar_equal", num_inputs=1, attrs=("value",)), - OpSchema("scalar_greater", num_inputs=1, attrs=("value",)), - OpSchema("scalar_greater_equal", num_inputs=1, attrs=("value",)), - OpSchema("scalar_less", num_inputs=1, attrs=("value",)), - OpSchema("scalar_less_equal", num_inputs=1, attrs=("value",)), - OpSchema("scalar_exp", num_inputs=1, backend_op="scalar_exp", is_fusible=True), - OpSchema("scalar_sqrt", num_inputs=1, backend_op="scalar_sqrt", is_fusible=True), - OpSchema("scalar_log", num_inputs=1, backend_op="scalar_log", is_fusible=True), - OpSchema("scalar_cos", num_inputs=1, backend_op="scalar_cos"), - OpSchema("scalar_sin", num_inputs=1, backend_op="scalar_sin"), - OpSchema("precision_cast", num_inputs=1, attrs=("dtype",), backend_op="precision_cast"), - OpSchema("view", num_inputs=1, attrs=("shape",), backend_op="reshape", is_layout=True), - OpSchema("permute", num_inputs=1, attrs=("permutation",), backend_op="permute", is_layout=True), - OpSchema("squeeze", num_inputs=1, attrs=("dim",), is_layout=True), - OpSchema("movedim", num_inputs=1, attrs=("source", "destination"), is_layout=True), - OpSchema("slice", num_inputs=1, attrs=("axis", "start", "end", "step"), backend_op="slice"), - OpSchema("split_with_sizes", num_inputs=1, attrs=("sizes", "axis")), - OpSchema("chunk", num_inputs=1, attrs=("chunks", "axis")), - OpSchema("ones", num_inputs=0, attrs=("shape", "dtype")), - OpSchema("one_hot", num_inputs=1, attrs=("num_classes",)), - OpSchema("tril", num_inputs=1, attrs=("diagonal",)), - OpSchema("unfold", num_inputs=1, attrs=("dimension", "size", "step")), - OpSchema("index", num_inputs=1, attrs=("axis", "index_value"), backend_op="index"), - OpSchema("advanced_index", num_inputs=None), - OpSchema( - "where", - num_inputs=None, - attrs=("true_is_scalar", "true_value", "false_is_scalar", "false_value"), - ), - OpSchema("masked_fill", num_inputs=2, attrs=("value",)), - OpSchema("masked_scatter", num_inputs=3), - OpSchema("getitem", num_inputs=1, attrs=("index",)), - OpSchema("cat", num_inputs=None, attrs=("axis",), backend_op="cat"), - OpSchema("matmul", num_inputs=2, backend_op="matmul"), - OpSchema("linear", num_inputs=None, attrs=("has_bias",), backend_op="matmul"), - OpSchema("addmm", num_inputs=3), - OpSchema("gather", num_inputs=2, attrs=("axis",), backend_op="gather"), - OpSchema("embedding", num_inputs=2, backend_op="embedding_from_tensor"), - OpSchema("sum", num_inputs=1, attrs=("axis", "keepdim"), backend_op="sum", is_reduction=True), - OpSchema("mean", num_inputs=1, attrs=("axis", "keepdim"), backend_op="mean", is_reduction=True), - OpSchema("variance", num_inputs=1, attrs=("axis", "keepdim"), backend_op="variance", is_reduction=True), - OpSchema("min", num_inputs=1, attrs=("axis", "keepdim"), backend_op="min", is_reduction=True), - OpSchema("max", num_inputs=1, attrs=("axis", "keepdim"), backend_op="max", is_reduction=True), - OpSchema("cumsum", num_inputs=1, attrs=("axis",), backend_op="cumsum"), - OpSchema("relu", num_inputs=1, backend_op="relu", is_fusible=True), - OpSchema("silu", num_inputs=1, backend_op="silu", is_fusible=True), - OpSchema("gelu", num_inputs=1, backend_op="gelu", is_fusible=True), - OpSchema("gelu_erf", num_inputs=1, backend_op="gelu_erf", is_fusible=True), - OpSchema("sigmoid", num_inputs=1, backend_op="sigmoid", is_fusible=True), - OpSchema("glu", num_inputs=1, attrs=("axis",), backend_op="glu"), - OpSchema("tanh", num_inputs=1, backend_op="tanh", is_fusible=True), - OpSchema("softplus", num_inputs=1), - OpSchema("softmax", num_inputs=1, attrs=("axis",), backend_op="softmax"), - OpSchema("layer_norm", num_inputs=None, attrs=("eps",), backend_op="layer_norm"), - OpSchema("group_norm", num_inputs=3, attrs=("num_groups", "eps"), backend_op="group_norm"), - OpSchema("batch_norm", num_inputs=5, attrs=("axis", "eps"), backend_op="batch_norm"), - OpSchema("conv1d", num_inputs=2, attrs=("stride", "padding", "dilation", "groups"), backend_op="conv1d"), - OpSchema("conv2d", num_inputs=2, attrs=("stride", "padding", "dilation", "groups"), backend_op="conv2d"), - OpSchema("pad", num_inputs=1, attrs=("pads", "value", "mode")), - OpSchema("lstm_cell", num_inputs=7, backend_op="lstm_cell"), - - # CUSTOM / FUSED OPS - OpSchema("dense_mlp_tq_fused", num_inputs=4, attrs=("product_scale",), backend_op="dense_mlp_tq_fused"), - OpSchema("rms_norm", num_inputs=2, attrs=("eps",), backend_op="rms_norm"), - OpSchema("rope", num_inputs=1, attrs=("theta", "position_offset"), backend_op="rope"), - OpSchema("rel_pos_bias", num_inputs=2, attrs=("scale",), backend_op="rel_pos_bias"), - OpSchema( - "attention", - num_inputs=3, - attrs=( - "scale", - "is_causal", - "window_size", - "mask", - "additive_mask", - "qkv_layout", - "q_layout", - "k_layout", - "v_layout", - "output_layout", - ), - backend_op="attention", - ), - OpSchema( - "attention_block", - num_inputs=None, - attrs=( - "scale", - "is_causal", - "window_size", - "has_mask", - "additive_mask", - "has_gate", - "has_bias", - "attention_output_shape", - "qkv_layout", - ), - is_fusible=True, - ), - OpSchema( - "self_attention_block", - num_inputs=None, - attrs=( - "scale", - "is_causal", - "window_size", - "has_mask", - "additive_mask", - "has_gate", - "has_bias", - "has_query_projection_bias", - "has_query_add", - "has_rel_query_add", - "has_key_projection_bias", - "has_value_projection_bias", - "has_rel_pos_bias", - "has_relative_key_projection_bias", - "query_shape", - "key_shape", - "value_shape", - "relative_key_shape", - "rel_pos_scale", - "attention_output_shape", - ), - is_fusible=True, - ), - OpSchema( - "conv_module", - num_inputs=None, - attrs=( - "eps", - "has_pointwise1_bias", - "has_depthwise_bias", - "has_pointwise2_bias", - "depthwise_kernel_size", - "depthwise_padding", - ), - is_fusible=True, - ), - OpSchema( - "gated_deltanet_prefill", - num_inputs=None, - attrs=( - "num_k_heads", - "num_v_heads", - "key_dim", - "value_dim", - "eps", - "chunk_size", - "has_z", - "has_dt_bias", - "has_a_log", - "has_conv", - ), - is_fusible=True, - ), - OpSchema( - "gated_deltanet_decode", - num_inputs=None, - attrs=( - "num_k_heads", - "num_v_heads", - "key_dim", - "value_dim", - "eps", - "has_z", - "has_dt_bias", - "has_a_log", - "has_conv", - ), - is_fusible=True, - ), -) - - -ALL_OPS: dict[str, OpSchema] = {op.name: op for op in OPS} - - -# Import/capture spellings reduced into the canonical registry above. -ALIASES: dict[str, str] = { - "reshape": "view", - "flatten": "view", - "unsqueeze": "view", - "expand": "view", - "contiguous": "view", - "identity": "view", - "transpose": "permute", - "concat": "cat", - "type_as": "precision_cast", - "scaled_dot_product_attention": "attention", - "cos": "scalar_cos", - "sin": "scalar_sin", -} - - -PRECISION_OP = "precision_cast" -SUPPORTED_PRECISIONS: tuple[str, ...] = ("int4", "int8", "fp16", "fp32") - - -def canonicalize_op(name: str) -> str: - return ALIASES.get(name, name) - - -def get_op(name: str) -> OpSchema: - return ALL_OPS[canonicalize_op(name)] - - -def has_op(name: str) -> bool: - return canonicalize_op(name) in ALL_OPS diff --git a/python/cactus/transpile/optimize_graph.py b/python/cactus/transpile/optimize_graph.py deleted file mode 100644 index ea2714d0d..000000000 --- a/python/cactus/transpile/optimize_graph.py +++ /dev/null @@ -1,1653 +0,0 @@ -from __future__ import annotations - -from dataclasses import dataclass -import re - -import torch - -from cactus.transpile.canonicalize.cleanup import canonicalize_exported_graph -from cactus.transpile.canonicalize.utils import normalize_dtype_name -from cactus.transpile.canonicalize.utils import remove_node -from cactus.transpile.canonicalize.utils import rebuild_graph -from cactus.transpile.fusion import match_attention -from cactus.transpile.fusion import match_attention_block -from cactus.transpile.fusion import match_conv_module -from cactus.transpile.fusion import match_gated_deltanet -from cactus.transpile.fusion import match_gated_mlp -from cactus.transpile.fusion import match_lstm_cell -from cactus.transpile.fusion import match_rel_pos_bias -from cactus.transpile.fusion import match_rms_norm -from cactus.transpile.fusion import match_rope -from cactus.transpile.fusion import match_self_attention_block -from cactus.transpile.fusion.common import producer -from cactus.transpile.fusion.common import strip_layout_passthrough -from cactus.transpile.fusion.common import strip_passthrough -from cactus.transpile.fusion.linear import match_linear -from cactus.transpile.fusion.rope import _extract_rope_angle_source -from cactus.transpile.graph_ir import IRGraph -from cactus.transpile.graph_ir import IRNode -from cactus.transpile.graph_ir import IRValue -from cactus.transpile.graph_ir import verify_ir -from cactus.transpile.model_patterns import GOLD_PATTERNS -from cactus.transpile.model_profiles import profile_for_family -from cactus.transpile.normalize import dtype_to_ir - - -@dataclass(frozen=True) -class DetectedPattern: - name: str - anchor_node_id: str - node_ids: tuple[str, ...] - value_ids: tuple[str, ...] - details: dict[str, object] - - -@dataclass(frozen=True) -class FusionConfig: - enable_gated_deltanet: bool = True - enable_lstm_cell: bool = True - enable_rms_norm: bool = True - enable_rope: bool = True - enable_rel_pos_bias: bool = True - enable_attention: bool = True - enable_self_attention_block: bool = True - enable_attention_block: bool = True - enable_conv_module: bool = True - enable_add_clipped: bool = True - enable_dense_mlp_tq_fused: bool = True - - -def optimize_graph(graph: IRGraph, *, max_passes: int = 8, config: FusionConfig | None = None, - precompute_rope: bool = True) -> IRGraph: - config = config or FusionConfig() - verify_ir(graph) - canonicalize_exported_graph(graph) - enable_attention_block_fusions = not ( - _is_gemma4_graph(graph) or _is_whisper_seq2seq_decoder_graph(graph) - ) - - for _ in range(max_passes): - changed = False - if config.enable_gated_deltanet and fuse_gated_deltanet(graph): - changed = True - if config.enable_lstm_cell and fuse_lstm_cells(graph): - changed = True - if config.enable_rms_norm: - if fuse_rms_norm(graph): - changed = True - if fuse_rms_norm_scale_multiply(graph): - changed = True - if config.enable_rope and fuse_rope(graph): - changed = True - if config.enable_rel_pos_bias and fuse_rel_pos_bias(graph): - changed = True - if config.enable_attention and fuse_attention(graph): - changed = True - if normalize_attention_layouts(graph): - changed = True - if ( - enable_attention_block_fusions - and config.enable_self_attention_block - and fuse_self_attention_blocks(graph) - ): - changed = True - if ( - enable_attention_block_fusions - and config.enable_attention_block - and fuse_attention_blocks(graph) - ): - changed = True - if normalize_gemma4_decoder_attention_semantics(graph): - changed = True - if normalize_cached_decoder_attention_hints(graph): - changed = True - if config.enable_dense_mlp_tq_fused and fuse_dense_mlp_tq(graph): - changed = True - if config.enable_conv_module and fuse_conv_modules(graph): - changed = True - if config.enable_add_clipped and fuse_add_clipped(graph): - changed = True - if not changed: - break - canonicalize_exported_graph(graph) - - if precompute_rope and precompute_rope_tables(graph): - canonicalize_exported_graph(graph) - - annotate_gold_patterns(graph) - _prune_unused_inputs(graph) - verify_ir(graph) - return graph - - -def fuse_gated_deltanet(graph: IRGraph) -> bool: - changed = False - for node_id in list(graph.order): - node = graph.nodes.get(node_id) - if node is None: - continue - - match = match_gated_deltanet(graph, node) - if match is None: - continue - - inputs = [ - match.normalized_input_value_id, - match.qkv_weight_value_id, - match.a_weight_value_id, - match.b_weight_value_id, - match.norm_weight_value_id, - ] - if match.z_weight_value_id is not None: - inputs.append(match.z_weight_value_id) - if match.dt_bias_value_id is not None: - inputs.append(match.dt_bias_value_id) - if match.a_log_value_id is not None: - inputs.append(match.a_log_value_id) - if match.conv_weight_value_id is not None: - inputs.append(match.conv_weight_value_id) - - node.op = f"gated_deltanet_{match.mode}" - node.inputs = inputs - node.attrs = { - "num_k_heads": int(match.num_k_heads), - "num_v_heads": int(match.num_v_heads), - "key_dim": int(match.key_dim), - "value_dim": int(match.value_dim), - "eps": float(match.eps), - "chunk_size": int(match.chunk_size), - "has_z": bool(match.z_weight_value_id is not None), - "has_dt_bias": bool(match.dt_bias_value_id is not None), - "has_a_log": bool(match.a_log_value_id is not None), - "has_conv": bool(match.conv_weight_value_id is not None), - } - node.kind = "semantic" - node.meta["gated_deltanet_mode"] = match.mode - changed = True - - if changed: - rebuild_graph(graph) - return changed - - -def fuse_lstm_cells(graph: IRGraph) -> bool: - changed = False - for node_id in list(graph.order): - node = graph.nodes.get(node_id) - if node is None or node.op != "multiply": - continue - - match = match_lstm_cell(graph, node) - if match is None: - continue - - bias_hh_value_id = match.bias_hh_value_id - if bias_hh_value_id is None: - bias_hh_value_id = _materialize_zero_like_constant( - graph, - match.bias_ih_value_id, - suffix="lstm_bias_hh_zero", - ) - - node.op = "lstm_cell" - node.inputs = [ - match.x_value_id, - match.h_prev_value_id, - match.c_prev_value_id, - match.weight_ih_value_id, - match.weight_hh_value_id, - match.bias_ih_value_id, - bias_hh_value_id, - ] - node.outputs = [match.h_output_value_id, match.c_output_value_id] - node.attrs = {} - node.kind = "semantic" - node.meta["lstm_cell_nodes"] = match.node_ids - node.meta["lstm_cell_bias_hh_zero"] = bool(match.bias_hh_value_id is None) - - for other_node_id in match.node_ids: - if other_node_id == node.id: - continue - graph.nodes.pop(other_node_id, None) - if other_node_id in graph.order: - graph.order.remove(other_node_id) - changed = True - - if changed: - rebuild_graph(graph) - return changed - - -def annotate_gold_patterns(graph: IRGraph) -> list[DetectedPattern]: - _clear_gold_pattern_annotations(graph) - patterns = [ - *_detect_gated_mlps(graph), - *_detect_decoder_attentions(graph), - *_detect_transformer_blocks(graph), - ] - - for pattern in patterns: - for node_id in pattern.node_ids: - node = graph.nodes.get(node_id) - if node is None: - continue - node.meta.setdefault("gold_patterns", []) - node.meta["gold_patterns"].append(pattern.name) - - anchor = graph.nodes.get(pattern.anchor_node_id) - if anchor is not None: - anchor.meta["gold_pattern_anchor"] = True - anchor.meta["gold_pattern_details"] = pattern.details - - graph.meta["gold_patterns_catalog"] = tuple(pattern.name for pattern in GOLD_PATTERNS) - graph.meta["detected_gold_patterns"] = patterns - return patterns - - -def summarize_detected_gold_patterns(graph: IRGraph) -> dict[str, int]: - summary: dict[str, int] = {} - for pattern in annotate_gold_patterns(graph): - summary[pattern.name] = summary.get(pattern.name, 0) + 1 - return summary - - -def fuse_rms_norm(graph: IRGraph) -> bool: - changed = False - for node_id in list(graph.order): - node = graph.nodes.get(node_id) - if node is None or node.op not in {"multiply", "type_as", "precision_cast"}: - continue - - match = match_rms_norm(graph, node) - if match is None: - continue - input_value = graph.values.get(match.input_value_id) - if input_value is not None and normalize_dtype_name(input_value.dtype) == "fp32": - # Keep FP32 normalization paths unfused so they can execute with FP32 - # reductions and elementwise math instead of the FP16-only RMS kernel. - continue - - weight_value_id = match.weight_value_id - if weight_value_id is None: - if input_value is None or input_value.shape is None or not input_value.shape: - continue - hidden_dim = input_value.shape[-1] - if not isinstance(hidden_dim, int) or hidden_dim <= 0: - continue - weight_value_id = _materialize_ones_constant( - graph, - hidden_dim, - dtype=input_value.dtype, - suffix="rms_norm_ones", - ) - node.op = "rms_norm" - node.inputs = [match.input_value_id, weight_value_id] - node.attrs = { - "eps": float(match.eps), - "weight_offset": float(match.weight_offset), - } - node.kind = "semantic" - node.meta["rms_weight_offset"] = float(match.weight_offset) - node.meta["rms_input_value_id"] = match.input_value_id - node.meta["rms_weight_value_id"] = weight_value_id - changed = True - - if changed: - rebuild_graph(graph) - return changed - - -def fuse_rms_norm_scale_multiply(graph: IRGraph) -> bool: - if not _is_gemma4_graph(graph): - return False - - changed = False - for node_id in list(graph.order): - node = graph.nodes.get(node_id) - if node is None or node.op != "multiply" or len(node.inputs) != 2 or len(node.outputs) != 1: - continue - - rms_node: IRNode | None = None - rms_output_id: str | None = None - scale_value_id: str | None = None - for input_id, other_id in ((node.inputs[0], node.inputs[1]), (node.inputs[1], node.inputs[0])): - candidate = producer(graph, input_id) - if candidate is None or candidate.op != "rms_norm" or len(candidate.inputs) != 2: - continue - if not _is_materialized_ones_constant(graph, candidate.inputs[1]): - continue - rms_node = candidate - rms_output_id = input_id - scale_value_id = other_id - break - - if rms_node is None or rms_output_id is None or scale_value_id is None: - continue - if rms_node.outputs != [rms_output_id]: - continue - if _has_other_users(graph, rms_output_id, excluding=node.id) or rms_output_id in graph.outputs: - continue - - node.op = "rms_norm" - node.inputs = [rms_node.inputs[0], scale_value_id] - node.attrs = dict(rms_node.attrs) - node.kind = "semantic" - node.meta["rms_folded_scale_multiply"] = True - node.meta["rms_input_value_id"] = rms_node.inputs[0] - node.meta["rms_weight_value_id"] = scale_value_id - remove_node(graph, rms_node.id) - graph.values.pop(rms_output_id, None) - changed = True - - if changed: - rebuild_graph(graph) - return changed - - -def fuse_rope(graph: IRGraph) -> bool: - changed = False - for node_id in list(graph.order): - node = graph.nodes.get(node_id) - if node is None or not node.outputs: - continue - - match = match_rope(graph, node.outputs[0]) - if match is None or match.partial: - continue - - node.op = "rope" - node.inputs = [match.input_value_id] - node.attrs = { - "theta": float(match.theta), - "position_offset": int(match.position_offset), - } - node.kind = "semantic" - changed = True - - if changed: - rebuild_graph(graph) - return changed - - -def precompute_rope_tables(graph: IRGraph) -> bool: - """Replace the runtime cos/sin(position * inv_freq) rope angle with an fp64-precomputed fp16 table - gathered by position id, fixing fp16's inability to represent angles past ~2048 in long context.""" - - changed = False - for node_id in list(graph.order): - node = graph.nodes.get(node_id) - if node is None or node.op not in {"scalar_cos", "scalar_sin"} or len(node.inputs) != 1: - continue - # Only the text decoder's rope. Vision/audio encoders carry cos/sin with - # the same shape but an axial/2D position this scalar table cannot model. - if node.meta.get("component") != "decoder": - continue - - output_id = node.outputs[0] - output_value = graph.values.get(output_id) - if output_value is None or output_value.shape is None: - continue - index_shape = tuple(int(dim) for dim in output_value.shape[:-1]) - - match = _match_rope_angle_table_source(graph, node.inputs[0], index_shape) - if match is None: - continue - inv_freq_value_id, position_value_id = match - - inv_freq = graph.constants[inv_freq_value_id].detach().cpu().to(torch.float64).reshape(-1) - head_dim = int(inv_freq.numel()) * 2 - max_seq = _rope_table_max_seq(graph) - positions = torch.arange(max_seq, dtype=torch.float64).reshape(max_seq, 1) - freqs = positions * inv_freq.reshape(1, -1) - emb = torch.cat((freqs, freqs), dim=-1) - table = (torch.cos(emb) if node.op == "scalar_cos" else torch.sin(emb)).to(torch.float16) - assert table.shape[-1] == head_dim - - table_value_id = _materialize_rope_table_constant(graph, node.id, table) - output_value.dtype = "fp16" - node.op = "embedding" - node.inputs = [table_value_id, position_value_id] - node.attrs = {} - node.kind = "generic" - changed = True - - if changed: - rebuild_graph(graph) - return changed - - -def _stored_inv_freq_constant(graph: IRGraph, value_id: str) -> str | None: - stripped = strip_layout_passthrough(graph, value_id) - constant = graph.constants.get(stripped) - if isinstance(constant, torch.Tensor): - return stripped - return None - - -def _match_rope_angle_table_source(graph: IRGraph, value_id: str, index_shape: tuple[int, ...]) -> tuple[str, str] | None: - cat_node = producer(graph, strip_passthrough(graph, value_id)) - if cat_node is None or cat_node.op != "cat" or len(cat_node.inputs) != 2: - return None - if strip_layout_passthrough(graph, cat_node.inputs[0]) != strip_layout_passthrough(graph, cat_node.inputs[1]): - return None - angle_info = _extract_rope_angle_source(graph, cat_node.inputs[0]) - if angle_info is None: - return None - angle_node = angle_info["matmul_node"] - if len(angle_node.inputs) != 2: - return None - - inv_freq_value_id: str | None = None - position_value_id: str | None = None - for input_id in angle_node.inputs: - const_id = _stored_inv_freq_constant(graph, input_id) - if const_id is not None: - inv_freq_value_id = const_id - continue - position_value_id = _trace_integer_position_source(graph, input_id, index_shape) - - if inv_freq_value_id is None or position_value_id is None: - return None - return inv_freq_value_id, position_value_id - - -def _trace_integer_position_source(graph: IRGraph, value_id: str, index_shape: tuple[int, ...]) -> str | None: - current = value_id - visited: set[str] = set() - while current not in visited: - visited.add(current) - value = graph.values.get(current) - if ( - value is not None - and str(value.dtype) in {"int32", "int64", "i32", "i64"} - and value.shape is not None - and tuple(int(dim) for dim in value.shape) == index_shape - ): - return current - node = producer(graph, current) - if node is None or node.op not in {"precision_cast", "view", "reshape", "unsqueeze", "squeeze", "contiguous"} or len(node.inputs) != 1: - return None - current = node.inputs[0] - return None - - -def _rope_table_max_seq(graph: IRGraph) -> int: - max_seq = _coerce_optional_int(graph.meta.get("max_cache_seq_len")) - if max_seq is None or max_seq <= 0: - max_seq = _coerce_optional_int(graph.meta.get("max_position_embeddings")) - if max_seq is None or max_seq <= 0: - raise ValueError("rope table precompute requires max_cache_seq_len or max_position_embeddings in graph.meta") - return int(max_seq) - - -def _materialize_rope_table_constant(graph: IRGraph, node_id: str, table: torch.Tensor) -> str: - value_id = f"c_rope_table_{node_id}" - graph.constants[value_id] = table - graph.values[value_id] = IRValue( - id=value_id, - shape=tuple(int(dim) for dim in table.shape), - dtype=dtype_to_ir(table.dtype), - producer=None, - users=[], - ) - return value_id - - -def fuse_attention(graph: IRGraph) -> bool: - changed = False - for node_id in list(graph.order): - node = graph.nodes.get(node_id) - if node is None or node.op not in {"scaled_dot_product_attention", "attention"}: - continue - - match = match_attention(graph, node) - if match is None: - continue - - window_size = int(node.attrs.get("window_size", 0)) - if len(node.inputs) > 3: - mask_info = _extract_sliding_window_mask(graph, node.inputs[3]) - if mask_info is not None: - node.meta["mask_window_size_hint"] = int(mask_info["window_size"]) - if window_size == 0: - window_size = int(mask_info["window_size"]) - node.meta["window_size_source"] = "mask_pattern" - - if window_size != 0 and "window_size_source" not in node.meta: - node.meta["window_size_source"] = "import_attr" - - semantic_attrs = { - key: value - for key, value in node.attrs.items() - if key not in {"mask", "dropout_p", "scale", "is_causal"} - } - semantic_attrs.update( - { - "scale": float(node.attrs.get("scale", 0.0)), - "is_causal": bool(node.attrs.get("is_causal", True)), - "window_size": window_size, - } - ) - - if ( - node.op == "attention" - and node.inputs[:3] == [match.query_value_id, match.key_value_id, match.value_value_id] - and node.attrs == semantic_attrs - ): - continue - - node.op = "attention" - node.inputs = [match.query_value_id, match.key_value_id, match.value_value_id] - node.attrs = semantic_attrs - node.kind = "semantic" - changed = True - - if changed: - rebuild_graph(graph) - return changed - - -def fuse_rel_pos_bias(graph: IRGraph) -> bool: - changed = False - for node_id in list(graph.order): - node = graph.nodes.get(node_id) - if node is None: - continue - - match = match_rel_pos_bias(graph, node) - if match is None: - continue - - node.op = "rel_pos_bias" - node.inputs = [match.query_value_id, match.relative_key_value_id] - node.attrs = {"scale": float(match.scale)} - node.kind = "semantic" - node.meta["rel_pos_bias_nodes"] = match.node_ids - changed = True - - if changed: - rebuild_graph(graph) - return changed - - -def fuse_attention_blocks(graph: IRGraph) -> bool: - changed = False - for node_id in list(graph.order): - node = graph.nodes.get(node_id) - if node is None: - continue - - match = match_attention_block(graph, node) - if match is None: - continue - - qkv_inputs = [match.query_value_id, match.key_value_id, match.value_value_id] - qkv_layout = "bhsd" - native_qkv_inputs = [ - _strip_bhsd_to_bthd_permute_input(graph, value_id) - for value_id in qkv_inputs - ] - if all(value_id is not None for value_id in native_qkv_inputs): - qkv_inputs = [str(value_id) for value_id in native_qkv_inputs] - qkv_layout = "bthd" - - inputs = qkv_inputs - if match.mask_value_id is not None: - inputs.append(match.mask_value_id) - if match.gate_value_id is not None: - inputs.append(match.gate_value_id) - inputs.append(match.output_projection_weight_value_id) - if match.output_projection_bias_value_id is not None: - inputs.append(match.output_projection_bias_value_id) - - node.op = "attention_block" - node.inputs = inputs - node.attrs = { - "scale": float(match.scale), - "is_causal": bool(match.is_causal), - "window_size": int(match.window_size), - "has_mask": bool(match.mask_value_id is not None), - "additive_mask": bool(match.additive_mask), - "has_gate": bool(match.gate_value_id is not None), - "has_bias": bool(match.output_projection_bias_value_id is not None), - "attention_output_shape": tuple(int(dim) for dim in match.attention_output_shape), - "qkv_layout": qkv_layout, - } - node.kind = "semantic" - node.meta["attention_block_source"] = match.attention_node_id - changed = True - - if changed: - rebuild_graph(graph) - return changed - - -def _strip_bhsd_to_bthd_permute_input(graph: IRGraph, value_id: str) -> str | None: - value = graph.values.get(value_id) - if value is None or value.producer is None: - return None - node = graph.nodes.get(value.producer) - if node is None or node.op != "permute" or len(node.inputs) != 1: - return None - permutation = tuple(int(dim) for dim in node.attrs.get("permutation", ())) - if permutation != (0, 2, 1, 3): - return None - source_id = strip_passthrough(graph, node.inputs[0]) - source_value = graph.values.get(source_id) - if source_value is None or source_value.shape is None or value.shape is None: - return None - source_shape = tuple(int(dim) for dim in source_value.shape) - output_shape = tuple(int(dim) for dim in value.shape) - if len(source_shape) != 4 or len(output_shape) != 4: - return None - if ( - output_shape[0] == source_shape[0] - and output_shape[1] == source_shape[2] - and output_shape[2] == source_shape[1] - and output_shape[3] == source_shape[3] - ): - return source_id - return None - - -def normalize_attention_layouts(graph: IRGraph) -> bool: - """Erase PyTorch BHSD<->BTHD layout round-trips around attention ops.""" - - changed = False - for node_id in list(graph.order): - node = graph.nodes.get(node_id) - if node is None or node.op not in {"attention", "scaled_dot_product_attention"}: - continue - if len(node.inputs) < 3 or len(node.outputs) != 1: - continue - - node_changed = False - layouts = ["bhsd", "bhsd", "bhsd"] - for index, value_id in enumerate(node.inputs[:3]): - native_value_id = _strip_bhsd_to_bthd_permute_input(graph, value_id) - if native_value_id is None: - continue - node.inputs[index] = str(native_value_id) - layouts[index] = "bthd" - node_changed = True - if node_changed: - node.attrs["q_layout"] = layouts[0] - node.attrs["k_layout"] = layouts[1] - node.attrs["v_layout"] = layouts[2] - changed = True - - output_id = node.outputs[0] - output_value = graph.values.get(output_id) - output_users = list(output_value.users) if output_value is not None else [] - if len(output_users) != 1 or output_id in graph.outputs: - continue - output_user = graph.nodes.get(output_users[0]) - if output_user is None or output_user.op != "permute" or len(output_user.inputs) != 1 or len(output_user.outputs) != 1: - continue - permutation = tuple(int(dim) for dim in output_user.attrs.get("permutation", ())) - if permutation != (0, 2, 1, 3): - continue - - replacement_id = output_user.outputs[0] - replacement_value = graph.values.get(replacement_id) - if replacement_value is not None: - old_value = graph.values.get(output_id) - if old_value is not None: - old_value.shape = replacement_value.shape - old_value.dtype = replacement_value.dtype - old_value.meta.update(replacement_value.meta) - node.outputs = [replacement_id] - node.attrs["output_layout"] = "bthd" - remove_node(graph, output_user.id) - changed = True - - if changed: - rebuild_graph(graph) - return changed - - -def normalize_gemma4_decoder_attention_semantics(graph: IRGraph) -> bool: - if not _is_gemma4_graph(graph): - return False - - changed = False - if _assign_gemma4_decoder_attention_hints_from_graph_meta(graph): - changed = True - sliding_window = _graph_sliding_window(graph) - for node_id in list(graph.order): - node = graph.nodes.get(node_id) - if node is None or node.op not in {"attention", "scaled_dot_product_attention", "attention_block"}: - continue - layer_type = str(node.meta.get("attention_layer_type") or "").strip().lower() - if layer_type == "full_attention": - mask_input_index: int | None = None - if node.op == "attention_block": - if bool(node.attrs.get("has_mask", False)) and len(node.inputs) > 3: - mask_input_index = 3 - elif len(node.inputs) > 3: - mask_input_index = 3 - - if mask_input_index is not None: - mask_value_id = node.inputs[mask_input_index] - if _gemma4_can_elide_attention_mask(graph, mask_value_id): - del node.inputs[mask_input_index] - node.meta["gemma4_full_mask_elided"] = True - changed = True - - if node.op == "attention_block" and bool(node.attrs.get("has_mask", False)): - node.attrs["has_mask"] = False - changed = True - if not bool(node.attrs.get("is_causal", False)): - node.attrs["is_causal"] = True - changed = True - if "additive_mask" in node.attrs and not bool(node.attrs.get("additive_mask", False)): - node.attrs.pop("additive_mask", None) - changed = True - if int(node.attrs.get("window_size", 0)) == 0 and not bool(graph.meta.get("use_internal_kv_cache", False)): - seq_len = 0 - attention_output_shape = node.attrs.get("attention_output_shape") - if isinstance(attention_output_shape, (list, tuple)) and len(attention_output_shape) >= 2: - try: - seq_len = int(attention_output_shape[1]) - except (TypeError, ValueError): - seq_len = 0 - if seq_len <= 0 and node.inputs: - query_value = graph.values.get(node.inputs[0]) - if query_value is not None and query_value.shape is not None and len(query_value.shape) >= 3: - seq_len = int(query_value.shape[2]) - if seq_len > 0: - # Gemma4 full-attention layers are semantically unwindowed, but using a - # window_size >= seq_len avoids the Apple Accelerate full-window fast path - # that has been unstable for transpiled grouped-query Gemma4 graphs. - node.attrs["window_size"] = seq_len - node.meta["gemma4_full_attention_window_compat"] = True - changed = True - continue - if layer_type != "sliding_attention": - continue - if bool(node.attrs.get("additive_mask", False)): - continue - - mask_input_index: int | None = None - if node.op == "attention_block": - if bool(node.attrs.get("has_mask", False)) and len(node.inputs) > 3: - mask_input_index = 3 - elif len(node.inputs) > 3: - mask_input_index = 3 - - if mask_input_index is not None: - mask_value_id = node.inputs[mask_input_index] - if _gemma4_can_elide_attention_mask(graph, mask_value_id): - mask_info = _extract_sliding_window_mask(graph, mask_value_id) - mask_input_name = _graph_input_name(graph, mask_value_id) - del node.inputs[mask_input_index] - node.meta["gemma4_sliding_mask_elided"] = True - if mask_info is not None: - node.meta["mask_window_size_hint"] = int(mask_info["window_size"]) - if int(node.attrs.get("window_size", 0)) == 0 and int(mask_info["window_size"]) > 0: - node.attrs["window_size"] = int(mask_info["window_size"]) - node.meta.setdefault("window_size_source", "mask_pattern") - elif mask_input_name is not None: - node.meta["gemma4_elided_mask_input_name"] = mask_input_name - changed = True - - if node.op == "attention_block" and bool(node.attrs.get("has_mask", False)): - node.attrs["has_mask"] = False - changed = True - if not bool(node.attrs.get("is_causal", False)): - node.attrs["is_causal"] = True - changed = True - if ( - int(node.attrs.get("window_size", 0)) == 0 - and sliding_window is not None - and sliding_window > 0 - ): - node.attrs["window_size"] = int(sliding_window) - node.meta.setdefault("window_size_source", "layer_type_config") - changed = True - if "additive_mask" in node.attrs and not bool(node.attrs.get("additive_mask", False)): - node.attrs.pop("additive_mask", None) - changed = True - - if changed: - rebuild_graph(graph) - return changed - - -def normalize_cached_decoder_attention_hints(graph: IRGraph) -> bool: - if _is_gemma4_graph(graph): - return False - if not bool(graph.meta.get("use_internal_kv_cache", False)): - return False - changed = _assign_gemma4_decoder_attention_hints_from_graph_meta(graph) - sliding_window = _graph_sliding_window(graph) - if sliding_window is not None and sliding_window > 0: - for node_id in graph.order: - node = graph.nodes.get(node_id) - if node is None or node.op not in {"attention", "scaled_dot_product_attention", "attention_block"}: - continue - layer_type = str(node.meta.get("attention_layer_type") or "").strip().lower() - if layer_type not in {"sliding", "sliding_attention"}: - continue - if int(node.attrs.get("window_size", 0) or 0) == 0: - node.attrs["window_size"] = int(sliding_window) - node.meta.setdefault("window_size_source", "layer_type_config") - changed = True - if changed: - rebuild_graph(graph) - return changed - - -def _assign_gemma4_decoder_attention_hints_from_graph_meta(graph: IRGraph) -> bool: - component = str(graph.meta.get("component", "") or "").strip().lower() - if component not in {"decoder", "decoder_step", "decoder_prefill_chunk", "decoder_media_step"}: - return False - - layer_types = _graph_layer_types(graph) - if not layer_types: - return False - - attention_nodes: list[IRNode] = [] - for node_id in graph.order: - node = graph.nodes.get(node_id) - if node is None or node.op not in {"attention", "scaled_dot_product_attention", "attention_block"}: - continue - attention_nodes.append(node) - - if len(attention_nodes) != len(layer_types): - return False - - use_internal_cache = bool(graph.meta.get("use_internal_kv_cache", False)) - sliding_window = _graph_sliding_window(graph) - graph_cache_seq_len = _graph_max_cache_seq_len(graph) - changed = False - for layer_index, (node, layer_type) in enumerate(zip(attention_nodes, layer_types, strict=True)): - if "attention_layer_type" not in node.meta: - node.meta["attention_layer_type"] = str(layer_type) - changed = True - if "attention_layer_index" not in node.meta: - node.meta["attention_layer_index"] = int(layer_index) - changed = True - if use_internal_cache: - normalized_layer_type = str(layer_type).strip().lower() - node_cache_seq_len = None - if normalized_layer_type in {"sliding", "sliding_attention"} and sliding_window is not None and sliding_window > 0: - node_cache_seq_len = int(sliding_window) - elif normalized_layer_type in {"global", "full", "full_attention"} and graph_cache_seq_len is not None: - node_cache_seq_len = int(graph_cache_seq_len) - if node_cache_seq_len is not None and node.meta.get("max_cache_seq_len") != node_cache_seq_len: - node.meta["max_cache_seq_len"] = node_cache_seq_len - changed = True - return changed - - -def _gemma4_can_elide_attention_mask( - graph: IRGraph, - value_id: str, - *, - _visited: set[str] | None = None, -) -> bool: - mask_value = graph.values.get(value_id) - if mask_value is None: - return False - if mask_value.dtype in (None, "bool"): - return True - - input_name = _graph_input_name(graph, value_id) - if isinstance(input_name, str) and "attention_mask" in input_name: - return True - - if _visited is None: - _visited = set() - if value_id in _visited: - return False - _visited.add(value_id) - - node = producer(graph, strip_passthrough(graph, value_id)) - if node is None: - return False - if node.op in { - "logical_and", - "logical_or", - "scalar_equal", - "scalar_not_equal", - "greater", - "greater_equal", - "less", - "less_equal", - "aten.__and__.Tensor", - }: - return True - if node.op in {"precision_cast", "view", "expand", "permute", "slice", "index", "where"}: - return any(_gemma4_can_elide_attention_mask(graph, input_id, _visited=_visited) for input_id in node.inputs) - return False - - -def fuse_self_attention_blocks(graph: IRGraph) -> bool: - changed = False - for node_id in list(graph.order): - node = graph.nodes.get(node_id) - if node is None: - continue - - match = match_self_attention_block(graph, node) - if match is None: - continue - - inputs = [match.hidden_value_id, match.query_weight_value_id] - if match.query_projection_bias_value_id is not None: - inputs.append(match.query_projection_bias_value_id) - if match.query_add_value_id is not None: - inputs.append(match.query_add_value_id) - if match.rel_query_add_value_id is not None: - inputs.append(match.rel_query_add_value_id) - inputs.append(match.key_weight_value_id) - if match.key_projection_bias_value_id is not None: - inputs.append(match.key_projection_bias_value_id) - inputs.append(match.value_weight_value_id) - if match.value_projection_bias_value_id is not None: - inputs.append(match.value_projection_bias_value_id) - if match.mask_value_id is not None: - inputs.append(match.mask_value_id) - if match.relative_key_input_value_id is not None and match.relative_key_weight_value_id is not None: - inputs.append(match.relative_key_input_value_id) - inputs.append(match.relative_key_weight_value_id) - if match.relative_key_projection_bias_value_id is not None: - inputs.append(match.relative_key_projection_bias_value_id) - if match.gate_value_id is not None: - inputs.append(match.gate_value_id) - inputs.append(match.output_projection_weight_value_id) - if match.output_projection_bias_value_id is not None: - inputs.append(match.output_projection_bias_value_id) - - node.op = "self_attention_block" - node.inputs = inputs - node.attrs = { - "scale": float(match.scale), - "is_causal": bool(match.is_causal), - "window_size": int(match.window_size), - "has_mask": bool(match.mask_value_id is not None), - "additive_mask": bool(match.additive_mask), - "has_gate": bool(match.gate_value_id is not None), - "has_bias": bool(match.output_projection_bias_value_id is not None), - "has_query_projection_bias": bool(match.query_projection_bias_value_id is not None), - "has_query_add": bool(match.query_add_value_id is not None), - "has_rel_query_add": bool(match.rel_query_add_value_id is not None), - "has_key_projection_bias": bool(match.key_projection_bias_value_id is not None), - "has_value_projection_bias": bool(match.value_projection_bias_value_id is not None), - "has_rel_pos_bias": bool(match.relative_key_input_value_id is not None and match.relative_key_weight_value_id is not None), - "has_relative_key_projection_bias": bool(match.relative_key_projection_bias_value_id is not None), - "query_shape": tuple(int(dim) for dim in match.query_shape), - "key_shape": tuple(int(dim) for dim in match.key_shape), - "value_shape": tuple(int(dim) for dim in match.value_shape), - "relative_key_shape": tuple(int(dim) for dim in match.relative_key_shape) if match.relative_key_shape is not None else (), - "rel_pos_scale": float(match.rel_pos_scale) if match.rel_pos_scale is not None else 1.0, - "attention_output_shape": tuple(int(dim) for dim in match.attention_output_shape), - } - node.kind = "semantic" - node.meta["self_attention_block_source"] = match.attention_node_id - changed = True - - if changed: - rebuild_graph(graph) - return changed - - -def fuse_add_clipped(graph: IRGraph) -> bool: - changed = False - for node_id in list(graph.order): - node = graph.nodes.get(node_id) - if node is None or node.op != "add" or len(node.inputs) != 2: - continue - - lhs = strip_passthrough(graph, node.inputs[0]) - rhs = strip_passthrough(graph, node.inputs[1]) - if not (_looks_like_gemma_residual_add(graph, lhs, rhs) or _looks_like_gemma_residual_add(graph, rhs, lhs)): - continue - - node.op = "add_clipped" - node.kind = "semantic" - changed = True - - if changed: - rebuild_graph(graph) - return changed - - -def fuse_dense_mlp_tq(graph: IRGraph) -> bool: - changed = False - for node_id in list(graph.order): - node = graph.nodes.get(node_id) - if node is None: - continue - - match = _match_dense_mlp_tq(graph, node) - if match is None: - continue - - matched_node_ids = set(match["node_ids"]) - if not _matched_nodes_are_private(graph, matched_node_ids, keep_node_id=node.id): - continue - - node.op = "dense_mlp_tq_fused" - node.inputs = [ - str(match["input_value_id"]), - str(match["gate_weight_value_id"]), - str(match["up_weight_value_id"]), - str(match["down_weight_value_id"]), - ] - node.attrs = {"product_scale": float(match.get("product_scale") or 1.0)} - node.kind = "semantic" - node.meta["dense_mlp_tq_fused"] = True - node.meta["dense_mlp_tq_nodes"] = tuple(sorted(matched_node_ids)) - if match.get("product_scale") is not None: - node.meta["product_scale_from_export"] = float(match["product_scale"]) - - for other_node_id in matched_node_ids - {node.id}: - graph.nodes.pop(other_node_id, None) - if other_node_id in graph.order: - graph.order.remove(other_node_id) - changed = True - - if changed: - rebuild_graph(graph) - return changed - - -def _match_dense_mlp_tq(graph: IRGraph, node: IRNode) -> dict[str, object] | None: - down = match_linear(graph, node) - if down is None or down.bias_value_id is not None: - return None - if not _is_cq_weight_value(graph, down.weight_value_id): - return None - - mul_node = producer(graph, down.input_value_id) - if mul_node is None or mul_node.op != "multiply" or len(mul_node.inputs) != 2: - return None - - for activated_value_id, up_value_id in ((mul_node.inputs[0], mul_node.inputs[1]), (mul_node.inputs[1], mul_node.inputs[0])): - activation_node, product_scale, scale_node_id = _unwrap_gemma4_scaled_activation(graph, activated_value_id) - if activation_node is None or activation_node.op != "gelu" or len(activation_node.inputs) != 1: - continue - - gate_input_producer = producer(graph, strip_passthrough(graph, activation_node.inputs[0])) - up_producer = producer(graph, strip_passthrough(graph, up_value_id)) - if gate_input_producer is None or up_producer is None: - continue - - gate = match_linear(graph, gate_input_producer) - up = match_linear(graph, up_producer) - if gate is None or up is None: - continue - if gate.bias_value_id is not None or up.bias_value_id is not None: - continue - if strip_passthrough(graph, gate.input_value_id) != strip_passthrough(graph, up.input_value_id): - continue - if not _is_cq_weight_value(graph, gate.weight_value_id): - continue - if not _is_cq_weight_value(graph, up.weight_value_id): - continue - if not _dense_mlp_weight_shapes_match(graph, gate.weight_value_id, up.weight_value_id, down.weight_value_id): - continue - - node_ids = { - *gate.node_ids, - activation_node.id, - *up.node_ids, - mul_node.id, - *down.node_ids, - } - if scale_node_id is not None: - node_ids.add(scale_node_id) - - return { - "input_value_id": strip_passthrough(graph, gate.input_value_id), - "gate_weight_value_id": gate.weight_value_id, - "up_weight_value_id": up.weight_value_id, - "down_weight_value_id": down.weight_value_id, - "product_scale": product_scale, - "node_ids": tuple(sorted(node_ids)), - } - - return None - - -def _unwrap_gemma4_scaled_activation(graph: IRGraph, value_id: str) -> tuple[IRNode | None, float | None, str | None]: - node = producer(graph, strip_passthrough(graph, value_id)) - if node is None: - return None, None, None - if node.op == "scalar_multiply" and len(node.inputs) == 1: - scale = float(node.attrs.get("value", 1.0)) - inner = producer(graph, strip_passthrough(graph, node.inputs[0])) - return inner, scale, node.id - return node, None, None - - -def _is_cq_weight_value(graph: IRGraph, value_id: str) -> bool: - value = graph.values.get(value_id) - if value is None: - return False - path = value.meta.get("path") if isinstance(value.meta, dict) else None - if isinstance(path, str) and ".cq" in path: - return True - return False - - -def _dense_mlp_weight_shapes_match( - graph: IRGraph, - gate_weight_value_id: str, - up_weight_value_id: str, - down_weight_value_id: str, -) -> bool: - gate = graph.values.get(gate_weight_value_id) - up = graph.values.get(up_weight_value_id) - down = graph.values.get(down_weight_value_id) - if gate is None or up is None or down is None: - return False - if gate.shape is None or up.shape is None or down.shape is None: - return False - if len(gate.shape) != 2 or len(up.shape) != 2 or len(down.shape) != 2: - return False - return ( - int(gate.shape[0]) == int(up.shape[0]) - and int(gate.shape[1]) == int(up.shape[1]) - and int(down.shape[1]) == int(gate.shape[0]) - ) - - -def _matched_nodes_are_private(graph: IRGraph, matched_node_ids: set[str], *, keep_node_id: str) -> bool: - for node_id in matched_node_ids: - if node_id == keep_node_id: - continue - node = graph.nodes.get(node_id) - if node is None: - continue - for output_id in node.outputs: - value = graph.values.get(output_id) - if value is None: - continue - for user_id in value.users: - if user_id not in matched_node_ids: - return False - return True - - -def fuse_conv_modules(graph: IRGraph) -> bool: - changed = False - for node_id in list(graph.order): - node = graph.nodes.get(node_id) - if node is None: - continue - - match = match_conv_module(graph, node) - if match is None: - continue - - inputs = [ - match.input_value_id, - match.pointwise1_weight_value_id, - ] - if match.pointwise1_bias_value_id is not None: - inputs.append(match.pointwise1_bias_value_id) - inputs.append(match.depthwise_weight_value_id) - if match.depthwise_bias_value_id is not None: - inputs.append(match.depthwise_bias_value_id) - inputs.extend( - [ - match.batch_norm_weight_value_id, - match.batch_norm_bias_value_id, - match.batch_norm_running_mean_value_id, - match.batch_norm_running_var_value_id, - match.pointwise2_weight_value_id, - ] - ) - if match.pointwise2_bias_value_id is not None: - inputs.append(match.pointwise2_bias_value_id) - - node.op = "conv_module" - node.inputs = inputs - node.attrs = { - "eps": float(match.eps), - "has_pointwise1_bias": bool(match.pointwise1_bias_value_id is not None), - "has_depthwise_bias": bool(match.depthwise_bias_value_id is not None), - "has_pointwise2_bias": bool(match.pointwise2_bias_value_id is not None), - "depthwise_kernel_size": int(match.depthwise_kernel_size), - "depthwise_padding": int(match.depthwise_padding), - } - node.kind = "semantic" - node.meta["conv_module_nodes"] = match.node_ids - changed = True - - if changed: - rebuild_graph(graph) - return changed - - -def _materialize_shifted_constant(graph: IRGraph, value_id: str, offset: float, *, suffix: str) -> str: - base = graph.constants[value_id] - if not isinstance(base, torch.Tensor): - raise NotImplementedError(f"expected tensor constant for {value_id}, got {type(base).__name__}") - - new_value_id = f"{value_id}_{suffix}" - if new_value_id in graph.constants: - return new_value_id - - shifted = base.detach().cpu() + offset - graph.constants[new_value_id] = shifted - meta: dict[str, object] = {} - source_value = graph.values.get(value_id) - if _is_gemma4_graph(graph) and float(offset) == 1.0 and source_value is not None: - path = source_value.meta.get("path") - kind = source_value.meta.get("kind") - source_name = source_value.meta.get("source_name") - if isinstance(path, str) and isinstance(kind, str) and isinstance(source_name, str): - meta = dict(source_value.meta) - meta["materialized_from_value_id"] = value_id - meta["materialized_by_op"] = "scalar_add" - graph.values[new_value_id] = IRValue( - id=new_value_id, - shape=tuple(shifted.shape), - dtype=dtype_to_ir(shifted.dtype), - producer=None, - users=[], - meta=meta, - ) - return new_value_id - - -def _materialize_zero_like_constant(graph: IRGraph, value_id: str, *, suffix: str) -> str: - base = graph.constants[value_id] - new_value_id = f"{value_id}_{suffix}" - if new_value_id in graph.constants: - return new_value_id - - if isinstance(base, torch.Tensor): - zero = torch.zeros_like(base.detach().cpu()) - else: - value = graph.values.get(value_id) - if value is None or value.shape is None or value.dtype is None: - raise NotImplementedError( - f"expected tensor constant metadata for {value_id}, got {type(base).__name__}" - ) - dtype_map = { - "bool": torch.bool, - "u8": torch.uint8, - "i8": torch.int8, - "i16": torch.int16, - "i32": torch.int32, - "i64": torch.int64, - "fp16": torch.float16, - "bf16": torch.bfloat16, - "fp32": torch.float32, - "fp64": torch.float64, - } - torch_dtype = dtype_map.get(str(value.dtype)) - if torch_dtype is None: - raise NotImplementedError(f"unsupported IR dtype for zero-like constant: {value.dtype}") - zero = torch.zeros(tuple(int(dim) for dim in value.shape), dtype=torch_dtype) - - graph.constants[new_value_id] = zero - graph.values[new_value_id] = IRValue( - id=new_value_id, - shape=tuple(zero.shape), - dtype=dtype_to_ir(zero.dtype), - producer=None, - users=[], - ) - return new_value_id - - -def _materialize_ones_constant(graph: IRGraph, size: int, *, dtype: str | None, suffix: str) -> str: - dtype_map = { - "fp16": torch.float16, - "bf16": torch.bfloat16, - "fp32": torch.float32, - "fp64": torch.float64, - } - torch_dtype = dtype_map.get(dtype, torch.float32) - new_value_id = f"c_{suffix}_{size}_{dtype or 'fp32'}" - if new_value_id in graph.constants: - return new_value_id - - value = torch.ones((size,), dtype=torch_dtype) - graph.constants[new_value_id] = value - graph.values[new_value_id] = IRValue( - id=new_value_id, - shape=(size,), - dtype=dtype_to_ir(value.dtype), - producer=None, - users=[], - ) - return new_value_id - - -def _is_gemma4_graph(graph: IRGraph) -> bool: - family = str(graph.meta.get("adapter_family") or graph.meta.get("family") or "").lower() - profile = profile_for_family(family) - return profile is not None and profile.family == "gemma4" - - -def _is_whisper_seq2seq_decoder_graph(graph: IRGraph) -> bool: - family = str(graph.meta.get("adapter_family") or graph.meta.get("family") or "").lower() - profile = profile_for_family(family) - task = str(graph.meta.get("task") or "").lower() - component = str(graph.meta.get("component") or "").lower() - return profile is not None and profile.family == "whisper" and task == "seq2seq_transcription" and component == "decoder" - - -def _prune_unused_inputs(graph: IRGraph) -> bool: - kept_inputs: list[str] = [] - changed = False - for value_id in graph.inputs: - value = graph.values.get(value_id) - if value_id in graph.outputs: - kept_inputs.append(value_id) - continue - if value is not None and value.users: - kept_inputs.append(value_id) - continue - changed = True - - if not changed: - return False - - graph.inputs = kept_inputs - rebuild_graph(graph) - return True - - -def _graph_input_name(graph: IRGraph, value_id: str) -> str | None: - input_names = _coerce_string_tuple(graph.meta.get("input_names")) - match = re.match(r"^v_args_(\d+)$", value_id) - if match is not None: - input_index = int(match.group(1)) - if input_index < len(input_names): - return input_names[input_index] - try: - input_index = list(graph.inputs).index(value_id) - except ValueError: - return None - if input_index < len(input_names): - return input_names[input_index] - return None - - -def _is_materialized_ones_constant(graph: IRGraph, value_id: str) -> bool: - value = graph.constants.get(value_id) - if not isinstance(value, torch.Tensor): - return False - if value.numel() == 0: - return False - return bool(torch.all(value.detach().cpu() == 1).item()) - - -def _has_other_users(graph: IRGraph, value_id: str, *, excluding: str) -> bool: - for other_id, other in graph.nodes.items(): - if other_id == excluding: - continue - if value_id in other.inputs: - return True - return False - - -def _clear_gold_pattern_annotations(graph: IRGraph) -> None: - for node in graph.nodes.values(): - node.meta.pop("gold_patterns", None) - node.meta.pop("gold_pattern_anchor", None) - node.meta.pop("gold_pattern_details", None) - - -def _detect_gated_mlps(graph: IRGraph) -> list[DetectedPattern]: - patterns: list[DetectedPattern] = [] - for node_id in graph.order: - node = graph.nodes.get(node_id) - if node is None: - continue - - match = match_gated_mlp(graph, node) - if match is None: - continue - - pattern_name = "gated_mlp_gelu" if match.activation == "gelu" else "gated_mlp_silu" - patterns.append( - DetectedPattern( - name=pattern_name, - anchor_node_id=node.id, - node_ids=match.node_ids, - value_ids=(match.input_value_id,), - details={ - "activation": match.activation, - "input_value_id": match.input_value_id, - "gate_weight_value_id": match.gate_weight_value_id, - "up_weight_value_id": match.up_weight_value_id, - "down_weight_value_id": match.down_weight_value_id, - }, - ) - ) - return patterns - - -def _detect_decoder_attentions(graph: IRGraph) -> list[DetectedPattern]: - patterns: list[DetectedPattern] = [] - for node_id in graph.order: - node = graph.nodes.get(node_id) - if node is None or node.op not in {"scaled_dot_product_attention", "attention"}: - continue - - match = match_attention(graph, node) - if match is None: - continue - - q_input, k_input, v_input = match.source_input_value_ids - pattern_name = "decoder_attention_gqa" if match.has_gqa_repeat else "decoder_attention" - - node_ids = set(match.node_ids) - details = { - "input_value_ids": (q_input, k_input, v_input), - "q_linear_weight_value_id": match.weight_value_ids[0], - "k_linear_weight_value_id": match.weight_value_ids[1], - "v_linear_weight_value_id": match.weight_value_ids[2], - "has_rope": bool(match.has_rope), - "has_qk_rms_norm": bool(match.has_qk_norm), - "has_gqa_repeat": bool(match.has_gqa_repeat), - "is_causal": bool(match.is_causal), - "scale": float(match.scale), - "window_size_hint": int(match.window_size), - } - - if len(node.inputs) > 3: - details["mask_value_id"] = node.inputs[3] - mask_info = _extract_sliding_window_mask(graph, node.inputs[3]) - if mask_info is not None: - details["mask_pattern"] = "sliding_window_attention_mask" - details["window_size_hint"] = int(mask_info["window_size"]) - node_ids.update(mask_info["node_ids"]) - - patterns.append( - DetectedPattern( - name=pattern_name, - anchor_node_id=node.id, - node_ids=tuple(sorted(node_ids)), - value_ids=(q_input, k_input, v_input), - details=details, - ) - ) - return patterns - - -def _detect_transformer_blocks(graph: IRGraph) -> list[DetectedPattern]: - patterns: list[DetectedPattern] = [] - for node_id in graph.order: - node = graph.nodes.get(node_id) - if node is None or node.op not in {"add", "add_clipped"} or len(node.inputs) != 2: - continue - - rhs_pattern = _find_anchor_pattern(graph, node.inputs[1], {"gated_mlp_gelu", "gated_mlp_silu"}) - lhs_pattern = _find_anchor_pattern(graph, node.inputs[0], {"gated_mlp_gelu", "gated_mlp_silu"}) - mlp_pattern = rhs_pattern or lhs_pattern - if mlp_pattern is None: - continue - - residual_value_id = node.inputs[0] if rhs_pattern is not None else node.inputs[1] - attn_pattern = _find_anchor_pattern( - graph, - residual_value_id, - {"decoder_attention_gqa", "decoder_attention", "gemma4_partial_rope_attention"}, - ) - if attn_pattern is None: - continue - - node_ids = {node.id} - for candidate in graph.nodes.values(): - gold_patterns = candidate.meta.get("gold_patterns", []) - if attn_pattern in gold_patterns or mlp_pattern in gold_patterns: - node_ids.add(candidate.id) - - patterns.append( - DetectedPattern( - name="decoder_block_post_attn_norm" if node.op == "add_clipped" else "decoder_block_simple_residual", - anchor_node_id=node.id, - node_ids=tuple(sorted(node_ids)), - value_ids=(residual_value_id,), - details={ - "residual_value_id": residual_value_id, - "attention_pattern": attn_pattern, - "mlp_pattern": mlp_pattern, - "residual_op": node.op, - }, - ) - ) - return patterns - - -def _looks_like_gemma_residual_add(graph: IRGraph, residual_value_id: str, branch_value_id: str) -> bool: - branch_node = producer(graph, strip_passthrough(graph, branch_value_id)) - if branch_node is None or branch_node.op != "rms_norm": - return False - return strip_passthrough(graph, residual_value_id) != strip_passthrough(graph, branch_node.inputs[0]) - - -def _find_anchor_pattern(graph: IRGraph, value_id: str, names: set[str]) -> str | None: - current = value_id - visited: set[str] = set() - while current not in visited: - visited.add(current) - current = strip_passthrough(graph, current) - node = producer(graph, current) - if node is None: - return None - - for name in node.meta.get("gold_patterns", []): - if name in names and node.meta.get("gold_pattern_anchor"): - return name - - if len(node.inputs) != 1: - return None - current = node.inputs[0] - return None - - -def _extract_sliding_window_mask(graph: IRGraph, value_id: str) -> dict[str, object] | None: - node_ids: set[str] = set() - top_and = producer(graph, strip_passthrough(graph, value_id)) - if top_and is None or top_and.op not in {"aten.__and__.Tensor", "logical_and"}: - return None - node_ids.add(top_and.id) - - stack = list(top_and.inputs) - saw_diff = False - saw_cumsum = False - window_candidates: list[int] = [] - - while stack: - current = strip_passthrough(graph, stack.pop()) - node = producer(graph, current) - if node is None or node.id in node_ids: - continue - node_ids.add(node.id) - - if node.op in {"aten.diff.default", "diff"}: - saw_diff = True - elif node.op in {"aten.cumsum.default", "cumsum"}: - saw_cumsum = True - elif node.op == "scalar_subtract": - maybe_window = int(node.attrs.get("value", 0)) - if maybe_window > 0: - window_candidates.append(maybe_window) - - stack.extend(node.inputs) - - if not saw_diff or not saw_cumsum: - return None - - return { - "window_size": max((candidate for candidate in window_candidates if candidate > 1), default=0), - "node_ids": tuple(sorted(node_ids)), - } - - -def _graph_layer_types(graph: IRGraph) -> tuple[str, ...]: - layer_types = _coerce_string_tuple(graph.meta.get("layer_types")) - if layer_types: - return layer_types - - providers = graph.meta.get("transpile_metadata_providers") - if isinstance(providers, dict): - for provider_meta in providers.values(): - if not isinstance(provider_meta, dict): - continue - layer_types = _coerce_string_tuple(provider_meta.get("layer_types")) - if layer_types: - return layer_types - return () - - -def _graph_sliding_window(graph: IRGraph) -> int | None: - sliding_window = _coerce_optional_int(graph.meta.get("sliding_window")) - if sliding_window is not None: - return sliding_window - - providers = graph.meta.get("transpile_metadata_providers") - if isinstance(providers, dict): - for provider_meta in providers.values(): - if not isinstance(provider_meta, dict): - continue - sliding_window = _coerce_optional_int(provider_meta.get("sliding_window")) - if sliding_window is not None: - return sliding_window - return None - - -def _graph_max_cache_seq_len(graph: IRGraph) -> int | None: - max_cache_seq_len = _coerce_optional_int(graph.meta.get("max_cache_seq_len")) - if max_cache_seq_len is not None: - return max_cache_seq_len - - providers = graph.meta.get("transpile_metadata_providers") - if isinstance(providers, dict): - for provider_meta in providers.values(): - if not isinstance(provider_meta, dict): - continue - max_cache_seq_len = _coerce_optional_int(provider_meta.get("max_cache_seq_len")) - if max_cache_seq_len is not None: - return max_cache_seq_len - return None - - -def _coerce_string_tuple(value: object) -> tuple[str, ...]: - if not isinstance(value, (tuple, list)): - return () - result: list[str] = [] - for item in value: - if not isinstance(item, str): - return () - result.append(item) - return tuple(result) - - -def _coerce_optional_int(value: object) -> int | None: - if value is None or isinstance(value, bool): - return None - if isinstance(value, int): - return int(value) - return None diff --git a/python/cactus/transpile/runtime_compat.py b/python/cactus/transpile/runtime_compat.py deleted file mode 100644 index 028682a96..000000000 --- a/python/cactus/transpile/runtime_compat.py +++ /dev/null @@ -1,376 +0,0 @@ -from __future__ import annotations - -import ctypes -import importlib -import sys -from typing import Any - -import numpy as np - - -class _MissingFFIFunction: - _cactus_missing_symbol = True - - def __init__(self, name: str): - self.__name__ = name - self.argtypes = None - self.restype = ctypes.c_int - - def __call__(self, *args: Any, **kwargs: Any) -> int: - raise RuntimeError(f"Cactus runtime is missing required symbol: {self.__name__}") - - -_ORIG_CDLL_GETATTR = ctypes.CDLL.__getattr__ - - -def _patched_cdll_getattr(self: ctypes.CDLL, name: str): - try: - return _ORIG_CDLL_GETATTR(self, name) - except AttributeError: - if not name.startswith("cactus_"): - raise - missing = _MissingFFIFunction(name) - setattr(self, name, missing) - return missing - - -def _load_runtime_module(): - if "cactus.bindings.cactus" in sys.modules: - return sys.modules["cactus.bindings.cactus"] - - ctypes.CDLL.__getattr__ = _patched_cdll_getattr - try: - return importlib.import_module("cactus.bindings.cactus") - finally: - ctypes.CDLL.__getattr__ = _ORIG_CDLL_GETATTR - - -def _patch_graph_runtime(cactus_module) -> None: - Graph = cactus_module.Graph - if getattr(Graph, "_transpile_runtime_compat_patched", False): - return - - _lib = cactus_module._lib - _err = cactus_module._err - cactus_node_t = cactus_module.cactus_node_t - - def _has_symbol(name: str) -> bool: - symbol = getattr(_lib, name, None) - return symbol is not None and not getattr(symbol, "_cactus_missing_symbol", False) - - def _ensure_compare_tensor(self, tensor): - tensor = self._ensure_tensor(tensor) - if int(tensor.dtype) == int(Graph.FP16): - return tensor - return self.precision_cast(tensor, Graph.FP16) - - def _ensure_scalar_tensor(self, tensor): - tensor = self._ensure_tensor(tensor) - if int(tensor.dtype) == int(Graph.FP16): - return tensor - return self.precision_cast(tensor, Graph.FP16) - - def _ensure_fp16_activation(self, tensor): - tensor = self._ensure_tensor(tensor) - if int(tensor.dtype) == int(Graph.FP32): - return self.precision_cast(tensor, Graph.FP16) - return tensor - - def _approx_nonzero_mask(self, tensor): - tensor = _ensure_scalar_tensor(self, tensor) - # Gemma4 compare paths in v2 only need stable 0/1-style masks for - # discrete values such as token-type ids and prebuilt boolean masks. - magnitude = self.abs(tensor) - shifted = self.scalar_add(magnitude, -0.5) - sharpened = self.scalar_multiply(shifted, 16.0) - return self.sigmoid(sharpened) - - orig_conv2d = Graph.conv2d - orig_not_equal = Graph.not_equal - orig_scalar_not_equal = Graph.scalar_not_equal - orig_transpose = Graph.transpose - orig_permute = Graph.permute - orig_scalar_add = Graph.scalar_add - orig_scalar_subtract = Graph.scalar_subtract - orig_scalar_multiply = Graph.scalar_multiply - orig_scalar_divide = Graph.scalar_divide - orig_scalar_exp = Graph.scalar_exp - orig_scalar_sqrt = Graph.scalar_sqrt - orig_scalar_cos = Graph.scalar_cos - orig_scalar_sin = Graph.scalar_sin - orig_scalar_log = Graph.scalar_log - orig_clamp = getattr(Graph, "clamp", None) - orig_add = Graph.add - orig_subtract = Graph.subtract - orig_multiply = Graph.multiply - orig_divide = Graph.divide - orig_concat = Graph.concat - orig_cat = Graph.cat - orig_abs = Graph.abs - orig_pow = Graph.pow - orig_expand = Graph.expand - - def matmul(self, a, b, pretransposed_rhs=False, backend=Graph.CPU, output_dtype=None): - a = _ensure_fp16_activation(self, a) - b = _ensure_fp16_activation(self, b) - out = cactus_node_t() - rc = _lib.cactus_graph_matmul( - self.h, - cactus_node_t(a.id), - cactus_node_t(b.id), - ctypes.c_bool(bool(pretransposed_rhs)), - ctypes.c_int32(int(backend)), - ctypes.byref(out), - ) - if rc != 0: - raise RuntimeError(_err("graph_matmul failed")) - result = self._tensor_from_node(out.value) - if output_dtype is not None and int(output_dtype) != int(result.dtype): - result = self.precision_cast(result, int(output_dtype)) - return result - - def gather(self, tensor, indices, axis=0): - tensor = self._ensure_tensor(tensor) - indices = self._ensure_tensor(indices) - if int(axis) != 0: - raise NotImplementedError( - f"transpiler runtime compatibility only supports gather(axis=0), got axis={axis}" - ) - out = cactus_node_t() - rc = _lib.cactus_graph_gather( - self.h, - cactus_node_t(tensor.id), - cactus_node_t(indices.id), - ctypes.byref(out), - ) - if rc != 0: - raise RuntimeError("graph_gather failed") - return self._tensor_from_node(out.value) - - def conv2d(self, x, weight, bias=None, stride=1, padding=0, dilation=1, groups=1): - def _pair(value: Any) -> tuple[int, int]: - if isinstance(value, (tuple, list)): - if len(value) != 2: - raise ValueError(f"expected pair for conv2d parameter, got {value!r}") - return int(value[0]), int(value[1]) - return int(value), int(value) - - stride_hw = _pair(stride) - padding_hw = _pair(padding) - dilation_hw = _pair(dilation) - groups_int = int(groups) - kernel_hw = tuple(int(dim) for dim in self._ensure_tensor(weight).shape[-2:]) - - if _has_symbol("cactus_graph_conv2d"): - return orig_conv2d( - self, - x, - weight, - bias=bias, - stride=stride, - padding=padding, - dilation=dilation, - groups=groups, - ) - - if ( - kernel_hw == (3, 3) - and stride_hw == (2, 2) - and padding_hw == (1, 1) - and dilation_hw == (1, 1) - ): - if groups_int == 1: - return self.conv2d_k3s2p1(x, weight, bias=bias) - if len(weight.shape) >= 1 and groups_int == int(weight.shape[0]): - return self.conv2d_depthwise_k3s2p1(x, weight, bias=bias) - - if ( - kernel_hw == (3, 3) - and stride_hw == (1, 1) - and padding_hw == (1, 1) - and dilation_hw == (1, 1) - and groups_int == 1 - and _has_symbol("cactus_graph_conv2d_k3s1p1") - ): - return self.conv2d_k3s1p1(x, weight, bias=bias) - - if ( - kernel_hw == (1, 1) - and stride_hw == (1, 1) - and padding_hw == (0, 0) - and dilation_hw == (1, 1) - and groups_int == 1 - ): - return self.conv2d_pointwise_1x1(x, weight, bias=bias) - - raise NotImplementedError( - "v2 runtime does not expose generic conv2d for this configuration: " - f"kernel={kernel_hw} stride={stride_hw} padding={padding_hw} " - f"dilation={dilation_hw} groups={groups_int}" - ) - - def scalar_not_equal(self, x, value): - if _has_symbol("cactus_graph_scalar_not_equal"): - return orig_scalar_not_equal(self, x, value) - x = _ensure_compare_tensor(self, x) - delta = self.scalar_add(x, -float(value)) - return _approx_nonzero_mask(self, delta) - - def not_equal(self, a, b): - if _has_symbol("cactus_graph_not_equal"): - return orig_not_equal(self, a, b) - a = _ensure_compare_tensor(self, a) - b = _ensure_compare_tensor(self, b) - delta = self.subtract(a, b) - return _approx_nonzero_mask(self, delta) - - def transpose(self, x, backend=Graph.CPU): - return orig_transpose(self, _ensure_scalar_tensor(self, x), backend=backend) - - def permute(self, x, permutation, backend=Graph.CPU): - return orig_permute(self, _ensure_scalar_tensor(self, x), permutation, backend=backend) - - def scalar_add(self, x, value): - return orig_scalar_add(self, self._ensure_tensor(x), value) - - def scalar_subtract(self, x, value): - return orig_scalar_subtract(self, self._ensure_tensor(x), value) - - def scalar_multiply(self, x, value): - return orig_scalar_multiply(self, self._ensure_tensor(x), value) - - def scalar_divide(self, x, value): - return orig_scalar_divide(self, self._ensure_tensor(x), value) - - def scalar_exp(self, x): - return orig_scalar_exp(self, self._ensure_tensor(x)) - - def scalar_sqrt(self, x): - return orig_scalar_sqrt(self, self._ensure_tensor(x)) - - def scalar_cos(self, x): - return orig_scalar_cos(self, self._ensure_tensor(x)) - - def scalar_sin(self, x): - return orig_scalar_sin(self, self._ensure_tensor(x)) - - def scalar_log(self, x): - return orig_scalar_log(self, self._ensure_tensor(x)) - - def clamp(self, x, lo, hi): - if orig_clamp is None: - raise RuntimeError("Cactus runtime is missing required symbol: cactus_graph_clamp") - return orig_clamp(self, _ensure_scalar_tensor(self, x), lo, hi) - - def add(self, a, b): - return orig_add(self, _ensure_scalar_tensor(self, a), _ensure_scalar_tensor(self, b)) - - def subtract(self, a, b): - return orig_subtract(self, _ensure_scalar_tensor(self, a), _ensure_scalar_tensor(self, b)) - - def multiply(self, a, b): - return orig_multiply(self, _ensure_scalar_tensor(self, a), _ensure_scalar_tensor(self, b)) - - def divide(self, a, b): - return orig_divide(self, _ensure_scalar_tensor(self, a), _ensure_scalar_tensor(self, b)) - - def concat(self, a, b, axis=0): - return orig_concat(self, _ensure_scalar_tensor(self, a), _ensure_scalar_tensor(self, b), axis=axis) - - def cat(self, tensors, axis=0): - legalized = [_ensure_scalar_tensor(self, tensor) for tensor in tensors] - return orig_cat(self, legalized, axis=axis) - - def abs(self, x): - return orig_abs(self, self._ensure_tensor(x)) - - def pow(self, x, exponent): - return orig_pow(self, self._ensure_tensor(x), exponent) - - def expand(self, x, shape): - x = self._ensure_tensor(x) - shape = tuple(int(v) for v in shape) - if tuple(int(dim) for dim in x.shape) == shape: - return x - if _has_symbol("cactus_graph_expand"): - return orig_expand(self, x, shape) - - # v2 builds do not always expose the generic expand symbol. For - # broadcast-only expands, adding an embedded zero tensor with the target - # shape gives the same logical result while staying inside existing - # Cactus graph ops. - dtype_map = { - int(Graph.FP16): np.float16, - int(Graph.FP32): np.float32, - } - x_dtype = int(x.dtype) - if x_dtype not in dtype_map: - raise NotImplementedError( - "expand fallback only supports floating-point tensors; " - f"got dtype={x_dtype} shape={tuple(int(dim) for dim in x.shape)} -> {shape}" - ) - zero = self.input(shape, dtype=x_dtype) - self.set_input(zero, np.zeros(shape, dtype=dtype_map[x_dtype]), dtype=x_dtype) - self.mark_embedded_input(zero) - return self.add(x, zero) - - Graph.matmul = matmul - Graph.gather = gather - Graph.conv2d = conv2d - Graph.scalar_not_equal = scalar_not_equal - Graph.not_equal = not_equal - Graph.transpose = transpose - Graph.permute = permute - Graph.scalar_add = scalar_add - Graph.scalar_subtract = scalar_subtract - Graph.scalar_multiply = scalar_multiply - Graph.scalar_divide = scalar_divide - Graph.scalar_exp = scalar_exp - Graph.scalar_sqrt = scalar_sqrt - Graph.scalar_cos = scalar_cos - Graph.scalar_sin = scalar_sin - Graph.scalar_log = scalar_log - Graph.clamp = clamp - Graph.add = add - Graph.subtract = subtract - Graph.multiply = multiply - Graph.divide = divide - Graph.concat = concat - Graph.cat = cat - Graph.abs = abs - Graph.pow = pow - Graph.expand = expand - Graph._transpile_runtime_compat_patched = True - - -_cactus_module = _load_runtime_module() - -if hasattr(_cactus_module, "_lib"): - _lib_obj = _cactus_module._lib - if hasattr(_lib_obj, "cactus_graph_matmul"): - _lib_obj.cactus_graph_matmul.argtypes = [ - ctypes.c_void_p, - _cactus_module.cactus_node_t, - _cactus_module.cactus_node_t, - ctypes.c_bool, - ctypes.c_int32, - ctypes.POINTER(_cactus_module.cactus_node_t), - ] - _lib_obj.cactus_graph_matmul.restype = ctypes.c_int - if hasattr(_lib_obj, "cactus_graph_gather"): - _lib_obj.cactus_graph_gather.argtypes = [ - ctypes.c_void_p, - _cactus_module.cactus_node_t, - _cactus_module.cactus_node_t, - ctypes.POINTER(_cactus_module.cactus_node_t), - ] - _lib_obj.cactus_graph_gather.restype = ctypes.c_int - -_patch_graph_runtime(_cactus_module) - -_lib = _cactus_module._lib -_err = _cactus_module._err -cactus_node_t = _cactus_module.cactus_node_t -cactus_tensor_info_t = _cactus_module.cactus_tensor_info_t -Graph = _cactus_module.Graph -Tensor = _cactus_module.Tensor diff --git a/python/cactus/transpile/runtime_support.py b/python/cactus/transpile/runtime_support.py deleted file mode 100644 index 8b966f77d..000000000 --- a/python/cactus/transpile/runtime_support.py +++ /dev/null @@ -1,178 +0,0 @@ -from __future__ import annotations - -import builtins -import importlib.util -import sys -from dataclasses import dataclass -from functools import lru_cache -from pathlib import Path - -import torch - -from cactus.transpile.model_profiles import GEMMA4_PROFILE -from cactus.transpile.model_profiles import ModelProfile -from cactus.transpile.model_profiles import profile_for_model_type - - -_TORCHVISION_COMPAT_LIBRARIES: list[object] = [] - - -@dataclass -class PreparedInputs: - names: tuple[str, ...] - tensors: tuple[torch.Tensor, ...] - metadata: dict[str, object] - - -def _transformers_supports_model_module(module_name: str) -> bool: - try: - return importlib.util.find_spec(module_name) is not None - except Exception: - return False - - -def _candidate_external_site_packages() -> list[Path]: - candidates: list[Path] = [] - major_minor = f"python{sys.version_info.major}.{sys.version_info.minor}" - pyenv_versions = Path.home() / ".pyenv" / "versions" - if not pyenv_versions.exists(): - return candidates - for version_dir in sorted(pyenv_versions.iterdir(), reverse=True): - site_packages = version_dir / "lib" / major_minor / "site-packages" - if site_packages.exists(): - candidates.append(site_packages) - return candidates - - -def ensure_transformers_supports_profile(profile: ModelProfile | None) -> str | None: - target_module = None if profile is None else profile.transformer_module - if not target_module or _transformers_supports_model_module(target_module): - return None - - for site_packages in _candidate_external_site_packages(): - candidate = site_packages / Path(*target_module.split(".")) - if not candidate.with_suffix(".py").exists() and not candidate.is_dir(): - continue - sys.path.insert(0, str(site_packages)) - for module_name in list(sys.modules): - root_name = module_name.split(".", 1)[0] - if root_name in {"transformers", "huggingface_hub", "tokenizers"}: - del sys.modules[module_name] - if _transformers_supports_model_module(target_module): - return str(site_packages) - try: - sys.path.remove(str(site_packages)) - except ValueError: - pass - return None - - -def ensure_transformers_supports_model_type(model_type: str) -> str | None: - return ensure_transformers_supports_profile(profile_for_model_type(model_type)) - - -def ensure_transformers_supports_gemma4() -> str | None: - return ensure_transformers_supports_profile(GEMMA4_PROFILE) - - -def _patch_torchvision_missing_nms_op() -> str | None: - """Keep mismatched torch/torchvision installs importable for processors.""" - - try: - import torchvision # type: ignore # noqa: F401 - return None - except RuntimeError as exc: - if "operator torchvision::nms does not exist" not in str(exc): - return None - except Exception: - return None - - try: - import torchvision.extension as tv_extension # type: ignore - - if bool(getattr(tv_extension, "_HAS_OPS", False)): - return None - except Exception: - pass - - try: - library = torch.library.Library("torchvision", "DEF") - library.define("nms(Tensor dets, Tensor scores, float iou_threshold) -> Tensor") - _TORCHVISION_COMPAT_LIBRARIES.append(library) - return "defined missing torchvision::nms operator for torchvision import compatibility" - except Exception: - return None - - -def patch_transformers_torchvision_probe() -> str | None: - has_torchvision = importlib.util.find_spec("torchvision") is not None - has_lzma = importlib.util.find_spec("_lzma") is not None - - if not has_torchvision: - return None - - if has_lzma: - nms_patch_note = _patch_torchvision_missing_nms_op() - return nms_patch_note - - base_note: str | None = None - try: - import backports.lzma as backports_lzma # type: ignore - - sys.modules.setdefault("lzma", backports_lzma) - base_note = "using backports.lzma because this Python build is missing _lzma" - nms_patch_note = _patch_torchvision_missing_nms_op() - if nms_patch_note: - return f"{base_note}; {nms_patch_note}" - return base_note - except Exception: - pass - - class _InterpolationModeStub: - NEAREST = "nearest" - BILINEAR = "bilinear" - BICUBIC = "bicubic" - LANCZOS = "lanczos" - - class _TorchvisionFunctionalStub: - InterpolationMode = _InterpolationModeStub - - def __getattr__(self, name: str): - raise RuntimeError( - "torchvision functionality is unavailable because this Python build is missing _lzma; " - f"attempted to access torchvision.transforms.functional.{name}" - ) - - builtins.F = _TorchvisionFunctionalStub() - builtins.tvF = builtins.F - - import transformers.utils as tf_utils # type: ignore - import transformers.utils.import_utils as tf_import_utils # type: ignore - - @lru_cache - def _disabled() -> bool: - return False - - tf_import_utils.is_torchvision_available = _disabled - tf_import_utils.is_torchvision_v2_available = _disabled - tf_utils.is_torchvision_available = _disabled - tf_utils.is_torchvision_v2_available = _disabled - return "disabled torchvision import checks because this Python build is missing _lzma" - - -def patch_torch_flex_attention_compat() -> str | None: - try: - import torch.nn.attention.flex_attention as flex_attention # type: ignore - except Exception: - return None - - if hasattr(flex_attention, "AuxRequest"): - return None - - class _AuxRequest: - def __init__(self, **kwargs): - for key, value in kwargs.items(): - setattr(self, key, value) - - flex_attention.AuxRequest = _AuxRequest # type: ignore[attr-defined] - return "installed torch flex_attention AuxRequest compatibility stub" diff --git a/python/cactus/transpile/tdt_runtime.py b/python/cactus/transpile/tdt_runtime.py deleted file mode 100644 index 01275d7ec..000000000 --- a/python/cactus/transpile/tdt_runtime.py +++ /dev/null @@ -1,1028 +0,0 @@ -from __future__ import annotations - -from dataclasses import dataclass -from dataclasses import replace -import json -import os -from pathlib import Path -import re -from typing import Any - -import numpy as np -import torch -import torch.nn as nn -import torch.nn.functional as F - -from cactus.transpile.component_pipeline import ComponentModuleSpec -from cactus.transpile.audio_preprocess import prepare_native_parakeet_audio_features -from cactus.transpile.model_profiles import add_tensor_aliases -from cactus.transpile.model_profiles import PARAKEET_TDT_PROFILE - - -def _cfg_get(config: dict[str, Any], key: str, default: Any = None) -> Any: - value = config.get(key, default) - return default if value is None else value - - -def _load_tensor_state_dict(model_source: str) -> dict[str, torch.Tensor]: - root = Path(model_source) - safetensors_path = root / "model.safetensors" - if safetensors_path.exists(): - from safetensors.torch import load_file - - return _with_parakeet_tdt_aliases(dict(load_file(str(safetensors_path)))) - - bin_path = root / "pytorch_model.bin" - if bin_path.exists(): - loaded = torch.load(bin_path, map_location="cpu") - if isinstance(loaded, dict): - return _with_parakeet_tdt_aliases({ - str(key): value - for key, value in loaded.items() - if isinstance(value, torch.Tensor) - }) - raise RuntimeError(f"unsupported Parakeet TDT checkpoint format in {model_source}") - - -def _add_tdt_derived_aliases(state_dict: dict[str, torch.Tensor]) -> None: - def alias(target: str, source: str) -> None: - if target not in state_dict and source in state_dict: - state_dict[target] = state_dict[source] - - for index in range(4): - alias(f"decoder.prediction.dec_rnn.lstm.{index}.Wx", f"decoder.lstm.weight_ih_l{index}") - alias(f"decoder.prediction.dec_rnn.lstm.{index}.Wh", f"decoder.lstm.weight_hh_l{index}") - alias(f"decoder.prediction.dec_rnn.lstm.{index}.Wx", f"decoder.prediction.dec_rnn.lstm.weight_ih_l{index}") - alias(f"decoder.prediction.dec_rnn.lstm.{index}.Wh", f"decoder.prediction.dec_rnn.lstm.weight_hh_l{index}") - bias_key = f"decoder.prediction.dec_rnn.lstm.{index}.bias" - bias_ih = state_dict.get(f"decoder.lstm.bias_ih_l{index}") - bias_hh = state_dict.get(f"decoder.lstm.bias_hh_l{index}") - if bias_ih is None: - bias_ih = state_dict.get(f"decoder.prediction.dec_rnn.lstm.bias_ih_l{index}") - if bias_hh is None: - bias_hh = state_dict.get(f"decoder.prediction.dec_rnn.lstm.bias_hh_l{index}") - if bias_key not in state_dict and bias_ih is not None and bias_hh is not None: - state_dict[bias_key] = bias_ih + bias_hh - - for source_index, target_index in ((0, 0), (2, 2), (3, 3), (5, 5), (6, 6)): - alias( - f"encoder.pre_encode.conv.{target_index}.weight", - f"encoder.subsampling.layers.{source_index}.weight", - ) - alias( - f"encoder.pre_encode.conv.{target_index}.bias", - f"encoder.subsampling.layers.{source_index}.bias", - ) - alias("encoder.pre_encode.out.weight", "encoder.subsampling.linear.weight") - alias("encoder.pre_encode.out.bias", "encoder.subsampling.linear.bias") - - - -def _with_parakeet_tdt_aliases(state_dict: dict[str, torch.Tensor]) -> dict[str, torch.Tensor]: - """Apply profile-declared aliases needed by the local TDT fallback runtime.""" - - add_tensor_aliases( - state_dict, - PARAKEET_TDT_PROFILE, - derived_aliases=_add_tdt_derived_aliases, - ) - return state_dict - - -@dataclass -class ParakeetTDTConfig: - model_source: str - sample_rate: int - num_mel_bins: int - hidden_dim: int - num_layers: int - attention_heads: int - attention_head_dim: int - attention_scale: float - ff_intermediate_dim: int - conv_kernel_size: int - subsampling_factor: int - subsampling_conv_channels: int - predictor_hidden_dim: int - predictor_num_layers: int - joint_dim: int - num_tdt_durations: int - tdt_durations: tuple[int, ...] - blank_id: int - vocabulary: tuple[str, ...] - encoder_hidden_act: str - - -def prepare_parakeet_tdt_audio_features( - audio_file: str | Path, - *, - expected_frames: int | None, - expected_mels: int, - torch_dtype: torch.dtype, -) -> tuple[torch.Tensor, int]: - return prepare_native_parakeet_audio_features( - audio_file, - expected_frames=expected_frames, - expected_mels=expected_mels, - torch_dtype=torch_dtype, - ) - - -def greedy_decode_parakeet_tdt_token_ids( - *, - config: ParakeetTDTConfig, - encoder_hidden_states: np.ndarray, - initial_states: tuple[np.ndarray, ...], - step, -) -> list[int]: - hidden = np.ascontiguousarray(np.asarray(encoder_hidden_states)) - if hidden.ndim != 3 or hidden.shape[0] != 1: - raise ValueError( - f"expected encoder_hidden_states with shape [1, T, D], got {tuple(hidden.shape)}" - ) - predictor_layers = int(config.predictor_num_layers) - if len(initial_states) != predictor_layers * 2: - raise ValueError( - f"expected {predictor_layers * 2} initial state tensors, got {len(initial_states)}" - ) - - states = tuple(np.ascontiguousarray(np.asarray(state)) for state in initial_states) - durations = [int(value) for value in config.tdt_durations] - if not durations and int(config.num_tdt_durations) > 0: - durations = list(range(int(config.num_tdt_durations))) - duration_class_count = max(len(durations), int(config.num_tdt_durations)) - last_token = _initial_parakeet_tdt_blank_id(config) - emitted: list[int] = [] - time_index = 0 - - while time_index < int(hidden.shape[1]): - frame = np.ascontiguousarray(hidden[:, time_index, :]) - advanced = False - symbols_added = 0 - - while symbols_added < 10: - logits, next_states = step(frame, last_token, states) - logits_array = np.asarray(logits, dtype=np.float32) - total_classes = int(logits_array.shape[-1]) - active_duration_count = duration_class_count - if active_duration_count <= 0 or active_duration_count >= total_classes: - active_duration_count = max(1, min(total_classes - 1, int(config.num_tdt_durations) or len(durations) or 1)) - if len(durations) != active_duration_count: - durations = list(range(active_duration_count)) - token_class_count = total_classes - active_duration_count - blank_id = _effective_parakeet_tdt_blank_id( - config, - token_class_count=token_class_count, - ) - if last_token < 0 or last_token >= token_class_count: - last_token = blank_id - token_scores = logits_array[:, :token_class_count] - duration_scores = logits_array[:, token_class_count:] - next_token = int(np.argmax(token_scores[0])) - duration_index = int(np.argmax(duration_scores[0])) - skip = int(durations[min(duration_index, len(durations) - 1)]) if durations else 1 - - if next_token != blank_id: - emitted.append(next_token) - last_token = next_token - states = tuple(np.ascontiguousarray(np.asarray(state).copy()) for state in next_states) - - symbols_added += 1 - - if skip > 0: - time_index += skip - advanced = True - break - if next_token == blank_id: - time_index += 1 - advanced = True - break - - if not advanced: - time_index += 1 - - return emitted - - -def _initial_parakeet_tdt_blank_id(config: ParakeetTDTConfig) -> int: - blank_id = int(config.blank_id) - vocab_size = len(config.vocabulary) - if vocab_size > 0 and blank_id == vocab_size - 1: - return vocab_size - return blank_id - - -def _effective_parakeet_tdt_blank_id( - config: ParakeetTDTConfig, - *, - token_class_count: int, -) -> int: - blank_id = int(config.blank_id) - vocab_size = len(config.vocabulary) - if vocab_size > 0 and blank_id == vocab_size - 1 and token_class_count == vocab_size + 1: - return token_class_count - 1 - if 0 <= blank_id < token_class_count: - return blank_id - if vocab_size > 0 and token_class_count == vocab_size + 1: - return token_class_count - 1 - return max(0, token_class_count - 1) - - -def load_parakeet_tdt_config(model_source: str) -> ParakeetTDTConfig: - root_path = Path(model_source) - config_path = root_path / "config.json" - if not config_path.exists(): - raise FileNotFoundError(f"missing config.json for Parakeet TDT: {config_path}") - root = json.loads(config_path.read_text()) - encoder = root.get("encoder") or root.get("encoder_config") or {} - decoder = root.get("decoder") or {} - prediction = decoder.get("prediction") or decoder.get("prednet") or {} - joint = root.get("joint") or {} - jointnet = joint.get("jointnet") or {} - model_defaults = root.get("model_defaults") or {} - preprocessor = root.get("preprocessor") or {} - - hidden_dim = int(_cfg_get(root, "hidden_dim", _cfg_get(encoder, "d_model", _cfg_get(encoder, "hidden_size", 0)))) - attention_heads = int(_cfg_get(encoder, "n_heads", _cfg_get(encoder, "num_attention_heads", 0))) - attention_head_dim = hidden_dim // max(attention_heads, 1) - tdt_durations = tuple( - int(value) - for value in _cfg_get(root, "tdt_durations", _cfg_get(root, "durations", _cfg_get(model_defaults, "tdt_durations", (0, 1, 2, 3, 4)))) - ) - vocabulary = tuple(str(value) for value in _cfg_get(joint, "vocabulary", ())) - if not vocabulary: - vocabulary = _load_parakeet_vocabulary(root_path) - decoder_vocab_size = int(_cfg_get(root, "vocab_size", _cfg_get(decoder, "vocab_size", len(vocabulary)))) - decoding = root.get("decoding") or {} - blank_id_raw = _cfg_get( - root, - "tdt_blank_id", - _cfg_get(root, "blank_token_id", _cfg_get(decoding, "blank_id", _cfg_get(decoder, "blank_id", None))), - ) - if blank_id_raw is None: - blank_id = decoder_vocab_size - else: - try: - blank_id = int(blank_id_raw) - except (TypeError, ValueError): - blank_id = decoder_vocab_size - if blank_id < 0: - blank_id = decoder_vocab_size - - return ParakeetTDTConfig( - model_source=model_source, - sample_rate=int(_cfg_get(preprocessor, "sample_rate", 16000)), - num_mel_bins=int(_cfg_get(root, "num_mel_bins", _cfg_get(preprocessor, "features", _cfg_get(encoder, "feat_in", _cfg_get(encoder, "num_mel_bins", 128))))), - hidden_dim=hidden_dim, - num_layers=int(_cfg_get(root, "num_layers", _cfg_get(encoder, "n_layers", _cfg_get(encoder, "num_hidden_layers", 0)))), - attention_heads=attention_heads, - attention_head_dim=attention_head_dim, - attention_scale=float(attention_head_dim ** -0.5 if attention_head_dim > 0 else 1.0), - ff_intermediate_dim=int( - _cfg_get( - root, - "ffn_intermediate_dim", - _cfg_get( - encoder, - "ffn_hidden_size", - _cfg_get(encoder, "intermediate_size", round(hidden_dim * float(_cfg_get(encoder, "ff_expansion_factor", 4.0)))), - ), - ) - ), - conv_kernel_size=int(_cfg_get(root, "conv_kernel_size", _cfg_get(encoder, "conv_kernel_size", 9))), - subsampling_factor=int(_cfg_get(root, "subsampling_factor", _cfg_get(encoder, "subsampling_factor", 8))), - subsampling_conv_channels=int( - _cfg_get(root, "subsampling_conv_channels", _cfg_get(encoder, "subsampling_conv_channels", 256)) - ), - predictor_hidden_dim=int( - _cfg_get(root, "predictor_hidden_dim", _cfg_get(root, "decoder_hidden_size", _cfg_get(prediction, "pred_hidden", _cfg_get(model_defaults, "pred_hidden", 640)))) - ), - predictor_num_layers=int( - _cfg_get(root, "predictor_num_layers", _cfg_get(root, "num_decoder_layers", _cfg_get(prediction, "pred_rnn_layers", 1))) - ), - joint_dim=int(_cfg_get(root, "tdt_joint_dim", _cfg_get(jointnet, "joint_hidden", _cfg_get(model_defaults, "joint_hidden", 640)))), - num_tdt_durations=int(_cfg_get(root, "tdt_num_durations", _cfg_get(model_defaults, "num_tdt_durations", len(tdt_durations)))), - tdt_durations=tdt_durations, - blank_id=blank_id, - vocabulary=vocabulary, - encoder_hidden_act=str(_cfg_get(root, "encoder_hidden_act", _cfg_get(encoder, "activation", _cfg_get(encoder, "hidden_act", "silu")))).lower(), - ) - - -def _load_parakeet_vocabulary(root: Path) -> tuple[str, ...]: - vocab_txt = root / "vocab.txt" - if vocab_txt.exists(): - pieces: list[str] = [] - for line in vocab_txt.read_text(encoding="utf-8").splitlines(): - if "\t" in line: - _, token = line.split("\t", 1) - pieces.append(token) - elif line: - pieces.append(line) - if pieces: - return tuple(pieces) - - tokenizer_json = root / "tokenizer.json" - if tokenizer_json.exists(): - try: - loaded = json.loads(tokenizer_json.read_text(encoding="utf-8")) - except Exception: - return () - model = loaded.get("model") if isinstance(loaded, dict) else None - vocab = model.get("vocab") if isinstance(model, dict) else None - if isinstance(vocab, dict): - ordered = sorted( - ((int(index), str(token)) for token, index in vocab.items()), - key=lambda item: item[0], - ) - return tuple(token for _, token in ordered) - return () - - -def _copy_linear_weight(linear: nn.Linear, weight: torch.Tensor, *, bias: torch.Tensor | None = None) -> None: - linear.weight.data.copy_(weight.to(dtype=linear.weight.dtype, device=linear.weight.device)) - if linear.bias is not None: - if bias is None: - linear.bias.data.zero_() - else: - linear.bias.data.copy_(bias.to(dtype=linear.bias.dtype, device=linear.bias.device)) - - -def _copy_conv2d_weight(conv: nn.Conv2d, weight: torch.Tensor, *, bias: torch.Tensor | None = None) -> None: - tensor = weight - if tuple(tensor.shape) != tuple(conv.weight.shape): - tensor = tensor.permute(0, 3, 1, 2).contiguous() - conv.weight.data.copy_(tensor.to(dtype=conv.weight.dtype, device=conv.weight.device)) - if conv.bias is not None: - if bias is None: - conv.bias.data.zero_() - else: - conv.bias.data.copy_(bias.to(dtype=conv.bias.dtype, device=conv.bias.device)) - - -def _copy_conv1d_weight(conv: nn.Conv1d, weight: torch.Tensor, *, bias: torch.Tensor | None = None) -> None: - tensor = weight - if tensor.ndim == 3 and tensor.shape[1] != conv.in_channels // conv.groups: - tensor = tensor.permute(0, 2, 1).contiguous() - conv.weight.data.copy_(tensor.to(dtype=conv.weight.dtype, device=conv.weight.device)) - if conv.bias is not None: - if bias is None: - conv.bias.data.zero_() - else: - conv.bias.data.copy_(bias.to(dtype=conv.bias.dtype, device=conv.bias.device)) - - -def _apply_activation(x: torch.Tensor, activation: str) -> torch.Tensor: - if "gelu" in activation: - return F.gelu(x) - if activation == "relu": - return F.relu(x) - return F.silu(x) - - -def _relative_position_embeddings(*, seq_len: int, hidden_dim: int, device: torch.device, dtype: torch.dtype) -> torch.Tensor: - rel_positions = torch.arange(seq_len - 1, -seq_len, -1, device=device, dtype=torch.float32) - half_dim = hidden_dim // 2 - inv_freq = 1.0 / ( - 10000.0 - ** (torch.arange(0, half_dim, device=device, dtype=torch.float32) * 2.0 / float(max(hidden_dim, 1))) - ) - angles = rel_positions.unsqueeze(1) * inv_freq.unsqueeze(0) - sin = torch.sin(angles) - cos = torch.cos(angles) - embeddings = torch.stack((sin, cos), dim=-1).reshape(2 * seq_len - 1, half_dim * 2) - if embeddings.shape[1] < hidden_dim: - pad = torch.zeros( - (embeddings.shape[0], hidden_dim - embeddings.shape[1]), - device=device, - dtype=embeddings.dtype, - ) - embeddings = torch.cat((embeddings, pad), dim=1) - elif embeddings.shape[1] > hidden_dim: - embeddings = embeddings[:, :hidden_dim] - return embeddings.to(dtype=dtype) - - -def _relative_position_bias(query: torch.Tensor, relative_key: torch.Tensor, *, scale: float) -> torch.Tensor: - batch, heads, seq_len, _ = query.shape - scores = torch.matmul(query, relative_key.transpose(-1, -2)) - rel_index = ( - torch.arange(seq_len, device=query.device).view(seq_len, 1) - - torch.arange(seq_len, device=query.device).view(1, seq_len) - + (seq_len - 1) - ) - rel_index = rel_index.view(1, 1, seq_len, seq_len).expand(batch, heads, seq_len, seq_len) - gathered = scores.gather(-1, rel_index) - return gathered * float(scale) - - -def _relative_position_bias_static(query: torch.Tensor, relative_key: torch.Tensor, *, scale: float) -> torch.Tensor: - batch, heads, seq_len, _ = query.shape - pos_len = 2 * seq_len - 1 - scores = torch.matmul(query, relative_key.transpose(-1, -2)) - x = F.pad(scores, (1, 0)) - x = x.reshape(batch, heads, pos_len + 1, seq_len) - x = x[:, :, 1:, :] - x = x.reshape(batch, heads, seq_len, pos_len) - return x[:, :, :, :seq_len] * float(scale) - - -class ParakeetTDTFeedForward(nn.Module): - def __init__(self, config: ParakeetTDTConfig, prefix: str, state_dict: dict[str, torch.Tensor]): - super().__init__() - self.linear1 = nn.Linear(config.hidden_dim, config.ff_intermediate_dim, bias=f"{prefix}.linear1.bias" in state_dict) - self.linear2 = nn.Linear(config.ff_intermediate_dim, config.hidden_dim, bias=f"{prefix}.linear2.bias" in state_dict) - _copy_linear_weight( - self.linear1, - state_dict[f"{prefix}.linear1.weight"], - bias=state_dict.get(f"{prefix}.linear1.bias"), - ) - _copy_linear_weight( - self.linear2, - state_dict[f"{prefix}.linear2.weight"], - bias=state_dict.get(f"{prefix}.linear2.bias"), - ) - - def forward(self, x: torch.Tensor, *, activation: str) -> torch.Tensor: - return self.linear2(_apply_activation(self.linear1(x), activation)) - - -class ParakeetTDTSelfAttention(nn.Module): - def __init__(self, config: ParakeetTDTConfig, prefix: str, state_dict: dict[str, torch.Tensor]): - super().__init__() - self.linear_q = nn.Linear(config.hidden_dim, config.hidden_dim, bias=f"{prefix}.linear_q.bias" in state_dict) - self.linear_k = nn.Linear(config.hidden_dim, config.hidden_dim, bias=f"{prefix}.linear_k.bias" in state_dict) - self.linear_v = nn.Linear(config.hidden_dim, config.hidden_dim, bias=f"{prefix}.linear_v.bias" in state_dict) - self.linear_out = nn.Linear(config.hidden_dim, config.hidden_dim, bias=f"{prefix}.linear_out.bias" in state_dict) - self.linear_pos = nn.Linear(config.hidden_dim, config.hidden_dim, bias=False) - _copy_linear_weight(self.linear_q, state_dict[f"{prefix}.linear_q.weight"], bias=state_dict.get(f"{prefix}.linear_q.bias")) - _copy_linear_weight(self.linear_k, state_dict[f"{prefix}.linear_k.weight"], bias=state_dict.get(f"{prefix}.linear_k.bias")) - _copy_linear_weight(self.linear_v, state_dict[f"{prefix}.linear_v.weight"], bias=state_dict.get(f"{prefix}.linear_v.bias")) - _copy_linear_weight( - self.linear_out, - state_dict[f"{prefix}.linear_out.weight"], - bias=state_dict.get(f"{prefix}.linear_out.bias"), - ) - _copy_linear_weight(self.linear_pos, state_dict[f"{prefix}.linear_pos.weight"]) - self.pos_bias_u = nn.Parameter(state_dict[f"{prefix}.pos_bias_u"].clone()) - self.pos_bias_v = nn.Parameter(state_dict[f"{prefix}.pos_bias_v"].clone()) - self.config = config - - def forward(self, x: torch.Tensor) -> torch.Tensor: - batch, seq_len, _ = x.shape - num_heads = self.config.attention_heads - head_dim = self.config.attention_head_dim - - q = self.linear_q(x).view(batch, seq_len, num_heads, head_dim) - k = self.linear_k(x).view(batch, seq_len, num_heads, head_dim) - v = self.linear_v(x).view(batch, seq_len, num_heads, head_dim) - - rel_pos = _relative_position_embeddings( - seq_len=seq_len, - hidden_dim=self.config.hidden_dim, - device=x.device, - dtype=x.dtype, - ) - rel_k = self.linear_pos(rel_pos).view(1, 2 * seq_len - 1, num_heads, head_dim) - - q_u = q + self.pos_bias_u.view(1, 1, num_heads, head_dim).to(dtype=x.dtype, device=x.device) - q_v = q + self.pos_bias_v.view(1, 1, num_heads, head_dim).to(dtype=x.dtype, device=x.device) - - q_u_heads = q_u.permute(0, 2, 1, 3) - q_v_heads = q_v.permute(0, 2, 1, 3) - k_heads = k.permute(0, 2, 1, 3) - v_heads = v.permute(0, 2, 1, 3) - rel_k_heads = rel_k.permute(0, 2, 1, 3) - - rel_bias = _relative_position_bias( - q_v_heads, - rel_k_heads, - scale=self.config.attention_scale, - ) - attn = F.scaled_dot_product_attention( - q_u_heads, - k_heads, - v_heads, - attn_mask=rel_bias, - dropout_p=0.0, - is_causal=False, - ) - attn = attn.permute(0, 2, 1, 3).reshape(batch, seq_len, self.config.hidden_dim) - return self.linear_out(attn) - - -class ParakeetTDTConformerConv(nn.Module): - def __init__(self, config: ParakeetTDTConfig, prefix: str, state_dict: dict[str, torch.Tensor]): - super().__init__() - hidden_dim = config.hidden_dim - kernel = config.conv_kernel_size - self.pointwise_conv1 = nn.Conv1d( - hidden_dim, - hidden_dim * 2, - kernel_size=1, - bias=f"{prefix}.pointwise_conv1.bias" in state_dict, - ) - self.depthwise_conv = nn.Conv1d( - hidden_dim, - hidden_dim, - kernel_size=kernel, - padding=kernel // 2, - groups=hidden_dim, - bias=f"{prefix}.depthwise_conv.bias" in state_dict, - ) - self.batch_norm = nn.BatchNorm1d(hidden_dim, eps=1e-5, affine=True, track_running_stats=True) - self.pointwise_conv2 = nn.Conv1d( - hidden_dim, - hidden_dim, - kernel_size=1, - bias=f"{prefix}.pointwise_conv2.bias" in state_dict, - ) - _copy_conv1d_weight( - self.pointwise_conv1, - state_dict[f"{prefix}.pointwise_conv1.weight"], - bias=state_dict.get(f"{prefix}.pointwise_conv1.bias"), - ) - _copy_conv1d_weight( - self.depthwise_conv, - state_dict[f"{prefix}.depthwise_conv.weight"], - bias=state_dict.get(f"{prefix}.depthwise_conv.bias"), - ) - _copy_conv1d_weight( - self.pointwise_conv2, - state_dict[f"{prefix}.pointwise_conv2.weight"], - bias=state_dict.get(f"{prefix}.pointwise_conv2.bias"), - ) - self.batch_norm.weight.data.copy_(state_dict[f"{prefix}.batch_norm.weight"].to(dtype=self.batch_norm.weight.dtype)) - self.batch_norm.bias.data.copy_(state_dict[f"{prefix}.batch_norm.bias"].to(dtype=self.batch_norm.bias.dtype)) - self.batch_norm.running_mean.data.copy_( - state_dict[f"{prefix}.batch_norm.running_mean"].to(dtype=self.batch_norm.running_mean.dtype) - ) - self.batch_norm.running_var.data.copy_( - state_dict[f"{prefix}.batch_norm.running_var"].to(dtype=self.batch_norm.running_var.dtype) - ) - self.config = config - - def forward(self, x: torch.Tensor, pad_mask: torch.Tensor | None = None) -> torch.Tensor: - x = x.transpose(1, 2) - x = self.pointwise_conv1(x) - x = F.glu(x, dim=1) - if pad_mask is not None: - x = x * pad_mask - x = self.depthwise_conv(x) - x = self.batch_norm(x) - x = _apply_activation(x, self.config.encoder_hidden_act) - x = self.pointwise_conv2(x) - return x.transpose(1, 2) - - -class ParakeetTDTEncoderLayer(nn.Module): - def __init__(self, config: ParakeetTDTConfig, layer_index: int, state_dict: dict[str, torch.Tensor]): - super().__init__() - prefix = f"encoder.layers.{layer_index}" - self.feed_forward1 = ParakeetTDTFeedForward(config, f"{prefix}.feed_forward1", state_dict) - self.feed_forward2 = ParakeetTDTFeedForward(config, f"{prefix}.feed_forward2", state_dict) - self.self_attn = ParakeetTDTSelfAttention(config, f"{prefix}.self_attn", state_dict) - self.conv = ParakeetTDTConformerConv(config, f"{prefix}.conv", state_dict) - self.norm_feed_forward1 = nn.LayerNorm(config.hidden_dim) - self.norm_self_att = nn.LayerNorm(config.hidden_dim) - self.norm_conv = nn.LayerNorm(config.hidden_dim) - self.norm_feed_forward2 = nn.LayerNorm(config.hidden_dim) - self.norm_out = nn.LayerNorm(config.hidden_dim) - for name in ("norm_feed_forward1", "norm_self_att", "norm_conv", "norm_feed_forward2", "norm_out"): - module = getattr(self, name) - module.weight.data.copy_(state_dict[f"{prefix}.{name}.weight"].to(dtype=module.weight.dtype)) - module.bias.data.copy_(state_dict[f"{prefix}.{name}.bias"].to(dtype=module.bias.dtype)) - self.config = config - - def forward(self, x: torch.Tensor) -> torch.Tensor: - x = x + 0.5 * self.feed_forward1(self.norm_feed_forward1(x), activation=self.config.encoder_hidden_act) - x = x + self.self_attn(self.norm_self_att(x)) - x = x + self.conv(self.norm_conv(x)) - x = x + 0.5 * self.feed_forward2(self.norm_feed_forward2(x), activation=self.config.encoder_hidden_act) - return self.norm_out(x) - - -class ParakeetTDTPreEncode(nn.Module): - def __init__(self, config: ParakeetTDTConfig, state_dict: dict[str, torch.Tensor]): - super().__init__() - channels = config.subsampling_conv_channels - self.conv = nn.ModuleList( - [ - nn.Conv2d(1, channels, kernel_size=3, stride=2, padding=1), - nn.Identity(), - nn.Conv2d(channels, channels, kernel_size=3, stride=2, padding=1, groups=channels), - nn.Conv2d(channels, channels, kernel_size=1, stride=1, padding=0), - nn.Identity(), - nn.Conv2d(channels, channels, kernel_size=3, stride=2, padding=1, groups=channels), - nn.Conv2d(channels, channels, kernel_size=1, stride=1, padding=0), - ] - ) - _copy_conv2d_weight(self.conv[0], state_dict["encoder.pre_encode.conv.0.weight"], bias=state_dict["encoder.pre_encode.conv.0.bias"]) - _copy_conv2d_weight(self.conv[2], state_dict["encoder.pre_encode.conv.2.weight"], bias=state_dict["encoder.pre_encode.conv.2.bias"]) - _copy_conv2d_weight(self.conv[3], state_dict["encoder.pre_encode.conv.3.weight"], bias=state_dict["encoder.pre_encode.conv.3.bias"]) - _copy_conv2d_weight(self.conv[5], state_dict["encoder.pre_encode.conv.5.weight"], bias=state_dict["encoder.pre_encode.conv.5.bias"]) - _copy_conv2d_weight(self.conv[6], state_dict["encoder.pre_encode.conv.6.weight"], bias=state_dict["encoder.pre_encode.conv.6.bias"]) - projected_width = config.num_mel_bins - for _ in range(3): - projected_width = (projected_width + 2 - 3) // 2 + 1 - self.out = nn.Linear(channels * projected_width, config.hidden_dim) - _copy_linear_weight( - self.out, - state_dict["encoder.pre_encode.out.weight"], - bias=state_dict["encoder.pre_encode.out.bias"], - ) - - def forward(self, input_features: torch.Tensor) -> torch.Tensor: - if input_features.ndim != 3: - raise ValueError(f"expected input_features [batch, frames, mels], got {tuple(input_features.shape)}") - x = input_features.unsqueeze(1) - x = F.relu(self.conv[0](x)) - x = F.relu(self.conv[3](self.conv[2](x))) - x = F.relu(self.conv[6](self.conv[5](x))) - x = x.permute(0, 2, 1, 3).reshape(x.shape[0], x.shape[2], -1) - return self.out(x) - - -class ParakeetTDTEncoder(nn.Module): - def __init__(self, config: ParakeetTDTConfig, state_dict: dict[str, torch.Tensor]): - super().__init__() - self.pre_encode = ParakeetTDTPreEncode(config, state_dict) - self.layers = nn.ModuleList( - [ParakeetTDTEncoderLayer(config, index, state_dict) for index in range(config.num_layers)] - ) - - def forward(self, input_features: torch.Tensor) -> torch.Tensor: - x = self.pre_encode(input_features) - for layer in self.layers: - x = layer(x) - return x - - -class ParakeetTDTANESelfAttention(nn.Module): - def __init__(self, attn: "ParakeetTDTSelfAttention", seq_len: int): - super().__init__() - self.linear_q = attn.linear_q - self.linear_k = attn.linear_k - self.linear_v = attn.linear_v - self.linear_out = attn.linear_out - self.linear_pos = attn.linear_pos - self.pos_bias_u = attn.pos_bias_u - self.pos_bias_v = attn.pos_bias_v - self.config = attn.config - weight = attn.linear_pos.weight - rel_pos = _relative_position_embeddings( - seq_len=int(seq_len), - hidden_dim=self.config.hidden_dim, - device=weight.device, - dtype=weight.dtype, - ) - with torch.no_grad(): - rel_k = attn.linear_pos(rel_pos).view( - 1, 2 * int(seq_len) - 1, self.config.attention_heads, self.config.attention_head_dim - ).flip(1) - self.register_buffer("_rel_k", rel_k.detach(), persistent=False) - - def forward(self, x: torch.Tensor, key_bias: torch.Tensor | None = None) -> torch.Tensor: - batch, seq_len, _ = x.shape - num_heads = self.config.attention_heads - head_dim = self.config.attention_head_dim - - q = self.linear_q(x).view(batch, seq_len, num_heads, head_dim) - k = self.linear_k(x).view(batch, seq_len, num_heads, head_dim) - v = self.linear_v(x).view(batch, seq_len, num_heads, head_dim) - - q_u = q + self.pos_bias_u.view(1, 1, num_heads, head_dim).to(dtype=x.dtype, device=x.device) - q_v = q + self.pos_bias_v.view(1, 1, num_heads, head_dim).to(dtype=x.dtype, device=x.device) - - rel_k_heads = self._rel_k.to(dtype=x.dtype, device=x.device).permute(0, 2, 1, 3) - rel_bias = _relative_position_bias_static( - q_v.permute(0, 2, 1, 3), - rel_k_heads, - scale=self.config.attention_scale, - ) - if key_bias is not None: - rel_bias = rel_bias + key_bias - attn = F.scaled_dot_product_attention( - q_u.permute(0, 2, 1, 3), - k.permute(0, 2, 1, 3), - v.permute(0, 2, 1, 3), - attn_mask=rel_bias, - dropout_p=0.0, - is_causal=False, - ) - attn = attn.permute(0, 2, 1, 3).reshape(batch, seq_len, self.config.hidden_dim) - return self.linear_out(attn) - - -class ParakeetTDTANEEncoderLayer(nn.Module): - def __init__(self, layer: "ParakeetTDTEncoderLayer", seq_len: int): - super().__init__() - self.feed_forward1 = layer.feed_forward1 - self.feed_forward2 = layer.feed_forward2 - self.self_attn = ParakeetTDTANESelfAttention(layer.self_attn, seq_len) - self.conv = layer.conv - self.norm_feed_forward1 = layer.norm_feed_forward1 - self.norm_self_att = layer.norm_self_att - self.norm_conv = layer.norm_conv - self.norm_feed_forward2 = layer.norm_feed_forward2 - self.norm_out = layer.norm_out - self.config = layer.config - - def forward(self, x: torch.Tensor, conv_mask: torch.Tensor, key_bias: torch.Tensor) -> torch.Tensor: - x = x + 0.5 * self.feed_forward1(self.norm_feed_forward1(x), activation=self.config.encoder_hidden_act) - x = x + self.self_attn(self.norm_self_att(x), key_bias) - x = x + self.conv(self.norm_conv(x), conv_mask) - x = x + 0.5 * self.feed_forward2(self.norm_feed_forward2(x), activation=self.config.encoder_hidden_act) - return self.norm_out(x) - - -class ParakeetTDTANEEncoder(nn.Module): - def __init__(self, encoder: "ParakeetTDTEncoder", seq_len: int): - super().__init__() - self.pre_encode = encoder.pre_encode - self.layers = nn.ModuleList( - [ParakeetTDTANEEncoderLayer(layer, seq_len) for layer in encoder.layers] - ) - factor = max(1, int(encoder.layers[0].config.subsampling_factor)) - self._subsample_stages = max(0, factor.bit_length() - 1) - - def _frame_masks(self, input_features: torch.Tensor, seq_len: int) -> tuple[torch.Tensor, torch.Tensor]: - valid = (input_features.abs().sum(-1) > 0).to(input_features.dtype).sum(-1, keepdim=True) - for _ in range(self._subsample_stages): - valid = torch.floor((valid - 1.0) / 2.0) + 1.0 - positions = torch.arange(seq_len, device=input_features.device, dtype=input_features.dtype) - enc_valid = (positions.view(1, 1, seq_len) < valid.view(-1, 1, 1)).to(input_features.dtype) - key_bias = (1.0 - enc_valid).unsqueeze(1) * -1e4 - return enc_valid, key_bias - - def forward(self, input_features: torch.Tensor) -> torch.Tensor: - x = self.pre_encode(input_features) - conv_mask, key_bias = self._frame_masks(input_features, x.shape[1]) - for layer in self.layers: - x = layer(x, conv_mask, key_bias) - return x - - -class ParakeetTDTDecoderCell(nn.Module): - def __init__(self, hidden_dim: int, prefix: str, state_dict: dict[str, torch.Tensor]): - super().__init__() - self.Wx = nn.Parameter(state_dict[f"{prefix}.Wx"].clone()) - self.Wh = nn.Parameter(state_dict[f"{prefix}.Wh"].clone()) - self.bias = nn.Parameter(state_dict[f"{prefix}.bias"].clone()) - self.hidden_dim = hidden_dim - - def forward(self, x: torch.Tensor, h: torch.Tensor, c: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: - gates = F.linear(x, self.Wx) + F.linear(h, self.Wh) + self.bias - i, f, g, o = torch.chunk(gates, 4, dim=-1) - i = torch.sigmoid(i) - f = torch.sigmoid(f) - g = torch.tanh(g) - o = torch.sigmoid(o) - c_next = f * c + i * g - h_next = o * torch.tanh(c_next) - return h_next, c_next - - -class ParakeetTDTDecoderPrediction(nn.Module): - def __init__(self, config: ParakeetTDTConfig, state_dict: dict[str, torch.Tensor]): - super().__init__() - vocab_size = max(config.blank_id + 1, state_dict["decoder.prediction.embed.weight"].shape[0]) - self.embed = nn.Embedding(vocab_size, config.predictor_hidden_dim) - self.embed.weight.data.copy_(state_dict["decoder.prediction.embed.weight"].to(dtype=self.embed.weight.dtype)) - self.dec_rnn = nn.Module() - self.dec_rnn.lstm = nn.ModuleList( - [ - ParakeetTDTDecoderCell( - config.predictor_hidden_dim, - f"decoder.prediction.dec_rnn.lstm.{index}", - state_dict, - ) - for index in range(config.predictor_num_layers) - ] - ) - - def forward( - self, - token_ids: torch.Tensor, - state_h: tuple[torch.Tensor, ...], - state_c: tuple[torch.Tensor, ...], - ) -> tuple[torch.Tensor, tuple[torch.Tensor, ...], tuple[torch.Tensor, ...]]: - x = self.embed(token_ids) - if x.ndim == 3 and x.shape[1] == 1: - x = x[:, 0, :] - next_h: list[torch.Tensor] = [] - next_c: list[torch.Tensor] = [] - for index, cell in enumerate(self.dec_rnn.lstm): - h_new, c_new = cell(x, state_h[index], state_c[index]) - next_h.append(h_new) - next_c.append(c_new) - x = h_new - return x, tuple(next_h), tuple(next_c) - - -class ParakeetTDTJoint(nn.Module): - def __init__(self, config: ParakeetTDTConfig, state_dict: dict[str, torch.Tensor]): - super().__init__() - self.enc = nn.Linear(config.hidden_dim, config.joint_dim) - self.pred = nn.Linear(config.predictor_hidden_dim, config.joint_dim) - output_classes = int(state_dict["joint.joint_net.2.weight"].shape[0]) - self.joint_net = nn.Sequential( - nn.Identity(), - nn.ReLU(), - nn.Linear( - config.joint_dim, - output_classes, - ), - ) - _copy_linear_weight(self.enc, state_dict["joint.enc.weight"], bias=state_dict["joint.enc.bias"]) - _copy_linear_weight(self.pred, state_dict["joint.pred.weight"], bias=state_dict["joint.pred.bias"]) - _copy_linear_weight( - self.joint_net[2], - state_dict["joint.joint_net.2.weight"], - bias=state_dict["joint.joint_net.2.bias"], - ) - - def forward(self, encoder_frame: torch.Tensor, predictor_hidden: torch.Tensor) -> torch.Tensor: - return self.joint_net(F.relu(self.enc(encoder_frame) + self.pred(predictor_hidden))) - - -class ParakeetTDTDecoderStep(nn.Module): - def __init__(self, config: ParakeetTDTConfig, state_dict: dict[str, torch.Tensor]): - super().__init__() - self.prediction = ParakeetTDTDecoderPrediction(config, state_dict) - self.joint = ParakeetTDTJoint(config, state_dict) - self.config = config - - def forward(self, encoder_frame: torch.Tensor, token_ids: torch.Tensor, *state_tensors: torch.Tensor) -> tuple[torch.Tensor, ...]: - if len(state_tensors) != self.config.predictor_num_layers * 2: - raise ValueError( - f"expected {self.config.predictor_num_layers * 2} predictor state tensors, got {len(state_tensors)}" - ) - state_h = tuple(state_tensors[0::2]) - state_c = tuple(state_tensors[1::2]) - predictor_hidden, next_h, next_c = self.prediction(token_ids, state_h, state_c) - logits = self.joint(encoder_frame, predictor_hidden) - outputs: list[torch.Tensor] = [logits] - for h, c in zip(next_h, next_c, strict=True): - outputs.append(h) - outputs.append(c) - return tuple(outputs) - - -class ParakeetTDTLocalModel(nn.Module): - def __init__(self, config: ParakeetTDTConfig, state_dict: dict[str, torch.Tensor]): - super().__init__() - self.name_or_path = config.model_source - self.family = "parakeet_tdt" - self.config = config - self.encoder = ParakeetTDTEncoder(config, state_dict) - self.decoder = nn.Module() - self.decoder.prediction = ParakeetTDTDecoderPrediction(config, state_dict) - self.joint = ParakeetTDTJoint(config, state_dict) - self.decoder_step = ParakeetTDTDecoderStep(config, state_dict) - - def forward(self, input_features: torch.Tensor) -> torch.Tensor: - return self.encoder(input_features) - - def initial_decoder_state(self, *, batch_size: int, device: torch.device, dtype: torch.dtype) -> tuple[torch.Tensor, ...]: - state: list[torch.Tensor] = [] - for _ in range(self.config.predictor_num_layers): - state.append(torch.zeros((batch_size, self.config.predictor_hidden_dim), device=device, dtype=dtype)) - state.append(torch.zeros((batch_size, self.config.predictor_hidden_dim), device=device, dtype=dtype)) - return tuple(state) - - def greedy_decode_token_ids(self, input_features: torch.Tensor) -> list[int]: - with torch.no_grad(): - encoder_hidden = self.encoder(input_features) - batch = int(encoder_hidden.shape[0]) - if batch != 1: - raise ValueError("Parakeet TDT local greedy decode currently expects batch size 1") - states = self.initial_decoder_state( - batch_size=batch, - device=encoder_hidden.device, - dtype=encoder_hidden.dtype, - ) - state_arrays = tuple(state.detach().cpu().numpy() for state in states) - encoder_hidden_np = encoder_hidden.detach().cpu().numpy() - - def _step( - frame: np.ndarray, - token_id: int, - state_values: tuple[np.ndarray, ...], - ) -> tuple[np.ndarray, tuple[np.ndarray, ...]]: - frame_tensor = torch.from_numpy(frame).to(device=encoder_hidden.device, dtype=encoder_hidden.dtype) - token_tensor = torch.tensor([[token_id]], device=encoder_hidden.device, dtype=torch.long) - state_tensors = tuple( - torch.from_numpy(value).to(device=encoder_hidden.device, dtype=encoder_hidden.dtype) - for value in state_values - ) - outputs = self.decoder_step(frame_tensor, token_tensor, *state_tensors) - logits = outputs[0].detach().cpu().numpy() - next_states = tuple(output.detach().cpu().numpy() for output in outputs[1:]) - return logits, next_states - - return greedy_decode_parakeet_tdt_token_ids( - config=self.config, - encoder_hidden_states=encoder_hidden_np, - initial_states=state_arrays, - step=_step, - ) - - def decode_token_ids(self, token_ids: list[int]) -> str: - pieces: list[str] = [] - for token_id in token_ids: - if token_id < 0 or token_id >= len(self.config.vocabulary): - continue - piece = self.config.vocabulary[token_id] - if piece.startswith("<|") and piece.endswith("|>"): - continue - pieces.append(piece) - text = "".join(pieces).replace("▁", " ") - return re.sub(r"\s+", " ", text).strip() - - -def load_tdt_local_model(model_source: str, *, torch_dtype: torch.dtype) -> ParakeetTDTLocalModel: - config = load_parakeet_tdt_config(model_source) - state_dict = _load_tensor_state_dict(model_source) - predictor_vocab_size = int(state_dict["decoder.prediction.embed.weight"].shape[0]) - if config.blank_id < 0 or config.blank_id >= predictor_vocab_size: - config = replace(config, blank_id=predictor_vocab_size - 1) - elif config.blank_id == len(config.vocabulary) - 1 and predictor_vocab_size == len(config.vocabulary) + 1: - config = replace(config, blank_id=predictor_vocab_size - 1) - model = ParakeetTDTLocalModel(config, state_dict).eval() - model.to(dtype=torch_dtype) - return model - - -def build_parakeet_tdt_component_specs( - model: ParakeetTDTLocalModel, - *, - named_tensors: dict[str, torch.Tensor], - weights_dir: str | None = None, -) -> list[ComponentModuleSpec]: - input_features = named_tensors["input_features"] - example_hidden = model.encoder(input_features) - batch_size = int(example_hidden.shape[0]) - initial_states = model.initial_decoder_state( - batch_size=batch_size, - device=example_hidden.device, - dtype=example_hidden.dtype, - ) - example_frame = example_hidden[:, :1, :].reshape(batch_size, example_hidden.shape[-1]) - example_token_id = torch.full( - (batch_size,), - int(model.config.blank_id), - device=example_hidden.device, - dtype=torch.long, - ) - - state_input_keys: list[str] = [] - state_output_keys: list[str] = [] - for index in range(model.config.predictor_num_layers): - state_input_keys.extend((f"state_h_{index}", f"state_c_{index}")) - state_output_keys.extend((f"state_h_{index}", f"state_c_{index}")) - - common_graph_meta = { - "weights_dir": weights_dir, - "task": "tdt_transcription", - "adapter_family": "parakeet_tdt", - } - - factor = max(1, int(model.config.subsampling_factor)) - npu_frames = int(os.environ.get("CACTUS_NPU_PARAKEET_FRAMES", "3000")) - if npu_frames <= 0 or npu_frames % factor != 0: - raise ValueError(f"CACTUS_NPU_PARAKEET_FRAMES must be a positive multiple of {factor}, got {npu_frames}") - npu_input_features = torch.zeros( - (int(input_features.shape[0]), npu_frames, int(input_features.shape[-1])), - device=input_features.device, - dtype=input_features.dtype, - ) - valid_frames = min(int(input_features.shape[1]), npu_frames) - npu_input_features[:, :valid_frames, :] = input_features[:, :valid_frames, :] - with torch.no_grad(): - ane_seq_len = int(model.encoder.pre_encode(npu_input_features).shape[1]) - ane_encoder = ParakeetTDTANEEncoder(model.encoder, ane_seq_len) - - return [ - ComponentModuleSpec( - component="audio_encoder", - module=model.encoder, - example_inputs=(input_features,), - input_keys=("input_features",), - output_keys=("encoder_hidden_states",), - graph_meta={**common_graph_meta, "component": "audio_encoder"}, - metadata={"family": "parakeet_tdt", "task": "tdt_transcription"}, - npu_module=ane_encoder, - npu_example_inputs=(npu_input_features,), - npu_runtime_input_count=1, - ), - ComponentModuleSpec( - component="decoder", - module=model.decoder_step, - example_inputs=(example_frame, example_token_id, *initial_states), - input_keys=("encoder_frame", "token_ids", *state_input_keys), - output_keys=("step_logits", *state_output_keys), - graph_meta={**common_graph_meta, "component": "decoder"}, - metadata={"family": "parakeet_tdt", "task": "tdt_transcription"}, - ), - ] diff --git a/python/cactus/transpile/weight_binding.py b/python/cactus/transpile/weight_binding.py deleted file mode 100644 index 68cd5992f..000000000 --- a/python/cactus/transpile/weight_binding.py +++ /dev/null @@ -1,317 +0,0 @@ -from __future__ import annotations - -from dataclasses import dataclass -import json -import os -from pathlib import Path -from typing import Any - - -@dataclass(frozen=True) -class WeightBinding: - path: str - kind: str # "weight" | "embedding" - source_name: str - - -_TOKEN_EMBEDDING_OUTPUT_NAMES = { - "token_embeddings.weights", - "token_embeddings.cq2.weights", - "token_embeddings.cq3.weights", - "token_embeddings.cq4.weights", - "decoder_token_embeddings.weights", - "decoder_token_embeddings.cq2.weights", - "decoder_token_embeddings.cq3.weights", - "decoder_token_embeddings.cq4.weights", -} - - -def _candidate_model_dir_names(model_name_or_path: str) -> list[str]: - candidates: list[str] = [] - - def _add(name: str) -> None: - name = name.strip().lower() - if name and name not in candidates: - candidates.append(name) - - raw = model_name_or_path.strip() - if not raw: - return candidates - - path = Path(raw) - _add(path.name) - if "/" in raw: - _add(raw.split("/")[-1]) - - slug = raw.replace("/", "--") - _add(slug) - _add(slug.replace("--", "-")) - - for part in path.parts: - if not part.startswith("models--"): - continue - cache_name = part[len("models--") :] - _add(cache_name) - _add(cache_name.replace("--", "-")) - _add(cache_name.split("--")[-1]) - break - - return candidates - - -def _default_weights_dir_for_model_name(model_name_or_path: str) -> str | None: - if not model_name_or_path: - return None - from ..cli.common import weights_root - root = weights_root() - for model_dir_name in _candidate_model_dir_names(model_name_or_path): - candidate = root / model_dir_name - if candidate.exists(): - return str(candidate) - return None - - -def resolve_transpile_weights_dir(graph_meta: dict[str, object]) -> str | None: - explicit = graph_meta.get("weights_dir") - if isinstance(explicit, str) and explicit: - return explicit - - family = str(graph_meta.get("adapter_family", "")).upper() - family_env = f"CACTUS_TRANSPILER_WEIGHTS_DIR_{family}" - if family and family_env in os.environ and os.environ[family_env]: - return os.environ[family_env] - - generic = os.environ.get("CACTUS_TRANSPILER_WEIGHTS_DIR") - if generic: - return generic - - model_name_or_path = graph_meta.get("model_name_or_path") - if isinstance(model_name_or_path, str) and model_name_or_path: - return _default_weights_dir_for_model_name(model_name_or_path) - return None - - -def _manifest_source_aliases(name: str) -> tuple[str, ...]: - aliases: list[str] = [] - - def _add(candidate: str) -> None: - if candidate and candidate not in aliases: - aliases.append(candidate) - - def _add_with_wrappers(candidate: str) -> None: - _add(candidate) - for prefix in ( - "module.", - "module.model.", - "adapter.", - "adapter.model.", - ): - _add(f"{prefix}{candidate}") - - raw = name.strip() - if not raw: - return () - - _add_with_wrappers(raw) - - unwrapped = raw - changed = True - while changed: - changed = False - for wrapper in ("module.model.", "adapter.model.", "module.", "adapter."): - if unwrapped.startswith(wrapper): - unwrapped = unwrapped[len(wrapper) :] - changed = True - break - - if unwrapped.endswith("lm_head.weight"): - prefix = unwrapped[: -len("lm_head.weight")] - for tied_name in ( - f"{prefix}embed_tokens.weight", - f"{prefix}language_model.embed_tokens.weight", - f"model.{prefix}embed_tokens.weight", - f"model.{prefix}language_model.embed_tokens.weight", - "embed_tokens.weight", - "language_model.embed_tokens.weight", - "model.embed_tokens.weight", - "model.language_model.embed_tokens.weight", - ): - _add_with_wrappers(tied_name) - - def _allow_multimodal_tail_alias(tail: str) -> bool: - if tail.startswith(("audio_tower.", "vision_tower.", "embed_audio.", "embed_vision.")): - return os.environ.get("CACTUS_TRANSPILER_BIND_MULTIMODAL_TOWER_WEIGHTS") != "0" - return True - - def _add_multimodal_backbone_aliases(candidate: str) -> None: - for prefix in ( - "module.multimodal_backbone.", - "module.model.multimodal_backbone.", - "adapter.multimodal_backbone.", - "adapter.model.multimodal_backbone.", - "multimodal_backbone.", - ): - if candidate.startswith(prefix): - tail = candidate[len(prefix) :] - if not _allow_multimodal_tail_alias(tail): - continue - _add_with_wrappers(tail) - _add_with_wrappers(f"model.{tail}") - - for prefix in ( - "model.audio_tower.", - "model.vision_tower.", - "model.embed_audio.", - "model.embed_vision.", - "audio_tower.", - "vision_tower.", - "embed_audio.", - "embed_vision.", - ): - if candidate.startswith(prefix): - if candidate.startswith("model."): - tail = candidate[len("model.") :] - else: - tail = candidate - if not _allow_multimodal_tail_alias(tail): - continue - _add_with_wrappers(f"multimodal_backbone.{tail}") - _add_with_wrappers(f"module.multimodal_backbone.{tail}") - - for prefix in ("encoder.", "decoder.", "model.encoder.", "model.decoder."): - if raw.startswith(prefix): - _add_with_wrappers(raw[len(prefix) :]) - - for prefix in ("model.language_model.", "language_model."): - if raw.startswith(prefix): - tail = raw[len(prefix) :] - _add_with_wrappers(tail) - _add_with_wrappers(f"backbone.{tail}") - _add_with_wrappers(f"model.{tail}") - - if raw.startswith("model."): - tail = raw[len("model.") :] - _add_with_wrappers(tail) - - def _add_text_backbone_aliases(candidate: str) -> None: - for prefix in ( - "module.backbone.", - "module.model.backbone.", - "adapter.backbone.", - "adapter.model.backbone.", - "backbone.", - ): - if candidate.startswith(prefix): - tail = candidate[len(prefix) :] - _add_with_wrappers(tail) - _add_with_wrappers(f"model.{tail}") - _add_with_wrappers(f"model.language_model.{tail}") - - for candidate in tuple(aliases): - _add_text_backbone_aliases(candidate) - _add_multimodal_backbone_aliases(candidate) - - return tuple(aliases) - - -def _manifest_entry_kind(output_name: str, explicit_kind: Any, explicit_names: list[object]) -> str: - if isinstance(explicit_kind, str) and explicit_kind: - return explicit_kind - - output_text = output_name.lower() - explicit_text = " ".join(str(name).lower() for name in explicit_names if isinstance(name, str)) - if ( - output_name in _TOKEN_EMBEDDING_OUTPUT_NAMES - or "embedding" in output_text - or "token_embeddings" in output_text - or "embed_tokens.weight" in explicit_text - or "embed_positions.weight" in explicit_text - ): - return "embedding" - return "weight" - - -def _flatten_convert_manifest(root_manifest: dict[str, Any]) -> dict[str, object]: - flattened: dict[str, object] = {} - rows = root_manifest.get("weights") - if not isinstance(rows, list): - for name, entry in root_manifest.items(): - if not isinstance(name, str): - continue - for alias in _manifest_source_aliases(name): - flattened.setdefault(alias, entry) - return flattened - - for row in rows: - if not isinstance(row, dict): - continue - output_name = row.get("output_name") or row.get("filename") - if not isinstance(output_name, str) or not output_name: - continue - - explicit_names: list[object] = [ - row.get("source_name"), - row.get("hf_name"), - row.get("adapter_name"), - ] - source_names = row.get("source_names") - if isinstance(source_names, list): - explicit_names.extend(source_names) - - entry = { - "filename": output_name, - "kind": _manifest_entry_kind(output_name, row.get("kind"), explicit_names), - } - for source_name in explicit_names: - if not isinstance(source_name, str) or not source_name: - continue - for alias in _manifest_source_aliases(source_name): - flattened.setdefault(alias, entry) - return flattened - - -def _load_weights_manifest(root: Path) -> dict[str, object]: - manifest_path = root / "weights_manifest.json" - if not manifest_path.exists(): - return {} - try: - loaded_manifest = json.loads(manifest_path.read_text()) - except json.JSONDecodeError as exc: - raise RuntimeError(f"corrupt weights_manifest.json at {manifest_path}: {exc}") from exc - if not isinstance(loaded_manifest, dict): - raise RuntimeError(f"weights_manifest.json at {manifest_path} is not a JSON object") - return _flatten_convert_manifest(loaded_manifest) - - -def _binding_from_manifest_entry(root: Path, source_name: str, entry: object) -> WeightBinding | None: - if not isinstance(entry, dict): - return None - filename = entry.get("filename") - kind = entry.get("kind", "weight") - if not isinstance(filename, str) or not isinstance(kind, str): - return None - candidate = root / filename - if not candidate.exists(): - return None - return WeightBinding(path=str(candidate), kind=kind, source_name=source_name) - - -def resolve_weight_binding(*, weights_dir: str | None, source_name: str) -> WeightBinding | None: - """Resolve an exported parameter to a converted Cactus tensor file. - - The resolver intentionally trusts only `weights_manifest.json`. It may add - wrapper-local aliases such as `module.` or `adapter.model.`, but it does not - guess old model-specific filenames from layer names. - """ - - if not weights_dir: - return None - root = Path(weights_dir) - if not root.exists(): - return None - manifest = _load_weights_manifest(root) - for alias in _manifest_source_aliases(source_name): - binding = _binding_from_manifest_entry(root, alias, manifest.get(alias)) - if binding is not None: - return WeightBinding(path=binding.path, kind=binding.kind, source_name=source_name) - return None diff --git a/python/cactus/transpile/weight_compat.py b/python/cactus/transpile/weight_compat.py deleted file mode 100644 index dd9b02b65..000000000 --- a/python/cactus/transpile/weight_compat.py +++ /dev/null @@ -1,547 +0,0 @@ -from __future__ import annotations - -from dataclasses import replace -import shutil -import struct -from pathlib import Path - -import numpy as np - -from cactus.convert.cactus_adapters.tensor_io import CACTUS_ALIGNMENT -from cactus.convert.cactus_adapters.tensor_io import CACTUS_MAGIC -from cactus.convert.cactus_adapters.tensor_io import FLAG_INTERLEAVED -from cactus.convert.cactus_adapters.tensor_io import align_offset -from cactus.convert.model_adapters.naming import gemma4_scale_factor -from cactus.transpile.runtime_compat import Graph -from cactus.transpile.weight_binding import WeightBinding -from cactus.convert.quantization.cq import FLAG_ORTHOGONAL_ROTATION -from cactus.convert.quantization.cq import FLAG_INTERLEAVED_4ROW -from cactus.convert.quantization.cq import GROUP_SIZE as CQ_GROUP_SIZE -from cactus.convert.quantization.cq import PRECISION_CQ -from cactus.convert.quantization.cq import make_codebook -from cactus.convert.quantization.cq import make_hadamard_components -from cactus.convert.quantization.cq import make_hadamard_matrix -from cactus.convert.quantization.cq import make_orthogonal_rotation -from cactus.convert.quantization.cq import pack_indices_interleaved_4row -from cactus.convert.quantization.cq import pack_indices_lsb -from cactus.convert.quantization.cq import quantize_hadamard -from cactus.convert.quantization.cq import quantize_orthogonal -from cactus.convert.quantization.cq import write_cq_tensor - - -_HEADER_SIZE = 84 -_FLAG_EXTENDED_SHAPE = 1 << 4 -_INT8 = int(Graph.INT8) -_INT4 = int(Graph.CQ1) - -_TOKEN_EMBEDDING_FILENAMES = { - "token_embeddings.weights", - "decoder_token_embeddings.weights", -} -_PER_LAYER_EMBEDDING_FILENAMES = { - "embed_tokens_per_layer.weights", -} - - -class _OpenedTensor: - def __init__( - self, - *, - path: Path, - precision: int, - shape: tuple[int, ...], - data: np.memmap, - scales: np.memmap | None, - group_size: int, - num_groups: int, - is_interleaved: bool, - original_n: int, - alignment: int, - ) -> None: - self.path = path - self.precision = precision - self.shape = shape - self.data = data - self.scales = scales - self.group_size = group_size - self.num_groups = num_groups - self.is_interleaved = is_interleaved - self.original_n = original_n - self.alignment = alignment - - -def ensure_binding_compatible(binding: WeightBinding, source_tensor: object | None = None) -> WeightBinding: - binding = ensure_embedding_binding_compatible(binding) - return _ensure_legacy_int4_weight_binding_compatible(binding, source_tensor=source_tensor) - - -def ensure_embedding_binding_compatible(binding: WeightBinding) -> WeightBinding: - if binding.kind != "embedding": - return binding - - source_path = Path(binding.path).expanduser().resolve() - opened = _open_cactus_tensor_file(source_path) - if opened.precision != _INT8: - return binding - if len(opened.shape) != 2 or opened.scales is None or opened.group_size <= 0: - return binding - - config = _embedding_cache_config(source_path.name, source_name=binding.source_name) - if config is None: - return binding - - compat_path = source_path.with_name(source_path.stem + f".cq{config['bits']}.weights") - _materialize_cq_embedding_cache( - opened, - compat_path, - bits=int(config["bits"]), - rotation=str(config["rotation"]), - layout=str(config.get("layout", "row_major")), - ) - _cleanup_legacy_fp16_cache(source_path) - return WeightBinding(path=str(compat_path), kind=binding.kind, source_name=binding.source_name) - - -def _embedding_cache_config(filename: str, *, source_name: str) -> dict[str, object] | None: - normalized_source = str(source_name or "") - if filename in _PER_LAYER_EMBEDDING_FILENAMES or normalized_source.endswith("embed_tokens_per_layer.weight"): - return {"bits": 2, "rotation": "hadamard"} - if filename in _TOKEN_EMBEDDING_FILENAMES: - return {"bits": 4, "rotation": "orthogonal", "layout": "interleaved_4row"} - return {"bits": 4, "rotation": "orthogonal"} - - -def _ensure_legacy_int4_weight_binding_compatible( - binding: WeightBinding, - *, - source_tensor: object | None, -) -> WeightBinding: - source_path = Path(binding.path).expanduser().resolve() - opened = _open_cactus_tensor_file(source_path) - if not _is_legacy_packed_int4_tensor(opened): - return binding - - compat_path = source_path.with_name(source_path.stem + ".cq4.weights") - if compat_path.exists(): - if _can_materialize_compat_weight(source_tensor): - _materialize_source_cq_weight( - source_tensor, - compat_path, - bits=4, - rotation="hadamard", - scale_factor=_compat_weight_scale_factor(source_path), - source_mtime_ns=source_path.stat().st_mtime_ns, - ) - return WeightBinding(path=str(compat_path), kind=binding.kind, source_name=binding.source_name) - - if not _can_materialize_compat_weight(source_tensor): - raise RuntimeError( - "legacy packed INT4 weight is not directly executable in the v2 runtime and is " - f"missing a CQ companion file: {source_path.name}. " - "Rerun `cactus convert ...` so the CQ weights can be materialized, " - "or generate them through `cq_convert` first." - ) - - _materialize_source_cq_weight( - source_tensor, - compat_path, - bits=4, - rotation="hadamard", - scale_factor=_compat_weight_scale_factor(source_path), - source_mtime_ns=source_path.stat().st_mtime_ns, - ) - return WeightBinding(path=str(compat_path), kind=binding.kind, source_name=binding.source_name) - - -def _is_legacy_packed_int4_tensor(opened: _OpenedTensor) -> bool: - if opened.precision != _INT4: - return False - if len(opened.shape) != 2 or opened.scales is None or opened.group_size <= 0: - return False - if not opened.is_interleaved: - return False - return True - - -def _can_materialize_compat_weight(source_tensor: object | None) -> bool: - if source_tensor is None: - return False - if isinstance(source_tensor, (str, bytes, bytearray, Path)): - return False - if np.isscalar(source_tensor): - return False - return hasattr(source_tensor, "shape") - - -def _compat_weight_scale_factor(source_path: Path) -> float: - parent_name = source_path.parent.name.lower() - if "gemma-4" in parent_name or "gemma4" in parent_name: - return gemma4_scale_factor(source_path.name) - return 1.0 - -def _cleanup_legacy_fp16_cache(source_path: Path) -> None: - legacy = source_path.with_name(source_path.stem + ".fp16.weights") - if not legacy.exists(): - return - try: - legacy.unlink() - except OSError: - pass - - -def _materialize_source_cq_weight( - tensor: object, - out_path: Path, - *, - bits: int, - rotation: str, - scale_factor: float = 1.0, - source_mtime_ns: int | None = None, -) -> None: - if source_mtime_ns is not None and out_path.exists() and out_path.stat().st_mtime_ns >= source_mtime_ns: - return - - if rotation == "orthogonal": - cq = quantize_orthogonal(tensor, bits=bits) - else: - cq = quantize_hadamard(tensor, bits=bits, use_gptq=False) - if float(scale_factor) != 1.0: - cq = replace(cq, norms=(cq.norms.astype(np.float32) * float(scale_factor)).astype(np.float16)) - write_cq_tensor(out_path, cq) - - -def _open_cactus_tensor_file(path: str | Path) -> _OpenedTensor: - tensor_path = Path(path).expanduser().resolve() - with tensor_path.open("rb") as handle: - header = handle.read(_HEADER_SIZE) - if len(header) < _HEADER_SIZE: - raise RuntimeError(f"tensor file is too small for a Cactus header: {tensor_path}") - if header[:4] != CACTUS_MAGIC: - raise RuntimeError(f"tensor file is missing the CACT header: {tensor_path}") - - flags = struct.unpack_from(" 0) - - dtype_map = { - int(Graph.INT8): np.int8, - int(Graph.FP16): np.float16, - int(Graph.FP32): np.float32, - int(Graph.CQ1): np.uint8, - 3: np.uint8, - 4: np.uint8, - 5: np.uint8, - 6: np.uint8, - } - dtype = dtype_map.get(precision) - if dtype is None: - raise RuntimeError(f"unsupported tensor precision {precision} in {tensor_path}") - - aligned_header = align_offset(header_size, alignment) - scales_offset = aligned_header if scales_bytes > 0 else 0 - data_offset = align_offset(scales_offset + scales_bytes, alignment) if scales_bytes > 0 else aligned_header - - data_element_count = byte_size // np.dtype(dtype).itemsize - data = np.memmap(tensor_path, mode="r", dtype=dtype, offset=data_offset, shape=(data_element_count,)) - scales = None - if scales_bytes > 0: - scales = np.memmap( - tensor_path, - mode="r", - dtype=np.float16, - offset=scales_offset, - shape=(scales_bytes // np.dtype(np.float16).itemsize,), - ) - return _OpenedTensor( - path=tensor_path, - precision=precision, - shape=shape, - data=data, - scales=scales, - group_size=group_size, - num_groups=num_groups, - is_interleaved=bool(flags & FLAG_INTERLEAVED), - original_n=original_n, - alignment=alignment, - ) - - -def _materialize_cq_embedding_cache( - opened: _OpenedTensor, - out_path: Path, - *, - bits: int, - rotation: str, - layout: str = "row_major", -) -> None: - src_mtime_ns = opened.path.stat().st_mtime_ns - if out_path.exists() and out_path.stat().st_mtime_ns >= src_mtime_ns: - return - - if len(opened.shape) != 2: - raise RuntimeError(f"expected rank-2 embedding tensor, got shape={opened.shape}") - - rows = int(opened.original_n or opened.shape[0]) - cols = int(opened.shape[1]) - interleaved = layout == "interleaved_4row" - if interleaved: - if rotation != "orthogonal" or bits != 4: - raise RuntimeError(f"INTERLEAVED_4ROW embedding cache requires orthogonal CQ4, got rotation={rotation} bits={bits}") - if rows % 4 != 0 or cols % 32 != 0: - raise RuntimeError(f"INTERLEAVED_4ROW embedding cache requires rows % 4 == 0 and cols % 32 == 0, got {(rows, cols)}") - if rotation == "hadamard" and cols % CQ_GROUP_SIZE != 0: - raise RuntimeError( - f"Hadamard CQ embedding conversion requires hidden dim divisible by {CQ_GROUP_SIZE}, " - f"got {cols} for {opened.path.name}" - ) - - input_scale = _estimate_embedding_input_scale(opened, rows=rows, cols=cols) - precision = int(PRECISION_CQ[int(bits)]) - group_size = cols if rotation == "orthogonal" else CQ_GROUP_SIZE - num_groups = cols // group_size - codebook = make_codebook(group_size, bits).astype(np.float16) - input_scale = np.asarray(input_scale, dtype=np.float16).reshape(cols) - recip = np.minimum(1.0 / np.maximum(input_scale.astype(np.float32), 1e-8), 65504.0).astype(np.float16) - - chunk_rows = _recommended_chunk_rows(cols=cols, rotation=rotation) - norms_tmp = out_path.with_suffix(out_path.suffix + ".norms.tmp") - data_tmp = out_path.with_suffix(out_path.suffix + ".data.tmp") - final_tmp = out_path.with_suffix(out_path.suffix + ".tmp") - - flags = 0 - trailer_parts: list[bytes] = [] - if rotation == "orthogonal": - flags |= FLAG_ORTHOGONAL_ROTATION - if interleaved: - flags |= FLAG_INTERLEAVED_4ROW - trailer_parts.append(make_orthogonal_rotation(group_size, seed=1234).astype(np.float16).tobytes()) - else: - left, right, perm = make_hadamard_components(group_size, seed=1234) - trailer_parts.extend([ - left.tobytes(), - right.tobytes(), - perm.astype(" current: - handle.write(b"\x00" * (data_offset - current)) - with data_tmp.open("rb") as packed_handle: - shutil.copyfileobj(packed_handle, handle, length=8 * 1024 * 1024) - final_tmp.replace(out_path) - finally: - for temp_path in (norms_tmp, data_tmp): - try: - temp_path.unlink() - except FileNotFoundError: - pass - - -def _recommended_chunk_rows(*, cols: int, rotation: str) -> int: - if rotation == "orthogonal": - if cols >= 8192: - return 8 - if cols >= 4096: - return 16 - return 64 - if cols >= 8192: - return 32 - if cols >= 4096: - return 48 - return 96 - - -def _estimate_embedding_input_scale( - opened: _OpenedTensor, - *, - rows: int, - cols: int, - max_sample_rows: int = 2048, - sample_chunk_rows: int = 64, -) -> np.ndarray: - if rows <= 0 or cols <= 0: - return np.ones((cols,), dtype=np.float16) - - stride = max(1, rows // max(1, max_sample_rows)) - sampled = 0 - accum = np.zeros((cols,), dtype=np.float64) - for start in range(0, rows, stride * sample_chunk_rows): - stop = min(start + sample_chunk_rows, rows) - chunk = _dequantize_int8_rows(opened, start, stop).astype(np.float32, copy=False) - accum += np.abs(chunk).sum(axis=0, dtype=np.float64) - sampled += int(stop - start) - if sampled >= max_sample_rows: - break - - if sampled <= 0: - return np.ones((cols,), dtype=np.float16) - - mean_abs = np.maximum(accum / float(sampled), 1e-6) - raw = np.power(mean_abs, -0.5, dtype=np.float64) - raw /= np.exp(np.mean(np.log(np.clip(raw, 1e-6, None)))) - return np.clip(raw, 1.0 / 8.0, 8.0).astype(np.float16) - - -def _quantize_hadamard_chunk( - chunk: np.ndarray, - *, - bits: int, - input_scale: np.ndarray, -) -> tuple[np.ndarray, np.ndarray]: - work = chunk.astype(np.float32, copy=False) * input_scale.astype(np.float32, copy=False)[None, :] - rows, cols = work.shape - group_size = CQ_GROUP_SIZE - groups = cols // group_size - rotation = make_hadamard_matrix(group_size, seed=1234).astype(np.float32) - codebook = make_codebook(group_size, bits).astype(np.float32) - - indices = np.empty((rows, cols), dtype=np.uint8) - norms = np.empty((rows, groups), dtype=np.float16) - for group_index in range(groups): - lo = group_index * group_size - hi = lo + group_size - group = work[:, lo:hi] - row_norms = np.linalg.norm(group, axis=1).clip(min=1e-8).astype(np.float32, copy=False) - rotated = (group / row_norms[:, None]) @ rotation - idx = np.abs(rotated[..., None] - codebook[None, None, :]).argmin(axis=-1).astype(np.uint8) - indices[:, lo:hi] = idx - norms[:, group_index] = row_norms.astype(np.float16) - return pack_indices_lsb(indices, group_size, bits), norms - - -def _quantize_orthogonal_chunk( - chunk: np.ndarray, - *, - bits: int, - input_scale: np.ndarray, - interleaved: bool = False, -) -> tuple[np.ndarray, np.ndarray]: - work = chunk.astype(np.float32, copy=False) * input_scale.astype(np.float32, copy=False)[None, :] - rows, cols = work.shape - rotation = make_orthogonal_rotation(cols, seed=1234).astype(np.float32) - codebook = make_codebook(cols, bits).astype(np.float32) - norms = np.linalg.norm(work, axis=1).clip(min=1e-8).astype(np.float32, copy=False) - rotated = (work / norms[:, None]) @ rotation - indices = np.abs(rotated[..., None] - codebook[None, None, :]).argmin(axis=-1).astype(np.uint8) - if interleaved: - return pack_indices_interleaved_4row(indices, cols, bits), norms.astype(np.float16).reshape(rows, 1) - return pack_indices_lsb(indices, cols, bits), norms.astype(np.float16).reshape(rows, 1) - - -def _write_cq_header( - handle, - *, - rows: int, - cols: int, - precision: int, - packed_bytes: int, - scales_bytes: int, - group_size: int, - num_groups: int, - flags: int, -) -> None: - handle.write(CACTUS_MAGIC) - handle.write(struct.pack(" np.ndarray: - if opened.scales is None or opened.group_size <= 0: - raise RuntimeError(f"INT8 embedding file is missing grouped scales: {opened.path}") - - rows = stop - start - cols = int(opened.shape[1]) - num_groups = int(opened.num_groups) - row_group = int(opened.group_size) - padded_rows = int(opened.shape[0]) - - if opened.is_interleaved: - block = 4 - padded_cols = ((cols + block - 1) // block) * block - q_blocks = opened.data.reshape(padded_rows // block, padded_cols // block, block, block) - s_blocks = opened.scales.reshape(padded_rows // block, num_groups, block) - block_start = start // block - block_end = (stop + block - 1) // block - q = np.asarray(q_blocks[block_start:block_end], dtype=np.int8).transpose(0, 2, 1, 3).reshape(-1, padded_cols) - s = np.asarray(s_blocks[block_start:block_end], dtype=np.float16).astype(np.float32).transpose(0, 2, 1).reshape(-1, num_groups) - row_offset = start - block_start * block - q = q[row_offset : row_offset + rows, :cols] - s = s[row_offset : row_offset + rows] - else: - q = np.asarray(opened.data, dtype=np.int8).reshape(padded_rows, cols)[start:stop] - s = np.asarray(opened.scales, dtype=np.float16).astype(np.float32).reshape(padded_rows, num_groups)[start:stop] - - out = np.empty((rows, cols), dtype=np.float16) - for group_index in range(num_groups): - lo = group_index * row_group - hi = min(lo + row_group, cols) - if lo >= hi: - break - scales = s[:, group_index : group_index + 1] - out[:, lo:hi] = (q[:, lo:hi].astype(np.float32) * scales).astype(np.float16) - return out diff --git a/python/pyproject.toml b/python/pyproject.toml index ed62f66e2..056fe7167 100644 --- a/python/pyproject.toml +++ b/python/pyproject.toml @@ -36,7 +36,6 @@ dependencies = [ "Pillow>=11.0.0,<12", "torchvision>=0.23.0,<1", "sentencepiece>=0.2.0", - "coremltools>=9.0; sys_platform == 'darwin'", ] [project.optional-dependencies] diff --git a/python/tests/test_bindings.py b/python/tests/test_bindings.py index 764d7f30c..9fe854290 100644 --- a/python/tests/test_bindings.py +++ b/python/tests/test_bindings.py @@ -30,7 +30,7 @@ def test_cli_module_importable(self): assert callable(create_parser) def test_model_module_importable(self): - from cactus.cli.model import ensure_weights, ensure_bundle, TranspileOptions + from cactus.cli.model import ensure_weights, ensure_bundle assert callable(ensure_weights) assert callable(ensure_bundle) @@ -225,49 +225,6 @@ def test_empty_string(self): # ── Model ID resolution tests ────────────────────────────────────── -# ── TranspileOptions tests ────────────────────────────────────────── - - -class TestTranspileOptions: - """Test the TranspileOptions dataclass.""" - - def test_defaults(self): - from cactus.cli.model import TranspileOptions - opts = TranspileOptions() - assert opts.task == "auto" - assert opts.prompt is None - assert opts.image_files is None - assert opts.audio_file is None - assert opts.max_new_tokens is None - assert opts.component_pipeline == "auto" - assert opts.components is None - assert opts.system_prompt is None - assert opts.trust_remote_code is False - assert opts.local_files_only is False - assert opts.cache_context_length is None - - def test_custom_values(self): - from cactus.cli.model import TranspileOptions - opts = TranspileOptions( - task="causal_lm_logits", - prompt="Hello", - max_new_tokens=256, - trust_remote_code=True, - cache_context_length="131072", - ) - assert opts.task == "causal_lm_logits" - assert opts.prompt == "Hello" - assert opts.max_new_tokens == 256 - assert opts.trust_remote_code is True - assert opts.cache_context_length == "131072" - - def test_frozen(self): - from cactus.cli.model import TranspileOptions - opts = TranspileOptions() - with pytest.raises(AttributeError): - opts.task = "something" - - # ── CLI parser tests ──────────────────────────────────────────────── @@ -354,14 +311,12 @@ def test_no_command_prints_help(self): args = self.parser.parse_args([]) assert args.command is None - def test_convert_absorbs_graph_flags(self): + def test_convert_rejects_graph_flags(self): import pytest - args = self.parser.parse_args(["convert", "google/gemma-4-E2B-it", - "--task", "causal_lm_logits", "--dynamic-batch"]) - assert args.command == "convert" - assert args.model_id == "google/gemma-4-E2B-it" - assert args.task == "causal_lm_logits" - assert args.dynamic_batch is True + with pytest.raises(SystemExit): + self.parser.parse_args(["convert", "google/gemma-4-E2B-it", "--task", "causal_lm_logits"]) + with pytest.raises(SystemExit): + self.parser.parse_args(["convert", "google/gemma-4-E2B-it", "--dynamic-batch"]) with pytest.raises(SystemExit): self.parser.parse_args(["transpile", "google/gemma-4-E2B-it"]) diff --git a/python/tests/test_capture_jax.py b/python/tests/test_capture_jax.py deleted file mode 100644 index 4267d8560..000000000 --- a/python/tests/test_capture_jax.py +++ /dev/null @@ -1,197 +0,0 @@ -from __future__ import annotations - -import numpy as np -import pytest - -jax = pytest.importorskip("jax") -import jax.numpy as jnp - -from cactus.transpile.capture_jax import capture_jax_function -from cactus.transpile.capture_jax import capture_jax_function_with_params -from cactus.transpile.lower import transpile_ir - - -def _execute_ir(ir, *inputs: object) -> list[np.ndarray]: - graph = transpile_ir(ir) - graph.set_inputs([np.asarray(value) for value in inputs]) - return [output.numpy() for output in graph.execute()] - - -def _assert_close(actual: object, expected: object, *, atol: float = 8e-2, rtol: float = 8e-2) -> None: - np.testing.assert_allclose(np.asarray(actual, dtype=np.float32), np.asarray(expected, dtype=np.float32), atol=atol, rtol=rtol) - - -def test_capture_jax_handles_scalar_and_broadcast_boundaries() -> None: - scalar = jnp.asarray(2.0, dtype=jnp.float16) - x = jnp.asarray([[1.0, -2.0, 3.0]], dtype=jnp.float16) - - def fn(runtime_scalar, values): - constant = jnp.asarray(0.25, dtype=jnp.float16) - return values * runtime_scalar + jnp.broadcast_to(constant, values.shape) - - ir = capture_jax_function(fn, (scalar, x)) - got = _execute_ir(ir, scalar, x)[0] - - _assert_close(got, fn(scalar, x)) - - -def test_capture_jax_gqa_repeat_then_matmul_matches_jax() -> None: - q = jnp.ones((1, 4, 3, 2), dtype=jnp.float16) - k = jnp.asarray(np.arange(12, dtype=np.float16).reshape(1, 2, 3, 2) / 10.0) - - def fn(query, key): - repeated = jnp.repeat(key, 2, axis=1) - return jnp.matmul(query, jnp.swapaxes(repeated, -1, -2)) - - ir = capture_jax_function(fn, (q, k)) - got = _execute_ir(ir, q, k)[0] - - _assert_close(got, fn(q, k)) - - -@pytest.mark.parametrize("variant", ["mean_rsqrt", "sum_div_rsqrt", "pow_neg_half", "inv_sqrt", "weight_first"]) -def test_capture_jax_rms_norm_common_spellings_lower_to_rms_norm(variant: str) -> None: - x = jnp.asarray(np.linspace(-650.0, 640.0, 16, dtype=np.float16).reshape(1, 2, 8)) - weight = jnp.asarray(np.linspace(0.8, 1.2, 8, dtype=np.float16)) - - def fn(values): - squared = values.astype(jnp.float32) * values.astype(jnp.float32) - mean_square = jnp.mean(squared, axis=-1, keepdims=True) - if variant == "sum_div_rsqrt": - inv = jax.lax.rsqrt(jnp.sum(squared, axis=-1, keepdims=True) / values.shape[-1] + 1.0e-5) - return values * inv * weight - if variant == "pow_neg_half": - return values * ((mean_square + 1.0e-5) ** -0.5) * weight - if variant == "inv_sqrt": - return values * (1.0 / jnp.sqrt(mean_square + 1.0e-5)) * weight - if variant == "weight_first": - return weight * (values * jax.lax.rsqrt(mean_square + 1.0e-5)) - return values * jax.lax.rsqrt(mean_square + 1.0e-5) * weight - - ir = capture_jax_function(fn, (x,), constant_names=("weight",)) - got = _execute_ir(ir, x)[0] - - assert any(ir.nodes[node_id].op == "rms_norm" for node_id in ir.order) - _assert_close(got, fn(x)) - - -def test_capture_jax_zero_centered_rms_norm_lowers_to_rms_norm() -> None: - x = jnp.asarray(np.linspace(-73.0, 73.0, 512, dtype=np.float16).reshape(1, 1, 512)) - scale = jnp.zeros((512,), dtype=jnp.float16) - - def fn(values): - rms = jnp.sqrt(jnp.mean(values.astype(jnp.float32) ** 2, axis=-1, keepdims=True) + 1.0e-6) - return ((1.0 + scale) * values / rms).astype(jnp.float16) - - ir = capture_jax_function(fn, (x,), constant_names=("scale",)) - got = _execute_ir(ir, x)[0] - - assert any(ir.nodes[node_id].op == "rms_norm" for node_id in ir.order) - _assert_close(got, fn(x)) - - -@pytest.mark.parametrize("variant", ["mean_rsqrt", "sum_div_rsqrt", "sqrt_div", "no_bias"]) -def test_capture_jax_layer_norm_common_spellings_lower_to_layer_norm(variant: str) -> None: - x = jnp.asarray(np.linspace(-650.0, 640.0, 16, dtype=np.float16).reshape(1, 2, 8)) - weight = jnp.asarray(np.linspace(0.8, 1.2, 8, dtype=np.float16)) - bias = jnp.asarray(np.linspace(-0.2, 0.2, 8, dtype=np.float16)) - - def fn(values): - centered = values - jnp.mean(values.astype(jnp.float32), axis=-1, keepdims=True) - if variant == "sum_div_rsqrt": - var = jnp.sum(centered.astype(jnp.float32) * centered.astype(jnp.float32), axis=-1, keepdims=True) / values.shape[-1] - normed = centered * jax.lax.rsqrt(var + 1.0e-5) - elif variant == "sqrt_div": - var = jnp.mean(centered.astype(jnp.float32) * centered.astype(jnp.float32), axis=-1, keepdims=True) - normed = centered / jnp.sqrt(var + 1.0e-5) - else: - var = jnp.mean(centered.astype(jnp.float32) * centered.astype(jnp.float32), axis=-1, keepdims=True) - normed = centered * jax.lax.rsqrt(var + 1.0e-5) - if variant == "no_bias": - return normed * weight - return normed * weight + bias - - ir = capture_jax_function(fn, (x,), constant_names=("weight", "bias")) - got = _execute_ir(ir, x)[0] - - assert any(ir.nodes[node_id].op == "layer_norm" for node_id in ir.order) - _assert_close(got, fn(x)) - - -def test_capture_jax_prenorm_residual_add_lowers_to_add_clipped() -> None: - x = jnp.asarray(np.linspace(-0.5, 0.5, 16, dtype=np.float16).reshape(1, 2, 8)) - weight = jnp.asarray(np.linspace(0.8, 1.2, 8, dtype=np.float16)) - - def fn(values): - rms = jnp.sqrt(jnp.mean(values.astype(jnp.float32) ** 2, axis=-1, keepdims=True) + 1.0e-6) - branch = (values / rms) * weight - return values + branch - - ir = capture_jax_function(fn, (x,), constant_names=("weight",)) - got = _execute_ir(ir, x)[0] - - assert any(ir.nodes[node_id].op == "add_clipped" for node_id in ir.order) - _assert_close(got, fn(x)) - - -def test_capture_jax_prefill_attention_pattern_lowers_to_attention() -> None: - q = jnp.asarray(np.linspace(-0.5, 0.5, 1 * 2 * 3 * 4, dtype=np.float16).reshape(1, 2, 3, 4)) - k = jnp.asarray(np.linspace(0.25, -0.25, 1 * 2 * 3 * 4, dtype=np.float16).reshape(1, 2, 3, 4)) - v = jnp.asarray(np.linspace(-0.75, 0.75, 1 * 2 * 3 * 4, dtype=np.float16).reshape(1, 2, 3, 4)) - mask = jnp.asarray(np.where(np.tril(np.ones((1, 2, 3, 3), dtype=np.bool_)), 0.0, -1.0e4), dtype=jnp.float16) - - def fn(query, key, value, additive_mask): - scores = (query @ key.transpose(0, 1, 3, 2)) / jnp.sqrt(jnp.asarray(4.0, dtype=jnp.float32)) - probs = jax.nn.softmax(scores + additive_mask, axis=-1) - return (probs @ value).transpose(0, 2, 1, 3).reshape(1, 3, 8) - - ir = capture_jax_function(fn, (q, k, v, mask)) - got = _execute_ir(ir, q, k, v, mask)[0] - - assert any(ir.nodes[node_id].op == "attention" for node_id in ir.order) - _assert_close(got, fn(q, k, v, mask)) - - -def test_capture_jax_attention_keeps_broadcast_mask_decomposed() -> None: - q = jnp.asarray(np.linspace(-0.5, 0.5, 2 * 4 * 5 * 8, dtype=np.float16).reshape(2, 4, 5, 8)) - k = jnp.asarray(np.linspace(0.25, -0.25, 2 * 4 * 5 * 8, dtype=np.float16).reshape(2, 4, 5, 8)) - v = jnp.asarray(np.linspace(-0.75, 0.75, 2 * 4 * 5 * 8, dtype=np.float16).reshape(2, 4, 5, 8)) - mask = jnp.asarray(np.where(np.tril(np.ones((1, 1, 5, 5), dtype=np.bool_)), 0.0, -1.0e4), dtype=jnp.float16) - - def fn(query, key, value, additive_mask): - scores = (query @ key.transpose(0, 1, 3, 2)) / jnp.sqrt(jnp.asarray(8.0, dtype=jnp.float32)) - probs = jax.nn.softmax(scores + additive_mask, axis=-1) - return (probs @ value).transpose(0, 2, 1, 3).reshape(2, 5, 32) - - ir = capture_jax_function(fn, (q, k, v, mask)) - got = _execute_ir(ir, q, k, v, mask)[0] - - assert not any(ir.nodes[node_id].op == "attention" for node_id in ir.order) - _assert_close(got, fn(q, k, v, mask)) - - -def test_capture_flax_module_smoke_when_available() -> None: - pytest.importorskip("flax") - from flax import linen as nn - - class TinyFlaxModule(nn.Module): - @nn.compact - def __call__(self, token_ids, features): - embedded = nn.Embed(8, 4)(token_ids) - projected = nn.Dense(4)(features) - hidden = nn.LayerNorm()(embedded + projected) - return nn.Dense(3)(jax.nn.silu(hidden)) - - model = TinyFlaxModule() - token_ids = jnp.asarray([[1, 2]], dtype=jnp.int32) - features = jnp.asarray([[[0.2, -0.1], [0.4, 0.3]]], dtype=jnp.float16) - params = model.init(jax.random.PRNGKey(0), token_ids, features)["params"] - - def fn(model_params, ids, feats): - return model.apply({"params": model_params}, ids, feats) - - ir = capture_jax_function_with_params(fn, params, (token_ids, features)) - got = _execute_ir(ir, token_ids, features)[0] - - assert any(ir.nodes[node_id].op == "layer_norm" for node_id in ir.order) - _assert_close(got, fn(params, token_ids, features), atol=1.2e-1, rtol=1.2e-1) diff --git a/python/tests/test_capture_jax_user_graph_bundle.py b/python/tests/test_capture_jax_user_graph_bundle.py deleted file mode 100644 index ac272cf74..000000000 --- a/python/tests/test_capture_jax_user_graph_bundle.py +++ /dev/null @@ -1,196 +0,0 @@ -from __future__ import annotations - -import json -from pathlib import Path - -import numpy as np -import pytest - -jax = pytest.importorskip("jax") -import jax.numpy as jnp - -from cactus.transpile.capture_jax import JaxGraphSpec -from cactus.transpile.jax_user_graph_bundle import build_jax_generation_graph_bundle -from cactus.transpile.jax_user_graph_bundle import build_jax_user_graph_bundle -from cactus.transpile.jax_user_graph_bundle import load_jax_user_graph_bundle - - -def _assert_close(actual: object, expected: object, *, atol: float = 8e-2, rtol: float = 8e-2) -> None: - np.testing.assert_allclose(np.asarray(actual, dtype=np.float32), np.asarray(expected, dtype=np.float32), atol=atol, rtol=rtol) - - -def _toy_export(): - params = { - "encoder_w": jnp.asarray([[0.2, -0.1, 0.4], [0.3, 0.5, -0.2]], dtype=jnp.float16), - "decoder_w": jnp.asarray([[0.1, -0.3], [0.4, 0.2], [-0.2, 0.6]], dtype=jnp.float16), - "out_b": jnp.asarray([0.01, -0.02], dtype=jnp.float16), - } - source = jnp.asarray([[[1.0, -0.5], [0.25, 0.75]]], dtype=jnp.float16) - target = jnp.asarray([[[0.4, -0.1, 0.2]]], dtype=jnp.float16) - - def encoder_fn(model_params, source_features): - encoded = source_features @ model_params["encoder_w"] - return encoded * jax.nn.sigmoid(encoded) - - def decoder_fn(model_params, target_features, encoder_out): - context = jnp.mean(encoder_out, axis=1, keepdims=True) - return (target_features + context) @ model_params["decoder_w"] + model_params["out_b"] - - encoder_out = encoder_fn(params, source) - specs = ( - JaxGraphSpec( - name="encoder", role="encoder", fn=encoder_fn, example_args=(source,), - input_names=("source_features",), output_names=("encoder_out",), - ), - JaxGraphSpec( - name="decoder", role="generic", fn=decoder_fn, example_args=(target, encoder_out), - input_names=("target_features", "encoder_out"), output_names=("logits",), - ), - ) - return params, specs, source, target, encoder_fn, decoder_fn - - -def _build_toy_bundle(tmp_path: Path, **kwargs): - params, specs, source, target, encoder_fn, decoder_fn = _toy_export() - result = build_jax_user_graph_bundle( - params=params, - specs=specs, - output_dir=tmp_path / "bundle", - model_id="toy-jax", - **kwargs, - ) - return result, params, source, target, encoder_fn, decoder_fn - - -def test_jax_user_graph_bundle_writes_manifest_graphs_and_weights(tmp_path: Path) -> None: - result, params, source, target, encoder_fn, decoder_fn = _build_toy_bundle( - tmp_path, - task="generic", - inputs_metadata={"owner": "client"}, - ) - - encoder_out = result.bundle.execute("encoder", source)[0].numpy() - logits = result.bundle.execute("decoder", target, encoder_out)[0].numpy() - manifest = json.loads(result.components_manifest_path.read_text()) - - assert manifest["model_source"] == "jax_user_graph" - assert manifest["component_order"] == ["encoder", "decoder"] - assert (tmp_path / "bundle/components/encoder/graph.cactus").exists() - assert (tmp_path / "bundle/components/decoder/graph.cactus").exists() - for component in manifest["components"]: - raw_ir_path, optimized_ir_path = (tmp_path / "bundle" / component[key] for key in ("raw_ir", "optimized_ir")) - assert raw_ir_path.exists() and optimized_ir_path.exists() - assert json.loads(raw_ir_path.read_text())["graph"]["meta"]["frontend"] == "jax" - assert json.loads(optimized_ir_path.read_text())["graph"]["outputs"] == component["outputs"] - assert (tmp_path / "bundle/weights_manifest.json").exists() - _assert_close(encoder_out, encoder_fn(params, source)) - _assert_close(logits, decoder_fn(params, target, encoder_fn(params, source))) - - -def test_jax_user_graph_bundle_loads_saved_graphs_and_mmap_weights(tmp_path: Path) -> None: - result, params, source, target, encoder_fn, decoder_fn = _build_toy_bundle(tmp_path) - loaded = load_jax_user_graph_bundle(result.output_dir) - - encoder_out = loaded.execute("encoder", source)[0] - logits = loaded.execute("decoder", target, encoder_out)[0].numpy() - - assert set(loaded.graphs) == {"encoder", "decoder"} - assert loaded.graphs["encoder"].logical_inputs == ["source_features"] - assert loaded.graphs["decoder"].logical_outputs == ["logits"] - _assert_close(encoder_out.numpy(), encoder_fn(params, source)) - _assert_close(logits, decoder_fn(params, target, encoder_fn(params, source))) - loaded.reset() - - -def test_jax_user_graph_bundle_supports_external_weights_dir(tmp_path: Path) -> None: - params = { - "w": jnp.asarray([[0.5], [-0.25]], dtype=jnp.float16), - "b": jnp.asarray([0.1], dtype=jnp.float16), - } - x = jnp.asarray([[2.0, -1.0]], dtype=jnp.float16) - - def fn(model_params, values): - return values @ model_params["w"] + model_params["b"] - - result = build_jax_user_graph_bundle( - params=params, - specs=(JaxGraphSpec(name="project", fn=fn, example_args=(x,), input_names=("x",), output_names=("y",)),), - output_dir=tmp_path / "bundle", - weights_dir=tmp_path / "weights", - model_id="external-weights", - ) - loaded = load_jax_user_graph_bundle(result.components_manifest_path) - manifest = json.loads(result.components_manifest_path.read_text()) - - assert result.weights_dir == tmp_path / "weights" - assert (tmp_path / "weights/weights_manifest.json").exists() - assert manifest["weights_dir"] == str(tmp_path / "weights") - _assert_close(loaded.execute("project", np.asarray(x, dtype=np.float32))[0].numpy(), fn(params, x)) - - -def test_jax_generation_decoder_step_fuses_attention_without_internal_cache(tmp_path: Path) -> None: - params = { - "wq": jnp.eye(16, dtype=jnp.float16), - "wk": jnp.eye(16, dtype=jnp.float16), - "wv": jnp.eye(16, dtype=jnp.float16), - } - x = jnp.asarray(np.linspace(-0.75, 0.75, 16, dtype=np.float16).reshape(1, 1, 16)) - - def decoder_step(model_params, values): - q = (values @ model_params["wq"]).reshape(1, 1, 2, 8).transpose(0, 2, 1, 3) - k = (values @ model_params["wk"]).reshape(1, 1, 2, 8).transpose(0, 2, 1, 3) - v = (values @ model_params["wv"]).reshape(1, 1, 2, 8).transpose(0, 2, 1, 3) - scores = (q @ k.transpose(0, 1, 3, 2)) / jnp.sqrt(jnp.asarray(8.0, dtype=jnp.float32)) - probs = jax.nn.softmax(scores, axis=-1).astype(jnp.float16) - return (probs @ v).transpose(0, 2, 1, 3).reshape(1, 1, 16) - - result = build_jax_generation_graph_bundle( - params=params, - decoder_step=JaxGraphSpec(name="decoder_step", fn=decoder_step, example_args=(x,), output_names=("hidden",)), - output_dir=tmp_path / "bundle", - model_id="decoder-step-attention", - ) - graph = result.bundle.graphs["decoder_step"].ir_graph - attention_nodes = [node for node in graph.nodes.values() if node.op == "attention"] - - assert "use_internal_kv_cache" not in graph.meta - assert len(attention_nodes) == 1 - assert result.bundle.graphs["decoder_step"].graph.cache_state_tensors == [] - _assert_close(result.bundle.execute("decoder_step", x)[0].numpy(), decoder_step(params, x)) - - -def test_jax_generation_decoder_step_fuses_cross_attention(tmp_path: Path) -> None: - params = { - "wq": jnp.eye(16, dtype=jnp.float16), - } - x = jnp.asarray(np.linspace(-0.75, 0.75, 16, dtype=np.float16).reshape(1, 1, 16)) - key = jnp.asarray(np.linspace(-0.5, 0.5, num=1 * 2 * 4 * 8).reshape(1, 2, 4, 8), dtype=jnp.float16) - value = jnp.asarray(np.linspace(0.75, -0.25, num=1 * 2 * 4 * 8).reshape(1, 2, 4, 8), dtype=jnp.float16) - mask = jnp.asarray([[[[True, True, True, False]]]], dtype=jnp.bool_) - - def decoder_step(model_params, values, cross_k, cross_v, cross_mask): - q = (values @ model_params["wq"]).reshape(1, 1, 2, 8).transpose(0, 2, 1, 3) - scores = (q @ cross_k.transpose(0, 1, 3, 2)) / jnp.sqrt(jnp.asarray(8.0, dtype=jnp.float32)) - scores = jnp.where(cross_mask, scores, jnp.finfo(scores.dtype).min) - probs = jax.nn.softmax(scores, axis=-1).astype(jnp.float16) - return (probs @ cross_v).transpose(0, 2, 1, 3).reshape(1, 1, 16) - - result = build_jax_generation_graph_bundle( - params=params, - decoder_step=JaxGraphSpec( - name="decoder_step", fn=decoder_step, example_args=(x, key, value, mask), output_names=("hidden",), - ), - output_dir=tmp_path / "bundle", - model_id="cross-attention-step", - ) - graph = result.bundle.graphs["decoder_step"].ir_graph - attention_nodes = [node for node in graph.nodes.values() if node.op == "attention"] - - assert len(attention_nodes) == 1 - assert attention_nodes[0].meta["rewritten_from"] == "jax_decoder_step_cross_attention" - assert len(attention_nodes[0].inputs) == 4 - assert result.bundle.graphs["decoder_step"].graph.cache_state_tensors == [] - _assert_close( - result.bundle.execute("decoder_step", x, key, value, mask)[0].numpy(), - decoder_step(params, x, key, value, mask), - ) diff --git a/python/tests/test_cli_transpile_defaults.py b/python/tests/test_cli_transpile_defaults.py deleted file mode 100644 index 592e029fd..000000000 --- a/python/tests/test_cli_transpile_defaults.py +++ /dev/null @@ -1,157 +0,0 @@ -from __future__ import annotations - -from pathlib import Path -from types import SimpleNamespace - -import pytest -from cactus import cli -from cactus.cli import convert as convert_cli -from cactus.transpile.model_adapters import _cache_context_length - - -class _ConfigWithText: - def get_text_config(self): - return SimpleNamespace(max_position_embeddings=128000) - - -def test_cache_context_length_uses_explicit_value() -> None: - model = SimpleNamespace(config=SimpleNamespace(max_position_embeddings=40960)) - - assert _cache_context_length( - model, - input_seq_len=2048, - cache_context_length="32768", - fallback_extra_tokens=512, - ) == 32768 - - -def test_cache_context_length_reads_top_level_config() -> None: - model = SimpleNamespace(config=SimpleNamespace(max_position_embeddings=40960)) - - assert _cache_context_length( - model, - input_seq_len=2048, - cache_context_length=None, - fallback_extra_tokens=512, - ) == 40960 - - -def test_cache_context_length_reads_text_config() -> None: - model = SimpleNamespace(config=_ConfigWithText()) - - assert _cache_context_length( - model, - input_seq_len=2048, - cache_context_length="auto", - fallback_extra_tokens=512, - ) == 128000 - - -def test_cache_context_length_falls_back_to_capture_size() -> None: - model = SimpleNamespace(config=SimpleNamespace()) - - assert _cache_context_length( - model, - input_seq_len=2048, - cache_context_length=None, - fallback_extra_tokens=512, - ) == 2560 - - -def test_cache_context_length_rejects_non_positive_explicit_value() -> None: - with pytest.raises(ValueError): - _cache_context_length( - SimpleNamespace(config=SimpleNamespace()), - input_seq_len=2048, - cache_context_length="0", - fallback_extra_tokens=512, - ) - - -def test_cmd_convert_builds_graph_after_weights(monkeypatch, tmp_path: Path) -> None: - """`cactus convert` quantizes weights and then builds the runtime graph in - one step; the graph build is pointed at the weights just written.""" - parser = cli.create_parser() - out = tmp_path / "out" - args = parser.parse_args(["convert", "google/gemma-4-E2B-it", str(out), "--reconvert"]) - - weight_calls: list[dict] = [] - - def _fake_ensure_weights(model_id, **kw): - weight_calls.append(kw) - Path(kw["output_dir"]).mkdir(parents=True, exist_ok=True) - return kw["output_dir"] - - transpile_calls: list[dict] = [] - - def _fake_run_transpile(model_id, **kw): - transpile_calls.append(kw) - return 0 - - import cactus.cli.model as model_mod - import cactus.cli.transpile as transpile_mod - monkeypatch.setattr(model_mod, "ensure_weights", _fake_ensure_weights) - monkeypatch.setattr(model_mod, "_default_multimodal_assets", lambda: ([], None)) - monkeypatch.setattr(transpile_mod, "run_transpile", _fake_run_transpile) - - rc = convert_cli.cmd_convert(args) - - assert rc == 0 - assert len(weight_calls) == 1 - assert len(transpile_calls) == 1 - extra_args = transpile_calls[0]["extra_args"] - assert "--weights-dir" in extra_args - assert str(out) in extra_args - - -def test_cmd_convert_weights_only_skips_graph(monkeypatch, tmp_path: Path) -> None: - """`--weights-only` stops after CQ quantization, before the runtime graph.""" - parser = cli.create_parser() - out = tmp_path / "out" - args = parser.parse_args(["convert", "google/gemma-4-E2B-it", str(out), "--weights-only"]) - - def _fake_ensure_weights(model_id, **kw): - Path(kw["output_dir"]).mkdir(parents=True, exist_ok=True) - return kw["output_dir"] - - def _fail_if_called(*a, **kw): - raise AssertionError("--weights-only must skip the runtime graph build") - - import cactus.cli.model as model_mod - import cactus.cli.transpile as transpile_mod - monkeypatch.setattr(model_mod, "ensure_weights", _fake_ensure_weights) - monkeypatch.setattr(transpile_mod, "run_transpile", _fail_if_called) - - rc = convert_cli.cmd_convert(args) - assert rc == 0 - - -def test_cli_convert_absorbs_graph_flags() -> None: - """The former `transpile` options now live on `cactus convert`; the separate - `transpile` subcommand no longer exists.""" - parser = cli.create_parser() - args = parser.parse_args(["convert", "google/gemma-4-E2B-it", - "--weights-dir", "/tmp/x", "--task", "causal_lm_logits", - "--dynamic-batch", "--max-slots", "8"]) - assert args.command == "convert" - assert args.model_id == "google/gemma-4-E2B-it" - assert args.weights_dir == "/tmp/x" - assert args.task == "causal_lm_logits" - assert args.dynamic_batch is True - assert args.max_slots == 8 - - with pytest.raises(SystemExit): - parser.parse_args(["transpile", "google/gemma-4-E2B-it"]) - - -def test_run_accepts_local_bundle_path() -> None: - """`cactus run` accepts a HF id (org/model) OR a local path. Bare names - like 'whisper-base' (no slash) are rejected.""" - parser = cli.create_parser() - args = parser.parse_args(["run", "/tmp/bundle", - "--prompt", "hi", "--system", "be brief", "--thinking"]) - assert args.command == "run" - assert args.model_id == "/tmp/bundle" - assert args.prompt == "hi" - assert args.system == "be brief" - assert args.thinking is True diff --git a/python/tests/test_component_partition.py b/python/tests/test_component_partition.py deleted file mode 100644 index a87ca03bd..000000000 --- a/python/tests/test_component_partition.py +++ /dev/null @@ -1,55 +0,0 @@ -from __future__ import annotations - -from cactus.transpile.component_partition import COMPONENT_LM_ENCODER -from cactus.transpile.component_partition import COMPONENT_VISION_ENCODER -from cactus.transpile.component_partition import classify_node_component -from cactus.transpile.component_partition import classify_value_component -from cactus.transpile.graph_ir import IRNode -from cactus.transpile.graph_ir import IRValue - - -def test_component_partition_ignores_torch_layer_names() -> None: - value = IRValue( - id="weight", - meta={ - "source_name": "model.vision_tower.encoder.layers.0.self_attn.q_proj.weight", - "torch_name": "model.layers.0.mlp.gate_proj.weight", - }, - ) - - assert classify_value_component( - "weight", - value, - family="", - task="multimodal_causal_lm_logits", - ) is None - - -def test_component_partition_uses_semantic_input_names() -> None: - assert ( - classify_value_component( - "pixel_values", - IRValue(id="pixel_values"), - family="", - task="multimodal_causal_lm_logits", - ) - == COMPONENT_VISION_ENCODER - ) - - -def test_component_partition_uses_ops_for_multimodal_merge() -> None: - node = IRNode( - id="embedding_0", - op="embedding", - inputs=["input_ids", "weight"], - outputs=["inputs_embeds"], - ) - - assert ( - classify_node_component( - node, - family="", - task="multimodal_causal_lm_logits", - ) - == COMPONENT_LM_ENCODER - ) diff --git a/python/tests/test_encoder_cross_kv_route.py b/python/tests/test_encoder_cross_kv_route.py deleted file mode 100644 index bba3627e9..000000000 --- a/python/tests/test_encoder_cross_kv_route.py +++ /dev/null @@ -1,79 +0,0 @@ -from __future__ import annotations - -from types import SimpleNamespace - -import torch - -from cactus.transpile.model_adapters import build_component_module_specs - - -class ToyNeedle(torch.nn.Module): - def __init__(self) -> None: - super().__init__() - self.config = SimpleNamespace( - model_type="needle", - decoder_start_token_id=1, - pad_token_id=0, - ) - - def cactus_source_encode( - self, - input_ids: torch.Tensor, - attention_mask: torch.Tensor, - ) -> tuple[torch.Tensor, torch.Tensor]: - hidden = input_ids.to(dtype=torch.float32).unsqueeze(-1).expand(-1, -1, 2).contiguous() - return hidden, attention_mask[:, None, None, :].to(dtype=torch.bool) - - def cactus_decoder_cross_kv( - self, - encoder_hidden_states: torch.Tensor, - encoder_attention_mask: torch.Tensor, - ) -> tuple[torch.Tensor, ...]: - del encoder_attention_mask - batch, seq, _ = encoder_hidden_states.shape - return tuple(torch.zeros((batch, seq, 1, 2), dtype=torch.float32) for _ in range(4)) - - def cactus_decoder_step( - self, - decoder_input_ids: torch.Tensor, - position_ids: torch.Tensor, - encoder_attention_mask: torch.Tensor, - *cross_kv: torch.Tensor, - ) -> torch.Tensor: - del position_ids, encoder_attention_mask, cross_kv - return torch.zeros((decoder_input_ids.shape[0], decoder_input_ids.shape[1], 8), dtype=torch.float32) - - -def test_needle_encoder_cross_kv_route_specs_are_metadata_driven() -> None: - specs = build_component_module_specs( - ToyNeedle(), - task="causal_lm_logits", - named_tensors={"input_ids": torch.tensor([[1, 2, 0]], dtype=torch.long)}, - ) - - assert specs is not None - assert [spec.component for spec in specs] == ["source_encoder", "decoder_cross_kv", "decoder_step"] - - source, cross_kv, decoder_step = specs - assert source.input_keys == ("input_ids", "attention_mask") - assert source.output_keys == ("encoder_hidden_states", "encoder_attention_mask") - assert source.metadata["runtime_route"] == "encoder_cross_kv_decoder_step" - assert source.metadata["runtime_role"] == "source_encoder" - assert source.metadata["source_kind"] == "text_tokens" - - assert cross_kv.input_keys == ("encoder_hidden_states", "encoder_attention_mask") - assert cross_kv.output_keys == ("cross_k_0", "cross_v_0", "cross_k_1", "cross_v_1") - assert cross_kv.metadata["runtime_role"] == "decoder_cross_kv" - - assert decoder_step.input_keys == ( - "decoder_input_ids", - "position_ids", - "encoder_attention_mask", - "cross_k_0", - "cross_v_0", - "cross_k_1", - "cross_v_1", - ) - assert decoder_step.output_keys == ("logits",) - assert decoder_step.metadata["runtime_role"] == "decoder_step" - assert decoder_step.graph_meta["use_internal_kv_cache"] is True diff --git a/python/tests/test_nomic_text_embedding.py b/python/tests/test_nomic_text_embedding.py deleted file mode 100644 index 716c89286..000000000 --- a/python/tests/test_nomic_text_embedding.py +++ /dev/null @@ -1,37 +0,0 @@ -import sys -from pathlib import Path - -sys.path.insert(0, str(Path(__file__).parent.parent)) - -from cactus.transpile.component_plan import infer_component_plan_from_config - - -def test_nomic_config_routes_to_text_embedding(): - config = {"model_type": "nomic_bert", "architectures": ["NomicBertModel"]} - plan = infer_component_plan_from_config(config, model_id="nomic-ai/nomic-embed-text-v2-moe") - assert plan is not None - assert plan.task == "text_embedding" - assert plan.components == ("text_embedding",) - assert plan.force_component_pipeline - - -def test_nomic_model_id_routes_to_text_embedding(): - plan = infer_component_plan_from_config({}, model_id="nomic-ai/nomic-embed-text-v2-moe") - assert plan is not None - assert plan.task == "text_embedding" - - -def test_family_key_detects_nomic(): - import torch - - from cactus.transpile.model_adapters import _family_key - - class _Cfg: - model_type = "nomic_bert" - - class _StubNomic(torch.nn.Module): - def __init__(self): - super().__init__() - self.config = _Cfg() - - assert _family_key(_StubNomic()) == "nomic" diff --git a/python/tests/test_optimize_gemma4_attention.py b/python/tests/test_optimize_gemma4_attention.py deleted file mode 100644 index 36fddc19d..000000000 --- a/python/tests/test_optimize_gemma4_attention.py +++ /dev/null @@ -1,328 +0,0 @@ -import torch - -from cactus.transpile.canonicalize.utils import rebuild_graph -from cactus.transpile.graph_ir import IRGraph -from cactus.transpile.graph_ir import IRNode -from cactus.transpile.graph_ir import IRValue -from cactus.transpile.optimize_graph import normalize_gemma4_decoder_attention_semantics -from cactus.transpile.optimize_graph import precompute_rope_tables - - -def test_normalize_gemma4_full_attention_uses_sequence_window_compat() -> None: - graph = IRGraph( - values={ - "query": IRValue(id="query", shape=(1, 8, 800, 512), dtype="fp16"), - "key": IRValue(id="key", shape=(1, 1, 800, 512), dtype="fp16"), - "value": IRValue(id="value", shape=(1, 1, 800, 512), dtype="fp16"), - "full_attention_mask": IRValue( - id="full_attention_mask", - shape=(1, 1, 800, 800), - dtype="bool", - ), - "proj": IRValue(id="proj", shape=(1536, 4096), dtype="fp16"), - "out": IRValue(id="out", shape=(1, 800, 1536), dtype="fp16", producer="attn"), - }, - nodes={ - "attn": IRNode( - id="attn", - op="attention_block", - inputs=["query", "key", "value", "full_attention_mask", "proj"], - outputs=["out"], - attrs={ - "has_mask": True, - "is_causal": False, - "window_size": 0, - "attention_output_shape": (1, 800, 4096), - }, - meta={"attention_layer_type": "full_attention"}, - ) - }, - order=["attn"], - inputs=["query", "key", "value", "full_attention_mask"], - outputs=["out"], - constants={}, - meta={ - "adapter_family": "gemma4", - "component": "decoder", - "input_names": ("query", "key", "value", "full_attention_mask"), - }, - ) - - changed = normalize_gemma4_decoder_attention_semantics(graph) - - assert changed is True - assert graph.nodes["attn"].inputs == ["query", "key", "value", "proj"] - assert graph.nodes["attn"].attrs["has_mask"] is False - assert graph.nodes["attn"].attrs["is_causal"] is True - assert graph.nodes["attn"].attrs["window_size"] == 800 - assert graph.nodes["attn"].meta["gemma4_full_attention_window_compat"] is True - - -def test_normalize_gemma4_sliding_attention_elides_runtime_mask() -> None: - graph = IRGraph( - values={ - "query": IRValue(id="query", shape=(1, 8, 800, 256), dtype="fp16"), - "key": IRValue(id="key", shape=(1, 8, 800, 256), dtype="fp16"), - "value": IRValue(id="value", shape=(1, 8, 800, 256), dtype="fp16"), - "sliding_attention_mask": IRValue( - id="sliding_attention_mask", - shape=(1, 1, 800, 800), - dtype="bool", - ), - "proj": IRValue(id="proj", shape=(1536, 2048), dtype="fp16"), - "out": IRValue(id="out", shape=(1, 800, 1536), dtype="fp16", producer="attn"), - }, - nodes={ - "attn": IRNode( - id="attn", - op="attention_block", - inputs=["query", "key", "value", "sliding_attention_mask", "proj"], - outputs=["out"], - attrs={ - "has_mask": True, - "is_causal": False, - "window_size": 0, - "attention_output_shape": (1, 800, 2048), - }, - meta={"attention_layer_type": "sliding_attention"}, - ) - }, - order=["attn"], - inputs=["query", "key", "value", "sliding_attention_mask"], - outputs=["out"], - constants={}, - meta={ - "adapter_family": "gemma4", - "component": "decoder", - "sliding_window": 512, - "input_names": ("query", "key", "value", "sliding_attention_mask"), - }, - ) - - changed = normalize_gemma4_decoder_attention_semantics(graph) - - assert changed is True - assert graph.nodes["attn"].inputs == ["query", "key", "value", "proj"] - assert graph.nodes["attn"].attrs["has_mask"] is False - assert graph.nodes["attn"].attrs["is_causal"] is True - assert graph.nodes["attn"].attrs["window_size"] == 512 - - -def test_normalize_gemma4_attention_recovers_layer_hints_from_graph_meta() -> None: - graph = IRGraph( - values={ - "query_0": IRValue(id="query_0", shape=(1, 8, 800, 256), dtype="fp16"), - "key_0": IRValue(id="key_0", shape=(1, 8, 800, 256), dtype="fp16"), - "value_0": IRValue(id="value_0", shape=(1, 8, 800, 256), dtype="fp16"), - "sliding_attention_mask": IRValue( - id="sliding_attention_mask", - shape=(1, 1, 800, 800), - dtype="bool", - ), - "out_0": IRValue(id="out_0", shape=(1, 8, 800, 256), dtype="fp16", producer="attn_0"), - "query_1": IRValue(id="query_1", shape=(1, 8, 800, 256), dtype="fp16"), - "key_1": IRValue(id="key_1", shape=(1, 8, 800, 256), dtype="fp16"), - "value_1": IRValue(id="value_1", shape=(1, 8, 800, 256), dtype="fp16"), - "full_attention_mask": IRValue( - id="full_attention_mask", - shape=(1, 1, 800, 800), - dtype="bool", - ), - "out_1": IRValue(id="out_1", shape=(1, 8, 800, 256), dtype="fp16", producer="attn_1"), - }, - nodes={ - "attn_0": IRNode( - id="attn_0", - op="attention", - inputs=["query_0", "key_0", "value_0", "sliding_attention_mask"], - outputs=["out_0"], - attrs={"is_causal": False, "window_size": 0}, - ), - "attn_1": IRNode( - id="attn_1", - op="attention", - inputs=["query_1", "key_1", "value_1", "full_attention_mask"], - outputs=["out_1"], - attrs={"is_causal": False, "window_size": 0}, - ), - }, - order=["attn_0", "attn_1"], - inputs=[ - "query_0", - "key_0", - "value_0", - "sliding_attention_mask", - "query_1", - "key_1", - "value_1", - "full_attention_mask", - ], - outputs=["out_1"], - constants={}, - meta={ - "adapter_family": "gemma4", - "component": "decoder", - "sliding_window": 512, - "layer_types": ("sliding_attention", "full_attention"), - }, - ) - - changed = normalize_gemma4_decoder_attention_semantics(graph) - - assert changed is True - assert graph.nodes["attn_0"].meta["attention_layer_type"] == "sliding_attention" - assert graph.nodes["attn_0"].meta["attention_layer_index"] == 0 - assert graph.nodes["attn_0"].inputs == ["query_0", "key_0", "value_0"] - assert graph.nodes["attn_0"].attrs["is_causal"] is True - assert graph.nodes["attn_0"].attrs["window_size"] == 512 - assert graph.nodes["attn_1"].meta["attention_layer_type"] == "full_attention" - assert graph.nodes["attn_1"].meta["attention_layer_index"] == 1 - assert graph.nodes["attn_1"].inputs == ["query_1", "key_1", "value_1"] - assert graph.nodes["attn_1"].attrs["is_causal"] is True - assert graph.nodes["attn_1"].attrs["window_size"] == 800 - - -def test_normalize_gemma4_cached_attention_sets_per_layer_cache_lengths() -> None: - graph = IRGraph( - values={ - "query_0": IRValue(id="query_0", shape=(1, 8, 1, 256), dtype="fp16"), - "key_0": IRValue(id="key_0", shape=(1, 8, 1, 256), dtype="fp16"), - "value_0": IRValue(id="value_0", shape=(1, 8, 1, 256), dtype="fp16"), - "out_0": IRValue(id="out_0", shape=(1, 8, 1, 256), dtype="fp16", producer="attn_0"), - "query_1": IRValue(id="query_1", shape=(1, 8, 1, 256), dtype="fp16"), - "key_1": IRValue(id="key_1", shape=(1, 8, 1, 256), dtype="fp16"), - "value_1": IRValue(id="value_1", shape=(1, 8, 1, 256), dtype="fp16"), - "out_1": IRValue(id="out_1", shape=(1, 8, 1, 256), dtype="fp16", producer="attn_1"), - }, - nodes={ - "attn_0": IRNode( - id="attn_0", - op="attention", - inputs=["query_0", "key_0", "value_0"], - outputs=["out_0"], - attrs={"is_causal": False, "window_size": 0}, - ), - "attn_1": IRNode( - id="attn_1", - op="attention", - inputs=["query_1", "key_1", "value_1"], - outputs=["out_1"], - attrs={"is_causal": False, "window_size": 0}, - ), - }, - order=["attn_0", "attn_1"], - inputs=["query_0", "key_0", "value_0", "query_1", "key_1", "value_1"], - outputs=["out_1"], - constants={}, - meta={ - "adapter_family": "gemma4", - "component": "decoder_step", - "use_internal_kv_cache": True, - "sliding_window": 512, - "max_cache_seq_len": 131072, - "layer_types": ("sliding_attention", "full_attention"), - }, - ) - - changed = normalize_gemma4_decoder_attention_semantics(graph) - - assert changed is True - assert graph.nodes["attn_0"].attrs["window_size"] == 512 - assert graph.nodes["attn_0"].meta["max_cache_seq_len"] == 512 - assert graph.nodes["attn_1"].attrs["window_size"] == 0 - assert graph.nodes["attn_1"].meta["max_cache_seq_len"] == 131072 - - -def test_normalize_gemma4_attention_elides_logical_fp16_mask() -> None: - graph = IRGraph( - values={ - "query": IRValue(id="query", shape=(1, 8, 800, 256), dtype="fp16"), - "key": IRValue(id="key", shape=(1, 8, 800, 256), dtype="fp16"), - "value": IRValue(id="value", shape=(1, 8, 800, 256), dtype="fp16"), - "mask_a": IRValue(id="mask_a", shape=(1, 1, 800, 800), dtype="bool"), - "mask_b": IRValue(id="mask_b", shape=(1, 1, 800, 800), dtype="bool"), - "mask_fp16": IRValue(id="mask_fp16", shape=(1, 1, 800, 800), dtype="fp16", producer="mask_and"), - "out": IRValue(id="out", shape=(1, 8, 800, 256), dtype="fp16", producer="attn"), - }, - nodes={ - "mask_and": IRNode( - id="mask_and", - op="logical_and", - inputs=["mask_a", "mask_b"], - outputs=["mask_fp16"], - ), - "attn": IRNode( - id="attn", - op="attention", - inputs=["query", "key", "value", "mask_fp16"], - outputs=["out"], - attrs={"is_causal": False, "window_size": 0}, - meta={"attention_layer_type": "sliding_attention"}, - ), - }, - order=["mask_and", "attn"], - inputs=["query", "key", "value", "mask_a", "mask_b"], - outputs=["out"], - constants={}, - meta={ - "adapter_family": "gemma4", - "component": "decoder", - "sliding_window": 512, - }, - ) - - changed = normalize_gemma4_decoder_attention_semantics(graph) - - assert changed is True - assert graph.nodes["attn"].inputs == ["query", "key", "value"] - assert graph.nodes["attn"].attrs["is_causal"] is True - assert graph.nodes["attn"].attrs["window_size"] == 512 - - -def test_precompute_rope_tables_replaces_runtime_angle_with_fp16_lookup() -> None: - head_dim = 8 - max_seq = 4096 - theta = 10000.0 - inv_freq = 1.0 / (theta ** (torch.arange(0, head_dim, 2, dtype=torch.float32) / head_dim)) - - graph = IRGraph( - values={ - "positions": IRValue(id="positions", shape=(1, 4), dtype="int32"), - "angle": IRValue(id="angle", shape=(1, 4, head_dim // 2), dtype="fp32", producer="matmul"), - "emb": IRValue(id="emb", shape=(1, 4, head_dim), dtype="fp32", producer="cat"), - "cos": IRValue(id="cos", shape=(1, 4, head_dim), dtype="fp32", producer="cos_node"), - }, - nodes={ - "matmul": IRNode(id="matmul", op="matmul", inputs=["positions", "inv_freq"], outputs=["angle"]), - "cat": IRNode(id="cat", op="cat", inputs=["angle", "angle"], outputs=["emb"], attrs={"dim": -1}), - "cos_node": IRNode( - id="cos_node", op="scalar_cos", inputs=["emb"], outputs=["cos"], - meta={"component": "decoder"}, - ), - }, - order=["matmul", "cat", "cos_node"], - inputs=["positions"], - outputs=["cos"], - constants={"inv_freq": inv_freq}, - meta={"component": "decoder", "max_cache_seq_len": max_seq}, - ) - rebuild_graph(graph) - - assert precompute_rope_tables(graph) is True - - cos_node = graph.nodes["cos_node"] - assert cos_node.op == "embedding" - table_id, position_id = cos_node.inputs - assert table_id.startswith("c_rope_table_") - assert position_id == "positions" - - table = graph.constants[table_id] - assert table.dtype == torch.float16 - assert tuple(table.shape) == (max_seq, head_dim) - assert graph.values["cos"].dtype == "fp16" - - # Positions past the fp16-angle range (>2048) stay distinct/representable in the table. - positions = torch.arange(max_seq, dtype=torch.float64).reshape(max_seq, 1) - freqs = positions * inv_freq.to(torch.float64).reshape(1, -1) - expected = torch.cos(torch.cat((freqs, freqs), dim=-1)).to(torch.float16) - assert torch.equal(table[[0, 1, 2049, max_seq - 1]], expected[[0, 1, 2049, max_seq - 1]]) diff --git a/python/tests/test_server_live.py b/python/tests/test_server_live.py index d75f2b738..664ceae8b 100644 --- a/python/tests/test_server_live.py +++ b/python/tests/test_server_live.py @@ -37,13 +37,13 @@ def _require_bundle(relative: Path, types: set[str]) -> Path: if not candidate.exists(): pytest.fail( f"Live-test model not found: {relative}\n" - f"Prepare it with `cactus convert google/gemma-4-E2B-it` from {PROJECT_ROOT}." + f"Prepare it with `cactus download google/gemma-4-E2B-it` from {PROJECT_ROOT}." ) if not _valid_bundle(candidate): pytest.fail( f"Live-test model is not a prepared v2 bundle: {relative}\n" "Expected config.txt and components/manifest.json.\n" - f"Prepare it with `cactus convert google/gemma-4-E2B-it` from {PROJECT_ROOT}." + f"Prepare it with `cactus download google/gemma-4-E2B-it` from {PROJECT_ROOT}." ) model_type = _read_model_type(candidate) if model_type not in types: @@ -333,7 +333,7 @@ def test_live_long_streaming_generation(live_server) -> None: def test_live_transcription_wav(live_server) -> None: base_url, _ = live_server - stt = _find_bundle(["parakeet-tdt-0.6b-v3-transpiled"], STT_TYPES) + stt = _find_bundle(["parakeet-tdt-0.6b-v3"], STT_TYPES) audio = ASSETS / "test.wav" assert audio.exists() with audio.open("rb") as f: @@ -349,7 +349,7 @@ def test_live_transcription_wav(live_server) -> None: def test_live_transcription_text_response(live_server) -> None: base_url, _ = live_server - stt = _find_bundle(["parakeet-tdt-0.6b-v3-transpiled"], STT_TYPES) + stt = _find_bundle(["parakeet-tdt-0.6b-v3"], STT_TYPES) audio = ASSETS / "test.wav" with audio.open("rb") as f: res = httpx.post( @@ -365,7 +365,7 @@ def test_live_transcription_text_response(live_server) -> None: def test_live_transcription_verbose_json_segments(live_server) -> None: base_url, _ = live_server - stt = _find_bundle(["parakeet-tdt-0.6b-v3-transpiled"], STT_TYPES) + stt = _find_bundle(["parakeet-tdt-0.6b-v3"], STT_TYPES) audio = ASSETS / "test.wav" with audio.open("rb") as f: res = httpx.post( @@ -386,7 +386,7 @@ def test_live_transcription_verbose_json_segments(live_server) -> None: def test_live_transcription_rejects_unsupported_format(live_server) -> None: base_url, _ = live_server - stt = _find_bundle(["parakeet-tdt-0.6b-v3-transpiled"], STT_TYPES) + stt = _find_bundle(["parakeet-tdt-0.6b-v3"], STT_TYPES) audio = ASSETS / "test.wav" with audio.open("rb") as f: res = httpx.post( @@ -401,7 +401,7 @@ def test_live_transcription_rejects_unsupported_format(live_server) -> None: def test_live_transcription_rejects_word_timestamps(live_server) -> None: base_url, _ = live_server - stt = _find_bundle(["parakeet-tdt-0.6b-v3-transpiled"], STT_TYPES) + stt = _find_bundle(["parakeet-tdt-0.6b-v3"], STT_TYPES) audio = ASSETS / "test.wav" with audio.open("rb") as f: res = httpx.post( @@ -420,7 +420,7 @@ def test_live_transcription_rejects_word_timestamps(live_server) -> None: def test_live_transcription_rejects_non_wav(live_server) -> None: base_url, _ = live_server - stt = _find_bundle(["parakeet-tdt-0.6b-v3-transpiled"], STT_TYPES) + stt = _find_bundle(["parakeet-tdt-0.6b-v3"], STT_TYPES) res = httpx.post( f"{base_url}/v1/audio/transcriptions", data={"model": stt.name}, @@ -440,7 +440,7 @@ def _find_embed_bundle() -> Path: return candidate raise RuntimeError( "No embedding (nomic/bert) bundle under weights/; " - "run `cactus convert nomic-ai/nomic-embed-text-v2-moe`" + "run `cactus download nomic-ai/nomic-embed-text-v2-moe`" ) diff --git a/python/tests/test_transpile_aten_ops.py b/python/tests/test_transpile_aten_ops.py deleted file mode 100644 index b0f0ad7f4..000000000 --- a/python/tests/test_transpile_aten_ops.py +++ /dev/null @@ -1,124 +0,0 @@ -from __future__ import annotations - -from collections import Counter - -import torch - -from cactus.transpile.aten_ops import canonical_torch_op -from cactus.transpile.capture_pytorch import capture_model -from cactus.transpile.graph_ir import IRGraph -from cactus.transpile.graph_ir import IRNode -from cactus.transpile.graph_ir import IRValue -from cactus.transpile.normalize import normalize_target -from cactus.transpile.optimize_graph import fuse_dense_mlp_tq - - -class OddlyNamedLinearBlock(torch.nn.Module): - def __init__(self) -> None: - super().__init__() - self.pear_tree = torch.nn.Linear(4, 3, bias=False) - - def forward(self, x: torch.Tensor) -> torch.Tensor: - return torch.nn.functional.silu(self.pear_tree(x)) - - -class SameOpsDifferentNames(torch.nn.Module): - def __init__(self) -> None: - super().__init__() - self.not_a_transformer_layer = torch.nn.ModuleDict( - {"banana": torch.nn.Linear(4, 3, bias=False)} - ) - - def forward(self, x: torch.Tensor) -> torch.Tensor: - return torch.nn.functional.silu(self.not_a_transformer_layer["banana"](x)) - - -def _op_counts(module: torch.nn.Module) -> Counter[str]: - example = torch.randn(2, 4) - captured = capture_model(module, (example,)) - return Counter(captured.ir_graph.nodes[node_id].op for node_id in captured.ir_graph.order) - - -def test_aten_targets_normalize_from_real_op_overloads() -> None: - assert canonical_torch_op(torch.ops.aten.add.Tensor) == "aten.add.Tensor" - assert normalize_target(torch.ops.aten.add.Tensor) == "add" - assert normalize_target(torch.ops.aten.linear.default) == "linear" - assert normalize_target("torch.ops.aten.silu.default") == "silu" - - -def test_import_is_based_on_aten_ops_not_module_names() -> None: - first = _op_counts(OddlyNamedLinearBlock()) - second = _op_counts(SameOpsDifferentNames()) - - assert first == second - assert first["linear"] == 1 - assert first["silu"] == 1 - - -def test_import_records_canonical_aten_op_metadata() -> None: - captured = capture_model(OddlyNamedLinearBlock(), (torch.randn(2, 4),)) - node_meta = [captured.ir_graph.nodes[node_id].meta for node_id in captured.ir_graph.order] - - assert {meta["aten_op"] for meta in node_meta} >= {"aten.linear.default", "aten.silu.default"} - assert all("torch_name" in meta for meta in node_meta) - assert all("module_paths" not in meta for meta in node_meta) - - -def test_dense_mlp_fusion_uses_topology_not_layer_names() -> None: - graph = IRGraph( - values={ - "x": IRValue(id="x", shape=(1, 4), dtype="fp16"), - "pear_weight": IRValue(id="pear_weight", shape=(8, 4), dtype="fp16", meta={"path": "pear.cq4.weights"}), - "banana_weight": IRValue(id="banana_weight", shape=(8, 4), dtype="fp16", meta={"path": "banana.cq4.weights"}), - "plum_weight": IRValue(id="plum_weight", shape=(4, 8), dtype="fp16", meta={"path": "plum.cq4.weights"}), - "gate": IRValue(id="gate", shape=(1, 8), dtype="fp16", producer="gate_linear"), - "activated": IRValue(id="activated", shape=(1, 8), dtype="fp16", producer="gate_gelu"), - "up": IRValue(id="up", shape=(1, 8), dtype="fp16", producer="up_linear"), - "product": IRValue(id="product", shape=(1, 8), dtype="fp16", producer="product"), - "out": IRValue(id="out", shape=(1, 4), dtype="fp16", producer="down_linear"), - }, - nodes={ - "gate_linear": IRNode( - id="gate_linear", - op="linear", - inputs=["x", "pear_weight"], - outputs=["gate"], - attrs={"has_bias": False}, - ), - "gate_gelu": IRNode( - id="gate_gelu", - op="gelu", - inputs=["gate"], - outputs=["activated"], - ), - "up_linear": IRNode( - id="up_linear", - op="linear", - inputs=["x", "banana_weight"], - outputs=["up"], - attrs={"has_bias": False}, - ), - "product": IRNode( - id="product", - op="multiply", - inputs=["activated", "up"], - outputs=["product"], - ), - "down_linear": IRNode( - id="down_linear", - op="linear", - inputs=["product", "plum_weight"], - outputs=["out"], - attrs={"has_bias": False}, - ), - }, - order=["gate_linear", "gate_gelu", "up_linear", "product", "down_linear"], - inputs=["x"], - outputs=["out"], - constants={}, - meta={}, - ) - - assert fuse_dense_mlp_tq(graph) is True - assert graph.nodes["down_linear"].op == "dense_mlp_tq_fused" - assert graph.nodes["down_linear"].inputs == ["x", "pear_weight", "banana_weight", "plum_weight"] diff --git a/python/tests/test_transpile_layer_keys.py b/python/tests/test_transpile_layer_keys.py deleted file mode 100644 index 1bd157ef0..000000000 --- a/python/tests/test_transpile_layer_keys.py +++ /dev/null @@ -1,68 +0,0 @@ -from cactus.transpile.graph_ir import IRGraph -from cactus.transpile.graph_ir import IRNode -from cactus.transpile.graph_ir import IRValue -from cactus.transpile.lower import _kv_cache_layer_key -from cactus.transpile.optimize_graph import normalize_cached_decoder_attention_hints - - -def _decoder_graph(component: str) -> IRGraph: - values: dict[str, IRValue] = {} - nodes: dict[str, IRNode] = {} - for index in range(2): - node_id = "n_scaled_dot_product_attention" + ("" if index == 0 else f"_{index}") - for name in ("query", "key", "value"): - values[f"{name}_{index}"] = IRValue(id=f"{name}_{index}", shape=(1, 8, 1, 64), dtype="fp16") - values[f"out_{index}"] = IRValue(id=f"out_{index}", shape=(1, 8, 1, 64), dtype="fp16", producer=node_id) - nodes[node_id] = IRNode( - id=node_id, - op="scaled_dot_product_attention", - inputs=[f"query_{index}", f"key_{index}", f"value_{index}"], - outputs=[f"out_{index}"], - ) - return IRGraph( - values=values, - nodes=nodes, - order=list(nodes), - inputs=[value_id for value_id in values if not value_id.startswith("out_")], - outputs=["out_1"], - constants={}, - meta={ - "component": component, - "use_internal_kv_cache": True, - "layer_types": ["full_attention"] * 2, - }, - ) - - -def test_kv_cache_layer_key_keeps_zero_layer_index() -> None: - node = IRNode( - id="n_scaled_dot_product_attention", - op="scaled_dot_product_attention", - inputs=["query_0", "key_0", "value_0"], - outputs=["out_0"], - ) - assert _kv_cache_layer_key(node) == "n_scaled_dot_product_attention" - node.meta["layer_index"] = 5 - assert _kv_cache_layer_key(node) == "5" - node.meta["attention_layer_index"] = 0 - assert _kv_cache_layer_key(node) == "0" - - -def test_media_step_receives_attention_layer_indices() -> None: - graph = _decoder_graph("decoder_media_step") - - changed = normalize_cached_decoder_attention_hints(graph) - - assert changed is True - assert graph.nodes["n_scaled_dot_product_attention"].meta["attention_layer_index"] == 0 - assert graph.nodes["n_scaled_dot_product_attention_1"].meta["attention_layer_index"] == 1 - - -def test_step_and_media_step_layer_keys_match() -> None: - keys = [] - for component in ("decoder_step", "decoder_media_step"): - graph = _decoder_graph(component) - normalize_cached_decoder_attention_hints(graph) - keys.append([_kv_cache_layer_key(graph.nodes[node_id]) for node_id in graph.order]) - - assert keys[0] == keys[1] == ["0", "1"] diff --git a/python/tests/test_transpile_weight_compat.py b/python/tests/test_transpile_weight_compat.py deleted file mode 100644 index 29fc9c8f1..000000000 --- a/python/tests/test_transpile_weight_compat.py +++ /dev/null @@ -1,361 +0,0 @@ -from __future__ import annotations - -from pathlib import Path -import json - -import numpy as np - -from cactus.convert.cactus_adapters.tensor_io import save_tensor_with_header -from cactus.transpile.weight_binding import WeightBinding -from cactus.transpile.weight_binding import resolve_weight_binding -from cactus.transpile.weight_compat import _open_cactus_tensor_file -from cactus.transpile.weight_compat import ensure_binding_compatible -from cactus.transpile.weight_compat import ensure_embedding_binding_compatible - - -def _write_grouped_int8_embedding(path: Path, rows: int, cols: int) -> None: - rng = np.random.default_rng(1234) - tensor = rng.standard_normal((rows, cols), dtype=np.float32) - save_tensor_with_header(tensor, path, precision="INT8", model_type="generic") - - -def test_token_embedding_binding_upgrades_to_cq4(tmp_path: Path) -> None: - source = tmp_path / "token_embeddings.weights" - _write_grouped_int8_embedding(source, rows=8, cols=128) - - legacy_fp16 = tmp_path / "token_embeddings.fp16.weights" - legacy_fp16.write_bytes(b"legacy") - - binding = WeightBinding( - path=str(source), - kind="embedding", - source_name="model.embed_tokens.weight", - ) - compat = ensure_embedding_binding_compatible(binding) - - assert compat.path.endswith(".cq4.weights") - assert compat.kind == "embedding" - assert not legacy_fp16.exists() - - opened = _open_cactus_tensor_file(compat.path) - assert opened.precision == 6 - assert opened.shape == (8, 128) - assert opened.scales is not None - - -def test_per_layer_embedding_binding_upgrades_to_cq2(tmp_path: Path) -> None: - source = tmp_path / "embed_tokens_per_layer.weights" - _write_grouped_int8_embedding(source, rows=8, cols=256) - - binding = WeightBinding( - path=str(source), - kind="embedding", - source_name="model.language_model.embed_tokens_per_layer.weight", - ) - compat = ensure_embedding_binding_compatible(binding) - - assert compat.path.endswith(".cq2.weights") - - opened = _open_cactus_tensor_file(compat.path) - assert opened.precision == 4 - assert opened.shape == (8, 256) - assert opened.scales is not None - - -def test_gemma4_per_layer_projection_binding_upgrades_legacy_int4(tmp_path: Path) -> None: - rng = np.random.default_rng(1234) - source = tmp_path / "per_layer_model_proj.weights" - tensor = rng.standard_normal((8, 128), dtype=np.float32) - save_tensor_with_header(tensor, source, precision="INT4", model_type="generic") - - binding = WeightBinding( - path=str(source), - kind="weight", - source_name="module.model.model.language_model.per_layer_model_projection.weight", - ) - compat = ensure_binding_compatible(binding, source_tensor=tensor) - - assert compat.path.endswith(".cq4.weights") - opened = _open_cactus_tensor_file(compat.path) - assert opened.precision == 6 - assert opened.shape == (8, 128) - assert opened.scales is not None - - -def test_decoder_binding_prefers_existing_cq4_companion(tmp_path: Path) -> None: - rng = np.random.default_rng(4321) - source = tmp_path / "layer_0_attn_q.weights" - tensor = rng.standard_normal((8, 128), dtype=np.float32) - save_tensor_with_header(tensor, source, precision="INT4", model_type="generic") - - companion = tmp_path / "layer_0_attn_q.cq4.weights" - save_tensor_with_header(tensor, companion, precision="FP16", model_type="generic") - - binding = WeightBinding( - path=str(source), - kind="weight", - source_name="module.backbone.layers.0.self_attn.q_proj.weight", - ) - compat = ensure_binding_compatible(binding, source_tensor=0) - - assert compat.path == str(companion) - - -def test_resolve_weight_binding_does_not_guess_legacy_filenames(tmp_path: Path) -> None: - (tmp_path / "token_embeddings.weights").write_bytes(b"") - (tmp_path / "layer_0_attn_q.weights").write_bytes(b"") - (tmp_path / "embed_tokens_per_layer.weights").write_bytes(b"") - (tmp_path / "per_layer_model_proj.weights").write_bytes(b"") - - assert ( - resolve_weight_binding( - weights_dir=str(tmp_path), - source_name="model.embed_tokens.weight", - ) - is None - ) - assert ( - resolve_weight_binding( - weights_dir=str(tmp_path), - source_name="model.layers.0.self_attn.q_proj.weight", - ) - is None - ) - assert ( - resolve_weight_binding( - weights_dir=str(tmp_path), - source_name="model.language_model.embed_tokens_per_layer.weight", - ) - is None - ) - assert ( - resolve_weight_binding( - weights_dir=str(tmp_path), - source_name="model.language_model.per_layer_model_projection.weight", - ) - is None - ) - - -def test_resolve_weight_binding_prefers_exact_manifest_without_name_guessing(tmp_path: Path) -> None: - (tmp_path / "exact.weights").write_bytes(b"") - (tmp_path / "token_embeddings.weights").write_bytes(b"") - (tmp_path / "weights_manifest.json").write_text( - json.dumps( - { - "v_adapter_named_anything": { - "filename": "exact.weights", - "kind": "weight", - } - } - ) - ) - - exact = resolve_weight_binding( - weights_dir=str(tmp_path), - source_name="v_adapter_named_anything", - ) - assert exact is not None - assert exact.path.endswith("exact.weights") - - guessed = resolve_weight_binding( - weights_dir=str(tmp_path), - source_name="model.embed_tokens.weight", - ) - assert guessed is None - - -def test_resolve_weight_binding_reads_convert_manifest_aliases(tmp_path: Path) -> None: - (tmp_path / "decoder_attn_q.cq4.weights").write_bytes(b"") - (tmp_path / "weights_manifest.json").write_text( - json.dumps( - { - "weights": [ - { - "source_name": "model.layers.0.self_attn.q_proj.weight", - "hf_name": "model.layers.0.self_attn.q_proj.weight", - "adapter_name": "backbone.layers.0.self_attn.q_proj.weight", - "source_names": ["adapter.backbone.layers.0.self_attn.q_proj.weight"], - "output_name": "decoder_attn_q.cq4.weights", - "component": "language", - } - ] - } - ) - ) - - binding = resolve_weight_binding( - weights_dir=str(tmp_path), - source_name="adapter.backbone.layers.0.self_attn.q_proj.weight", - ) - - assert binding is not None - assert binding.path.endswith("decoder_attn_q.cq4.weights") - - -def test_resolve_weight_binding_expands_manifest_wrapper_prefixes(tmp_path: Path) -> None: - (tmp_path / "layer_0_q.cq4.weights").write_bytes(b"") - (tmp_path / "weights_manifest.json").write_text( - json.dumps( - { - "model.layers.0.self_attn.q_proj.weight": { - "filename": "layer_0_q.cq4.weights", - "kind": "weight", - } - } - ) - ) - - binding = resolve_weight_binding( - weights_dir=str(tmp_path), - source_name="adapter.model.model.layers.0.self_attn.q_proj.weight", - ) - - assert binding is not None - assert binding.path.endswith("layer_0_q.cq4.weights") - - -def test_resolve_weight_binding_expands_component_local_manifest_aliases(tmp_path: Path) -> None: - (tmp_path / "encoder_layer_0_q.cq4.weights").write_bytes(b"") - (tmp_path / "weights_manifest.json").write_text( - json.dumps( - { - "encoder.layers.0.self_attn.q_proj.weight": { - "filename": "encoder_layer_0_q.cq4.weights", - "kind": "weight", - } - } - ) - ) - - binding = resolve_weight_binding( - weights_dir=str(tmp_path), - source_name="module.layers.0.self_attn.q_proj.weight", - ) - - assert binding is not None - assert binding.path.endswith("encoder_layer_0_q.cq4.weights") - - -def test_resolve_weight_binding_expands_language_model_backbone_aliases(tmp_path: Path) -> None: - (tmp_path / "layer_0_scalar.weights").write_bytes(b"") - (tmp_path / "weights_manifest.json").write_text( - json.dumps( - { - "model.language_model.layers.0.layer_scalar": { - "filename": "layer_0_scalar.weights", - "kind": "weight", - } - } - ) - ) - - binding = resolve_weight_binding( - weights_dir=str(tmp_path), - source_name="module.backbone.layers.0.layer_scalar", - ) - - assert binding is not None - assert binding.path.endswith("layer_0_scalar.weights") - - -def test_resolve_weight_binding_maps_module_backbone_to_model_manifest(tmp_path: Path) -> None: - (tmp_path / "layer_0_q.weights").write_bytes(b"") - (tmp_path / "weights_manifest.json").write_text( - json.dumps( - { - "weights": [ - { - "source_name": "model.layers.0.self_attn.q_proj.weight", - "hf_name": "model.layers.0.self_attn.q_proj.weight", - "adapter_name": "model.layers.0.self_attn.q_proj.weight", - "output_name": "layer_0_q.weights", - "component": "language", - } - ] - } - ) - ) - - for source_name in ( - "module.backbone.layers.0.self_attn.q_proj.weight", - "adapter.backbone.layers.0.self_attn.q_proj.weight", - ): - binding = resolve_weight_binding( - weights_dir=str(tmp_path), - source_name=source_name, - ) - - assert binding is not None, source_name - assert binding.path.endswith("layer_0_q.weights"), source_name - - -def test_resolve_weight_binding_maps_tied_lm_head_to_token_embedding_manifest(tmp_path: Path) -> None: - (tmp_path / "token_embeddings.weights").write_bytes(b"") - (tmp_path / "weights_manifest.json").write_text( - json.dumps( - { - "weights": [ - { - "source_name": "model.language_model.embed_tokens.weight", - "hf_name": "model.language_model.embed_tokens.weight", - "adapter_name": "model.language_model.embed_tokens.weight", - "source_names": ["model.language_model.embed_tokens.weight"], - "output_name": "token_embeddings.weights", - "component": "embedding", - } - ] - } - ) - ) - - for source_name in ("lm_head.weight", "model.lm_head.weight"): - binding = resolve_weight_binding( - weights_dir=str(tmp_path), - source_name=source_name, - ) - - assert binding is not None - assert binding.path.endswith("token_embeddings.weights") - assert binding.kind == "embedding" - - -def test_resolve_weight_binding_expands_nested_model_tower_aliases(tmp_path: Path) -> None: - (tmp_path / "vision_q.weights").write_bytes(b"") - (tmp_path / "weights_manifest.json").write_text( - json.dumps( - { - "model.vision_tower.encoder.layers.0.self_attn.q_proj.weight": { - "filename": "vision_q.weights", - "kind": "weight", - } - } - ) - ) - - binding = resolve_weight_binding( - weights_dir=str(tmp_path), - source_name="module.model.model.vision_tower.encoder.layers.0.self_attn.q_proj.weight", - ) - - assert binding is not None - assert binding.path.endswith("vision_q.weights") - - -def test_rank5_materialized_constant_embeds_in_cactus_graph(tmp_path: Path) -> None: - from cactus.transpile.runtime_compat import Graph - from cactus.transpile.runtime_compat import Tensor - - tensor = np.zeros((1, 2, 3, 4, 5), dtype=np.float16) - graph_path = tmp_path / "rank5.cactus" - - graph = Graph() - constant = graph.input(tensor.shape, Graph.FP16) - graph.set_input(constant, tensor) - graph.mark_embedded_input(constant) - graph.save(graph_path) - - loaded = Graph.load(graph_path) - loaded_constant = Tensor(loaded, constant.id, constant.shape, constant.dtype) - np.testing.assert_array_equal(loaded_constant.numpy(), tensor) - assert not list(tmp_path.glob("*.weights")) diff --git a/python/tests/transpile/__init__.py b/python/tests/transpile/__init__.py deleted file mode 100644 index c0dd68e2d..000000000 --- a/python/tests/transpile/__init__.py +++ /dev/null @@ -1 +0,0 @@ -"""Transpiler regression tests and developer tools.""" diff --git a/python/tests/transpile/tools/__init__.py b/python/tests/transpile/tools/__init__.py deleted file mode 100644 index b58eb5601..000000000 --- a/python/tests/transpile/tools/__init__.py +++ /dev/null @@ -1 +0,0 @@ -"""Developer-only transpiler comparison tools.""" diff --git a/python/tests/transpile/tools/compare_gemma4_multimodal.py b/python/tests/transpile/tools/compare_gemma4_multimodal.py deleted file mode 100644 index 14bee44e5..000000000 --- a/python/tests/transpile/tools/compare_gemma4_multimodal.py +++ /dev/null @@ -1,932 +0,0 @@ -from __future__ import annotations - -import argparse -import copy -import gc -import json -import os -import sys -import time -from dataclasses import asdict -from dataclasses import dataclass -from pathlib import Path -from typing import Any - -import numpy as np -import torch - -TOOLS_DIR = Path(__file__).resolve().parent -PYTHON_ROOT = TOOLS_DIR.parents[2] -PROJECT_ROOT = PYTHON_ROOT.parent - -sys.path.insert(0, str(PYTHON_ROOT)) - -from cactus.bindings.cactus import cactus_complete -from cactus.bindings.cactus import cactus_destroy -from cactus.bindings.cactus import cactus_init -from cactus.bindings.cactus import cactus_reset -from cactus.transpile.capture_pytorch import capture_model -from cactus.transpile.canonicalize.cleanup import canonicalize_exported_graph -from cactus.transpile.model_adapters import canonicalize_model_interface -from cactus.transpile.optimize_graph import FusionConfig -from cactus.transpile.optimize_graph import optimize_graph -from cactus.transpile.hf_model import TranspileWrapper -from cactus.transpile.hf_model import _graph_to_dict -from cactus.transpile.hf_model import _load_transformers_bundle -from cactus.transpile.hf_model import _lower_preoptimized_ir -from cactus.transpile.hf_model import _parse_dtype -from cactus.transpile.hf_model import _prepare_gemma4_multimodal_inputs -from cactus.transpile.hf_model import _serialize_json_compatible -from cactus.transpile.hf_model import _validate_weights_dir -from cactus.transpile.hf_model import _write_json - - -_DEFAULT_STOP_SEQUENCES = ("", "", "", "<|im_end|>") - - -@dataclass -class NumericSummary: - mean: float - min: float - max: float - median: float - - -@dataclass -class HandwrittenRun: - wall_ms: float - total_time_ms: float - time_to_first_token_ms: float - prefill_tps: float - decode_tps: float - decode_tokens: int - total_tokens: int - response: str - response_json: dict[str, Any] - - -@dataclass -class TranspiledRun: - preprocess_ms: float - execute_ms: float - total_ms: float - time_to_first_token_ms: float - prefill_tps: float - decode_tps: float - decode_tokens: int - total_tokens: int - response: str - stop_reason: str - generated_token_ids: list[int] - input_shapes: dict[str, list[int]] - - -def _numeric_summary(values: list[float]) -> NumericSummary: - arr = np.asarray(values, dtype=np.float64) - return NumericSummary( - mean=float(np.mean(arr)), - min=float(np.min(arr)), - max=float(np.max(arr)), - median=float(np.median(arr)), - ) - - -def _roundtrip_jsonable(value: Any) -> Any: - return json.loads(json.dumps(_serialize_json_compatible(value))) - - -def _resolve_default_weights_dir(model_id: str) -> Path | None: - env_path = os.environ.get("CACTUS_TEST_GEMMA4_MODEL") - if env_path: - candidate = Path(env_path).resolve() - if (candidate / "config.txt").exists(): - return candidate - - candidates = ( - PROJECT_ROOT / "weights" / "gemma4_int4", - PROJECT_ROOT / "weights" / model_id.split("/")[-1].lower(), - PROJECT_ROOT / "weights" / model_id.split("/")[-1].lower().replace("_", "-"), - ) - for candidate in candidates: - if (candidate / "config.txt").exists(): - return candidate - return None - - -def _require_processor_tokenizer(processor_or_tokenizer: object) -> object: - tokenizer_like = getattr(processor_or_tokenizer, "tokenizer", processor_or_tokenizer) - if not hasattr(tokenizer_like, "decode") or not hasattr(tokenizer_like, "encode"): - raise RuntimeError( - "Gemma4 compare script requires a tokenizer-like object with encode/decode support." - ) - return tokenizer_like - - -def _build_messages_payload( - *, - prompt: str, - image_files: tuple[str, ...], - audio_file: str | None, - system_prompt: str = "", -) -> list[dict[str, object]]: - messages: list[dict[str, object]] = [] - normalized_system = system_prompt.strip() - if normalized_system: - messages.append({"role": "system", "content": normalized_system}) - - user_message: dict[str, object] = { - "role": "user", - "content": prompt, - } - if image_files: - user_message["images"] = [str(Path(path).resolve()) for path in image_files] - if audio_file: - user_message["audio"] = [str(Path(audio_file).resolve())] - messages.append(user_message) - return messages - - -def _encode_stop_sequences(tokenizer: object, stop_sequences: tuple[str, ...]) -> list[list[int]]: - encoded: list[list[int]] = [] - encode = getattr(tokenizer, "encode", None) - if not callable(encode): - return encoded - for stop_sequence in stop_sequences: - try: - token_ids = list(encode(stop_sequence, add_special_tokens=False)) - except TypeError: - token_ids = list(encode(stop_sequence)) - if token_ids: - encoded.append([int(token_id) for token_id in token_ids]) - return encoded - - -def _has_suffix(token_ids: list[int], suffix: list[int]) -> bool: - if not suffix or len(token_ids) < len(suffix): - return False - return token_ids[-len(suffix) :] == suffix - - -def _trim_stop_suffix(token_ids: list[int], stop_sequences: list[list[int]]) -> bool: - for stop_sequence in stop_sequences: - if _has_suffix(token_ids, stop_sequence): - del token_ids[-len(stop_sequence) :] - return True - return False - - -def _decode_generated_text(tokenizer: object, token_ids: list[int], *, skip_special_tokens: bool) -> str: - decode = getattr(tokenizer, "decode", None) - if not callable(decode): - raise RuntimeError(f"tokenizer does not expose decode(): {type(tokenizer).__name__}") - try: - return str(decode(token_ids, skip_special_tokens=skip_special_tokens)) - except TypeError: - return str(decode(token_ids)) - - -def _summarize_handwritten_runs( - runs: list[HandwrittenRun], - *, - weights_dir: Path, -) -> dict[str, Any]: - if not runs: - raise ValueError("expected at least one handwritten run") - return { - "weights_dir": str(weights_dir), - "latest_response": runs[-1].response, - "total_time_ms": asdict(_numeric_summary([run.total_time_ms for run in runs])), - "time_to_first_token_ms": asdict(_numeric_summary([run.time_to_first_token_ms for run in runs])), - "prefill_tps": asdict(_numeric_summary([run.prefill_tps for run in runs])), - "decode_tps": asdict(_numeric_summary([run.decode_tps for run in runs])), - "decode_tokens": asdict(_numeric_summary([float(run.decode_tokens) for run in runs])), - "total_tokens": asdict(_numeric_summary([float(run.total_tokens) for run in runs])), - "wall_ms": asdict(_numeric_summary([run.wall_ms for run in runs])), - "runs": [_roundtrip_jsonable(asdict(run)) for run in runs], - } - - -def _summarize_transpiled_runs( - runs: list[TranspiledRun], - *, - compile_time_ms: float, - model_source: str, - weight_bindings: int, - prepared_input_shapes: dict[str, list[int]], -) -> dict[str, Any]: - if not runs: - raise ValueError("expected at least one transpiled run") - return { - "compile_time_ms": compile_time_ms, - "model_source": model_source, - "weight_bindings": weight_bindings, - "prepared_input_shapes": prepared_input_shapes, - "latest_response": runs[-1].response, - "preprocess_ms": asdict(_numeric_summary([run.preprocess_ms for run in runs])), - "execute_ms": asdict(_numeric_summary([run.execute_ms for run in runs])), - "total_ms": asdict(_numeric_summary([run.total_ms for run in runs])), - "time_to_first_token_ms": asdict(_numeric_summary([run.time_to_first_token_ms for run in runs])), - "prefill_tps": asdict(_numeric_summary([run.prefill_tps for run in runs])), - "decode_tps": asdict(_numeric_summary([run.decode_tps for run in runs])), - "decode_tokens": asdict(_numeric_summary([float(run.decode_tokens) for run in runs])), - "total_tokens": asdict(_numeric_summary([float(run.total_tokens) for run in runs])), - "runs": [_roundtrip_jsonable(asdict(run)) for run in runs], - } - - -def _prepare_repeat_inputs( - processor: object, - *, - tokenizer: object, - prompt: str, - image_files: tuple[str, ...], - audio_file: str | None, - torch_dtype: torch.dtype, - system_prompt: str, - enable_thinking_if_supported: bool, - max_new_tokens: int, -) -> tuple[object, int, float]: - start = time.perf_counter() - prepared = _prepare_gemma4_multimodal_inputs( - processor, - prompt=prompt, - image_files=image_files, - audio_file=audio_file, - torch_dtype=torch_dtype, - system_prompt=system_prompt, - enable_thinking_if_supported=enable_thinking_if_supported, - use_gemma4_chat_template=True, - ) - pad_token_id = getattr(tokenizer, "pad_token_id", None) - if pad_token_id is None: - pad_token_id = 0 - prompt_tokens = 0 - for name, tensor in zip(prepared.names, prepared.tensors): - if name == "input_ids": - prompt_tokens = int(tensor.shape[1]) - break - if prompt_tokens <= 0: - raise RuntimeError("Gemma4 multimodal prepared inputs did not include a non-empty input_ids tensor") - target_tokens = prompt_tokens + int(max_new_tokens) - - input_tensors: list[torch.Tensor] = [] - for name, tensor in zip(prepared.names, prepared.tensors): - if name in {"input_ids", "attention_mask", "token_type_ids"}: - padded_shape = list(tensor.shape) - padded_shape[1] = target_tokens - if name == "input_ids": - padded = torch.full( - padded_shape, - int(pad_token_id), - dtype=tensor.dtype, - device=tensor.device, - ) - else: - padded = torch.zeros( - padded_shape, - dtype=tensor.dtype, - device=tensor.device, - ) - padded[:, : tensor.shape[1]] = tensor - input_tensors.append(padded) - else: - input_tensors.append(tensor) - metadata = dict(prepared.metadata) - metadata["prompt_token_count"] = prompt_tokens - metadata["target_token_count"] = int(target_tokens) - metadata["input_shapes"] = { - name: list(tensor.shape) - for name, tensor in zip(prepared.names, input_tensors) - } - prepared = type(prepared)( - names=prepared.names, - tensors=tuple(input_tensors), - metadata=metadata, - ) - preprocess_ms = (time.perf_counter() - start) * 1000.0 - return prepared, prompt_tokens, preprocess_ms - - -def _run_handwritten_once( - *, - model_handle: Any, - messages_json: str, - options_json: str, -) -> HandwrittenRun: - cactus_reset(model_handle) - streamed_token_ids: list[int] = [] - - def _collect(_text: str, token_id: int) -> None: - streamed_token_ids.append(int(token_id)) - - start = time.perf_counter() - response_json_text = cactus_complete( - model_handle, - messages_json, - options_json, - "[]", - _collect, - pcm_data=None, - ) - wall_ms = (time.perf_counter() - start) * 1000.0 - response_payload = json.loads(response_json_text) - return HandwrittenRun( - wall_ms=wall_ms, - total_time_ms=float(response_payload.get("total_time_ms", wall_ms)), - time_to_first_token_ms=float(response_payload.get("time_to_first_token_ms", 0.0)), - prefill_tps=float(response_payload.get("prefill_tps", 0.0)), - decode_tps=float(response_payload.get("decode_tps", 0.0)), - decode_tokens=len(streamed_token_ids), - total_tokens=int(response_payload.get("total_tokens", len(streamed_token_ids))), - response=str(response_payload.get("response", "")), - response_json=response_payload, - ) - - -def _run_transpiled_generation_once( - *, - tg, - processor: object, - tokenizer: object, - prompt: str, - image_files: tuple[str, ...], - audio_file: str | None, - torch_dtype: torch.dtype, - system_prompt: str, - enable_thinking_if_supported: bool, - max_new_tokens: int, - stop_sequences: tuple[str, ...], - expected_input_shapes: dict[str, list[int]] | None = None, - prepared_override: object | None = None, - progress_label: str = "transpiled", - progress_every_tokens: int = 8, -) -> TranspiledRun: - if prepared_override is not None: - prepared = type(prepared_override)( - names=prepared_override.names, - tensors=tuple(tensor.clone() for tensor in prepared_override.tensors), - metadata=dict(prepared_override.metadata), - ) - prompt_tokens = int(prepared.metadata.get("prompt_token_count", 0)) - preprocess_ms = 0.0 - else: - prepared, prompt_tokens, preprocess_ms = _prepare_repeat_inputs( - processor, - tokenizer=tokenizer, - prompt=prompt, - image_files=image_files, - audio_file=audio_file, - torch_dtype=torch_dtype, - system_prompt=system_prompt, - enable_thinking_if_supported=enable_thinking_if_supported, - max_new_tokens=max_new_tokens, - ) - inputs_by_name = { - name: tensor.detach().cpu().numpy() - for name, tensor in zip(prepared.names, prepared.tensors) - } - if expected_input_shapes is not None: - for name, tensor in zip(prepared.names, prepared.tensors): - expected_shape = expected_input_shapes.get(name) - if expected_shape is None: - continue - actual_shape = list(tensor.shape) - if actual_shape != list(expected_shape): - raise RuntimeError( - f"runtime prepared input shape mismatch for {name}: " - f"compiled with {list(expected_shape)}, got {actual_shape}" - ) - if "input_ids" not in inputs_by_name: - raise RuntimeError("transpiled Gemma4 generation requires input_ids") - - target_tokens = int(prepared.metadata["target_token_count"]) - - eos_token_id = getattr(tokenizer, "eos_token_id", None) - - padded_inputs: dict[str, np.ndarray] = dict(inputs_by_name) - generated_ids: list[int] = [] - encoded_stop_sequences = _encode_stop_sequences(tokenizer, stop_sequences) - current_length = prompt_tokens - first_token_ms = 0.0 - execute_start = time.perf_counter() - stop_reason = "max_new_tokens" - progress_interval = max(1, int(progress_every_tokens)) - - for step_index in range(max_new_tokens): - runtime_inputs = [padded_inputs[name] for name in prepared.names] - if step_index == 0: - print("transpiled_step0_set_inputs_begin=true", flush=True) - tg.set_inputs(runtime_inputs) - if step_index == 0: - print("transpiled_step0_set_inputs_done=true", flush=True) - print("transpiled_step0_execute_begin=true", flush=True) - logits = tg.execute()[0].numpy() - if step_index == 0: - print("transpiled_step0_execute_done=true", flush=True) - if logits.ndim == 3: - token_position = current_length - 1 - if logits.shape[1] == 1: - token_position = 0 - next_token_id = int(np.argmax(logits[0, token_position])) - elif logits.ndim == 2: - next_token_id = int(np.argmax(logits[0])) - else: - raise RuntimeError(f"unexpected transpiled logits rank: {logits.ndim}") - generated_ids.append(next_token_id) - - if step_index == 0: - first_token_ms = preprocess_ms + (time.perf_counter() - execute_start) * 1000.0 - - if eos_token_id is not None and next_token_id == int(eos_token_id): - stop_reason = "eos_token" - break - if _trim_stop_suffix(generated_ids, encoded_stop_sequences): - stop_reason = "stop_sequence" - break - - if current_length >= target_tokens: - stop_reason = "context_limit" - break - - padded_inputs["input_ids"][0, current_length] = next_token_id - padded_inputs["attention_mask"][0, current_length] = 1 - padded_inputs["token_type_ids"][0, current_length] = 0 - if step_index == 0: - print("transpiled_step0_state_update_done=true", flush=True) - current_length += 1 - generated_count = len(generated_ids) - if generated_count % progress_interval == 0 or generated_count == max_new_tokens: - elapsed_ms = (time.perf_counter() - execute_start) * 1000.0 - print( - f"{progress_label}_decode_progress={generated_count}/{max_new_tokens} " - f"elapsed_ms={elapsed_ms:.3f}", - flush=True, - ) - - execute_ms = (time.perf_counter() - execute_start) * 1000.0 - response = _decode_generated_text(tokenizer, generated_ids, skip_special_tokens=True).strip() - if not response: - response = _decode_generated_text(tokenizer, generated_ids, skip_special_tokens=False).strip() - decode_time_ms = max(0.0, execute_ms + preprocess_ms - first_token_ms) - decode_tps = ( - ((len(generated_ids) - 1) * 1000.0) / decode_time_ms - if len(generated_ids) > 1 and decode_time_ms > 0.0 - else 0.0 - ) - prefill_tps = (prompt_tokens * 1000.0) / first_token_ms if first_token_ms > 0.0 else 0.0 - print( - f"{progress_label}_decode_done=true tokens={len(generated_ids)} " - f"stop_reason={stop_reason} execute_ms={execute_ms:.3f}", - flush=True, - ) - - return TranspiledRun( - preprocess_ms=preprocess_ms, - execute_ms=execute_ms, - total_ms=preprocess_ms + execute_ms, - time_to_first_token_ms=first_token_ms, - prefill_tps=prefill_tps, - decode_tps=decode_tps, - decode_tokens=len(generated_ids), - total_tokens=prompt_tokens + len(generated_ids), - response=response, - stop_reason=stop_reason, - generated_token_ids=generated_ids, - input_shapes={ - name: list(tensor.shape) - for name, tensor in zip(prepared.names, prepared.tensors) - }, - ) - - -def _print_summary(summary: dict[str, Any]) -> None: - print(f"model_id={summary['model_id']}") - print(f"prompt={summary['prompt']!r}") - if summary.get("image_files"): - print(f"image_files={summary['image_files']}") - if summary.get("audio_file"): - print(f"audio_file={summary['audio_file']}") - - handwritten = summary.get("handwritten") - if handwritten is not None: - hw_total = handwritten["total_time_ms"]["mean"] - hw_ttft = handwritten["time_to_first_token_ms"]["mean"] - hw_tokens = handwritten["decode_tokens"]["mean"] - hw_tps = hw_tokens * 1000.0 / max(hw_total, 1e-9) - print("handwritten:") - print(f" total_ms mean={hw_total:.3f}") - print(f" time_to_first_token_ms mean={hw_ttft:.3f}") - print(f" decode_tokens mean={hw_tokens:.3f}") - print(f" output_tokens_per_sec={hw_tps:.3f}") - print(f" latest_response={handwritten['latest_response']!r}") - - transpiled = summary.get("transpiled") - if transpiled is not None: - tr_total = transpiled["total_ms"]["mean"] - tr_ttft = transpiled["time_to_first_token_ms"]["mean"] - tr_tokens = transpiled["decode_tokens"]["mean"] - tr_tps = tr_tokens * 1000.0 / max(tr_total, 1e-9) - print("transpiled:") - print(f" compile_time_ms={transpiled['compile_time_ms']:.3f}") - print(f" total_ms mean={tr_total:.3f}") - print(f" time_to_first_token_ms mean={tr_ttft:.3f}") - print(f" decode_tokens mean={tr_tokens:.3f}") - print(f" output_tokens_per_sec={tr_tps:.3f}") - print(f" weight_bindings={transpiled['weight_bindings']}") - print(f" latest_response={transpiled['latest_response']!r}") - - comparison = summary.get("comparison") - if comparison is not None: - print("comparison:") - print( - " end_to_end_speedup_vs_handwritten=" - f"{comparison['end_to_end_speedup_vs_handwritten']:.6f}" - ) - print(f" responses_match={comparison['responses_match']}") - - -def main() -> int: - parser = argparse.ArgumentParser( - description=( - "Compare handwritten Cactus Gemma4 multimodal generation against the " - "transpiled Gemma4 multimodal graph on the same prompt, image, and audio." - ) - ) - parser.add_argument("--model-id", default="google/gemma-4-E2B-it") - parser.add_argument("--prompt", required=True) - parser.add_argument("--system-prompt", default="") - parser.add_argument("--image-file", action="append", default=[]) - parser.add_argument("--audio-file", default="") - parser.add_argument( - "--weights-dir", - default="", - help="Converted Cactus weights directory. Required for handwritten comparison and reused for transpiled mmap bindings.", - ) - parser.add_argument("--artifact-dir", default="", help="Optional directory for summary and transpile artifacts.") - parser.add_argument("--graph-filename", default="graph.cactus") - parser.add_argument("--warmup", type=int, default=1) - parser.add_argument("--repeats", type=int, default=1) - parser.add_argument("--max-new-tokens", type=int, default=96) - parser.add_argument( - "--transpiled-progress-every-tokens", - type=int, - default=1, - help=( - "Emit transpiled decode progress every N generated tokens. " - "Use 1 for per-token visibility on long CPU-bound Gemma4 runs." - ), - ) - parser.add_argument("--torch-dtype", default="float16") - parser.add_argument("--token", default=os.environ.get("HF_TOKEN")) - parser.add_argument("--trust-remote-code", action="store_true") - parser.add_argument("--local-files-only", action="store_true") - parser.add_argument("--skip-handwritten", action="store_true") - parser.add_argument("--skip-transpiled", action="store_true") - parser.add_argument("--enable-thinking", action="store_true") - parser.add_argument("--temperature", type=float, default=0.0) - parser.add_argument("--top-p", type=float, default=1.0) - parser.add_argument("--top-k", type=int, default=1) - parser.add_argument("--stop-sequence", action="append", default=[]) - parser.add_argument("--no-fuse-gated-deltanet", action="store_true") - parser.add_argument("--no-fuse-rms-norm", action="store_true") - parser.add_argument("--no-fuse-rope", action="store_true") - parser.add_argument("--no-fuse-attention", action="store_true") - parser.add_argument("--no-fuse-attention-block", action="store_true") - parser.add_argument("--no-fuse-add-clipped", action="store_true") - args = parser.parse_args() - - if args.skip_handwritten and args.skip_transpiled: - raise RuntimeError("cannot skip both handwritten and transpiled runs") - - image_files = tuple(str(Path(path).resolve()) for path in args.image_file if str(path).strip()) - audio_file = str(Path(args.audio_file).resolve()) if args.audio_file.strip() else None - if audio_file and not Path(audio_file).exists(): - raise FileNotFoundError(f"audio file does not exist: {audio_file}") - for image_file in image_files: - if not Path(image_file).exists(): - raise FileNotFoundError(f"image file does not exist: {image_file}") - - torch_dtype = _parse_dtype(args.torch_dtype) - artifact_dir = Path(args.artifact_dir).resolve() if args.artifact_dir else None - if artifact_dir is not None: - artifact_dir.mkdir(parents=True, exist_ok=True) - - weights_dir_path = None - if args.weights_dir.strip(): - weights_dir_path = _validate_weights_dir(args.weights_dir.strip(), model_id=args.model_id) - else: - weights_dir_path = _resolve_default_weights_dir(args.model_id) - - if not args.skip_handwritten and weights_dir_path is None: - raise RuntimeError( - "handwritten Gemma4 comparison requires a converted weights directory.\n" - "\n" - "Create one with something like:\n" - f" cactus convert {args.model_id} /path/to/gemma4_weights\n" - "\n" - "Then rerun with --weights-dir /path/to/gemma4_weights, or rerun with --skip-handwritten." - ) - - stop_sequences = tuple(args.stop_sequence) if args.stop_sequence else _DEFAULT_STOP_SEQUENCES - fusion_config = FusionConfig( - enable_gated_deltanet=not args.no_fuse_gated_deltanet, - enable_rms_norm=not args.no_fuse_rms_norm, - enable_rope=not args.no_fuse_rope, - enable_attention=not args.no_fuse_attention, - enable_attention_block=not args.no_fuse_attention_block, - enable_add_clipped=not args.no_fuse_add_clipped, - ) - - summary: dict[str, Any] = { - "model_id": args.model_id, - "prompt": args.prompt, - "system_prompt": args.system_prompt, - "image_files": list(image_files), - "audio_file": audio_file, - "weights_dir": str(weights_dir_path) if weights_dir_path is not None else None, - "warmup": args.warmup, - "repeats": args.repeats, - "max_new_tokens": args.max_new_tokens, - "transpiled_progress_every_tokens": int(args.transpiled_progress_every_tokens), - "torch_dtype": args.torch_dtype, - "trust_remote_code": bool(args.trust_remote_code), - "local_files_only": bool(args.local_files_only), - "enable_thinking": bool(args.enable_thinking), - "temperature": float(args.temperature), - "top_p": float(args.top_p), - "top_k": int(args.top_k), - "stop_sequences": list(stop_sequences), - } - - messages_payload = _build_messages_payload( - prompt=args.prompt, - image_files=image_files, - audio_file=audio_file, - system_prompt=args.system_prompt, - ) - messages_json = json.dumps(messages_payload) - options_json = json.dumps( - { - "temperature": float(args.temperature), - "top_p": float(args.top_p), - "top_k": int(args.top_k), - "max_tokens": int(args.max_new_tokens), - "stop_sequences": list(stop_sequences), - "enable_thinking_if_supported": bool(args.enable_thinking), - "auto_handoff": False, - "telemetry_enabled": False, - } - ) - - transpiled_runs: list[TranspiledRun] = [] - if not args.skip_transpiled: - model_source, processor_or_tokenizer, model, _ = _load_transformers_bundle( - model_id=args.model_id, - task="multimodal_causal_lm_logits", - torch_dtype=torch_dtype, - token=args.token, - trust_remote_code=args.trust_remote_code, - local_files_only=args.local_files_only, - ) - processor = processor_or_tokenizer - tokenizer = _require_processor_tokenizer(processor_or_tokenizer) - initial_prepared, _, initial_prepare_ms = _prepare_repeat_inputs( - processor, - tokenizer=tokenizer, - prompt=args.prompt, - image_files=image_files, - audio_file=audio_file, - torch_dtype=torch_dtype, - system_prompt=args.system_prompt, - enable_thinking_if_supported=args.enable_thinking, - max_new_tokens=args.max_new_tokens, - ) - canonical = canonicalize_model_interface( - model, - task="multimodal_causal_lm_logits", - input_names=initial_prepared.names, - weights_dir=str(weights_dir_path) if weights_dir_path is not None else None, - ) - if hasattr(canonical.module, "last_token_logits_only"): - canonical.module.last_token_logits_only = True - prime_static_features = getattr(canonical.module, "prime_static_multimodal_features", None) - if callable(prime_static_features): - prime_static_features(*initial_prepared.tensors) - wrapper = TranspileWrapper( - canonical.module, - weights_dir=str(weights_dir_path) if weights_dir_path is not None else None, - ).eval() - - compile_start = time.perf_counter() - print("transpile_capture_begin=true", flush=True) - prepare_cpu_float32_capture = getattr(canonical.module, "prepare_cpu_float32_capture", None) - restore_cpu_float32_capture = getattr(canonical.module, "restore_cpu_float32_capture", None) - if callable(prepare_cpu_float32_capture): - prepare_cpu_float32_capture() - try: - captured = capture_model(wrapper, initial_prepared.tensors) - finally: - if callable(restore_cpu_float32_capture): - restore_cpu_float32_capture() - print("transpile_capture_done=true", flush=True) - raw_ir_graph = copy.deepcopy(captured.ir_graph) - print("transpile_canonicalize_begin=true", flush=True) - canonicalize_exported_graph(captured.ir_graph) - print("transpile_canonicalize_done=true", flush=True) - print("transpile_optimize_begin=true", flush=True) - optimize_graph(captured.ir_graph, config=fusion_config) - print("transpile_optimize_done=true", flush=True) - optimized_ir_graph = copy.deepcopy(captured.ir_graph) - print("transpile_lower_begin=true", flush=True) - tg = _lower_preoptimized_ir(captured.ir_graph) - print("transpile_lower_done=true", flush=True) - compile_time_ms = (time.perf_counter() - compile_start) * 1000.0 - - op_counts: dict[str, int] = {} - for node_id in optimized_ir_graph.order: - op = optimized_ir_graph.nodes[node_id].op - op_counts[op] = op_counts.get(op, 0) + 1 - weight_bindings = sum( - 1 - for value in optimized_ir_graph.values.values() - if isinstance(value.meta, dict) and isinstance(value.meta.get("path"), str) - ) - - summary["transpile_compile"] = { - "model_source": model_source, - "raw_ir_nodes": len(raw_ir_graph.order), - "optimized_ir_nodes": len(optimized_ir_graph.order), - "weight_bindings": weight_bindings, - "compile_time_ms": compile_time_ms, - "initial_prepare_ms": initial_prepare_ms, - "ops": op_counts, - } - prompt_token_count = int(initial_prepared.metadata.get("prompt_token_count", 0)) - target_token_count = int(initial_prepared.metadata.get("target_token_count", 0)) - total_transpiled_runs = int(args.warmup) + int(args.repeats) - print( - "transpiled_decode_plan=" - f"prompt_tokens:{prompt_token_count} " - f"target_tokens:{target_token_count} " - f"max_new_tokens:{int(args.max_new_tokens)} " - f"runs:{total_transpiled_runs} " - f"progress_every_tokens:{int(args.transpiled_progress_every_tokens)}", - flush=True, - ) - - if weights_dir_path is not None and weight_bindings == 0: - raise RuntimeError( - f"No weight bindings were resolved from {weights_dir_path}\n" - "\n" - "The converted weights folder exists, but none of the captured constants matched weights_manifest.json.\n" - "Re-convert Gemma4 with the current converter before benchmarking." - ) - - if artifact_dir is not None: - _write_json( - artifact_dir / "raw_ir.json", - { - "model_id": args.model_id, - "model_source": model_source, - "task": "multimodal_causal_lm_logits", - "inputs": _serialize_json_compatible(initial_prepared.metadata), - "graph": _graph_to_dict(raw_ir_graph), - }, - ) - _write_json( - artifact_dir / "optimized_ir.json", - { - "model_id": args.model_id, - "model_source": model_source, - "task": "multimodal_causal_lm_logits", - "inputs": _serialize_json_compatible(initial_prepared.metadata), - "graph": _graph_to_dict(optimized_ir_graph), - }, - ) - print(f"saved_raw_ir={artifact_dir / 'raw_ir.json'}") - print(f"saved_optimized_ir={artifact_dir / 'optimized_ir.json'}") - if weight_bindings > 0: - skip_reason = ( - "skipped saved_graph because graph serialization does not yet preserve " - "mmap-bound weights/embeddings required by this transpiled Gemma4 graph" - ) - summary["transpile_compile"]["saved_graph_skipped_reason"] = skip_reason - print(f"note={skip_reason}") - else: - graph_path = artifact_dir / args.graph_filename - tg.graph.save(graph_path) - print(f"saved_graph={graph_path}") - - for _ in range(args.warmup): - warmup_index = _ + 1 - print(f"transpiled_warmup_begin={warmup_index}/{args.warmup}", flush=True) - _run_transpiled_generation_once( - tg=tg, - processor=processor, - tokenizer=tokenizer, - prompt=args.prompt, - image_files=image_files, - audio_file=audio_file, - torch_dtype=torch_dtype, - system_prompt=args.system_prompt, - enable_thinking_if_supported=args.enable_thinking, - max_new_tokens=args.max_new_tokens, - stop_sequences=stop_sequences, - expected_input_shapes=dict(initial_prepared.metadata.get("input_shapes", {})), - prepared_override=initial_prepared, - progress_label=f"transpiled_warmup{warmup_index}", - progress_every_tokens=args.transpiled_progress_every_tokens, - ) - for _ in range(args.repeats): - repeat_index = _ + 1 - print(f"transpiled_repeat_begin={repeat_index}/{args.repeats}", flush=True) - transpiled_runs.append( - _run_transpiled_generation_once( - tg=tg, - processor=processor, - tokenizer=tokenizer, - prompt=args.prompt, - image_files=image_files, - audio_file=audio_file, - torch_dtype=torch_dtype, - system_prompt=args.system_prompt, - enable_thinking_if_supported=args.enable_thinking, - max_new_tokens=args.max_new_tokens, - stop_sequences=stop_sequences, - expected_input_shapes=dict(initial_prepared.metadata.get("input_shapes", {})), - prepared_override=initial_prepared, - progress_label=f"transpiled_repeat{repeat_index}", - progress_every_tokens=args.transpiled_progress_every_tokens, - ) - ) - - summary["transpiled"] = _summarize_transpiled_runs( - transpiled_runs, - compile_time_ms=compile_time_ms, - model_source=model_source, - weight_bindings=weight_bindings, - prepared_input_shapes=dict(initial_prepared.metadata.get("input_shapes", {})), - ) - - # Keep the HF/export capture stack alive until the transpiled graph has - # completed execution. Gemma4 teardown before first execute has been - # unstable in practice, even though the lowered graph is already built. - del model - del wrapper - del canonical - del captured - gc.collect() - - handwritten_runs: list[HandwrittenRun] = [] - if not args.skip_handwritten and weights_dir_path is not None: - model_handle = cactus_init(str(weights_dir_path), None, False) - try: - for _ in range(args.warmup): - warmup_index = _ + 1 - print(f"handwritten_warmup_begin={warmup_index}/{args.warmup}", flush=True) - _run_handwritten_once( - model_handle=model_handle, - messages_json=messages_json, - options_json=options_json, - ) - for _ in range(args.repeats): - repeat_index = _ + 1 - print(f"handwritten_repeat_begin={repeat_index}/{args.repeats}", flush=True) - handwritten_runs.append( - _run_handwritten_once( - model_handle=model_handle, - messages_json=messages_json, - options_json=options_json, - ) - ) - finally: - cactus_destroy(model_handle) - summary["handwritten"] = _summarize_handwritten_runs( - handwritten_runs, - weights_dir=weights_dir_path, - ) - - if handwritten_runs and transpiled_runs: - handwritten_latest = handwritten_runs[-1].response.strip() - transpiled_latest = transpiled_runs[-1].response.strip() - summary["comparison"] = { - "end_to_end_speedup_vs_handwritten": ( - summary["handwritten"]["total_time_ms"]["mean"] - / max(summary["transpiled"]["total_ms"]["mean"], 1e-9) - ), - "responses_match": handwritten_latest == transpiled_latest, - "handwritten_latest_response": handwritten_latest, - "transpiled_latest_response": transpiled_latest, - } - - if artifact_dir is not None: - summary_path = artifact_dir / "summary.json" - summary_path.write_text(json.dumps(_roundtrip_jsonable(summary), indent=2, sort_keys=True) + "\n") - print(f"saved_summary={summary_path}") - - _print_summary(summary) - return 0 - - -if __name__ == "__main__": - raise SystemExit(main()) diff --git a/python/tests/transpile/tools/compare_parakeet_performance.py b/python/tests/transpile/tools/compare_parakeet_performance.py deleted file mode 100644 index eeaf06066..000000000 --- a/python/tests/transpile/tools/compare_parakeet_performance.py +++ /dev/null @@ -1,792 +0,0 @@ -from __future__ import annotations - -import argparse -import json -import os -import sys -import time -from collections import defaultdict -from dataclasses import asdict -from dataclasses import dataclass -from pathlib import Path -from typing import Any - -import numpy as np -import torch -from scipy.io import wavfile - -TOOLS_DIR = Path(__file__).resolve().parent -PYTHON_ROOT = TOOLS_DIR.parents[2] -PROJECT_ROOT = PYTHON_ROOT.parent - -sys.path.insert(0, str(PYTHON_ROOT)) - -from cactus.bindings.cactus import cactus_destroy -from cactus.bindings.cactus import cactus_init -from cactus.bindings.cactus import cactus_transcribe -from cactus.transpile.capture_pytorch import capture_model -from cactus.transpile.canonicalize.cleanup import canonicalize_exported_graph -from cactus.transpile.model_adapters import canonicalize_model_interface -from cactus.transpile.optimize_graph import FusionConfig -from cactus.transpile.optimize_graph import optimize_graph -from cactus.transpile.hf_model import TranspileWrapper -from cactus.transpile.hf_model import _ctc_greedy_decode_token_ids -from cactus.transpile.hf_model import _decode_token_ids -from cactus.transpile.hf_model import _infer_task_from_config -from cactus.transpile.hf_model import _load_optional_json -from cactus.transpile.hf_model import _load_optional_tokenizer -from cactus.transpile.hf_model import _load_transformers_bundle -from cactus.transpile.hf_model import _lower_preoptimized_ir -from cactus.transpile.hf_model import _parse_dtype -from cactus.transpile.hf_model import _prepare_audio_inputs -from cactus.transpile.hf_model import _validate_weights_dir - - -@dataclass -class NumericSummary: - mean: float - min: float - max: float - median: float - - -@dataclass -class ProfileEntry: - op: str - time_ms: float - shape: str - raw_line: str - - -@dataclass -class ProfileSection: - entries: list[ProfileEntry] - total_ms: float - - -@dataclass -class HandwrittenRun: - wall_ms: float - total_time_ms: float - time_to_first_token_ms: float - prefill_tps: float - decode_tps: float - decode_tokens: int - prefill_tokens: int - total_tokens: int - confidence: float - transcript: str - response_json: dict[str, Any] - - -@dataclass -class TranspiledRun: - preprocess_ms: float - execute_ms: float - decode_ms: float - total_ms: float - graph_only_ms: float - output_frames: int - output_vocab: int - decode_tokens: int - transcript: str | None - input_shapes: dict[str, list[int]] - - -def _numeric_summary(values: list[float]) -> NumericSummary: - arr = np.asarray(values, dtype=np.float64) - return NumericSummary( - mean=float(np.mean(arr)), - min=float(np.min(arr)), - max=float(np.max(arr)), - median=float(np.median(arr)), - ) - - -def _resolve_default_weights_dir(model_id: str) -> Path | None: - candidate = PROJECT_ROOT / "weights" / model_id.split("/")[-1].lower() - if (candidate / "config.txt").exists(): - return candidate - return None - - -def _audio_duration_seconds(audio_file: str) -> float: - sample_rate, samples = wavfile.read(audio_file) - if sample_rate <= 0: - raise ValueError(f"invalid sample rate in {audio_file}: {sample_rate}") - return float(len(samples)) / float(sample_rate) - - -def _profile_env(profile_file: Path | None): - class _ProfileContext: - def __enter__(self_inner): - self_inner.old_profile = os.environ.get("CACTUS_PROFILE") - self_inner.old_profile_file = os.environ.get("CACTUS_PROFILE_FILE") - if profile_file is None: - os.environ.pop("CACTUS_PROFILE", None) - os.environ.pop("CACTUS_PROFILE_FILE", None) - return self_inner - profile_file.parent.mkdir(parents=True, exist_ok=True) - if profile_file.exists(): - profile_file.unlink() - os.environ.pop("CACTUS_PROFILE", None) - os.environ["CACTUS_PROFILE_FILE"] = str(profile_file) - return self_inner - - def __exit__(self_inner, exc_type, exc, tb): - if self_inner.old_profile is None: - os.environ.pop("CACTUS_PROFILE", None) - else: - os.environ["CACTUS_PROFILE"] = self_inner.old_profile - - if self_inner.old_profile_file is None: - os.environ.pop("CACTUS_PROFILE_FILE", None) - else: - os.environ["CACTUS_PROFILE_FILE"] = self_inner.old_profile_file - return False - - return _ProfileContext() - - -def _parse_profile_file(path: Path) -> list[ProfileSection]: - if not path.exists(): - return [] - - lines = path.read_text().splitlines() - sections: list[ProfileSection] = [] - idx = 0 - while idx < len(lines): - if lines[idx].strip() != "=== Graph Execution Profile ===": - idx += 1 - continue - - entries: list[ProfileEntry] = [] - total_ms = 0.0 - idx += 1 - while idx < len(lines): - line = lines[idx].rstrip() - stripped = line.strip() - if stripped.startswith("Total execution time:"): - try: - total_ms = float(stripped.split(":", 1)[1].strip().split()[0]) - except Exception: - total_ms = 0.0 - elif stripped == "================================": - break - elif ( - stripped - and not stripped.startswith("Operation") - and not set(stripped).issubset({"-"}) - and not stripped.startswith("Total execution time:") - ): - parts = stripped.split() - if len(parts) >= 3: - op = parts[0] - try: - time_ms = float(parts[1]) - except ValueError: - time_ms = 0.0 - shape = parts[2] if parts[2].startswith("[") else "" - entries.append(ProfileEntry(op=op, time_ms=time_ms, shape=shape, raw_line=stripped)) - idx += 1 - - sections.append(ProfileSection(entries=entries, total_ms=total_ms)) - idx += 1 - - return sections - - -def _summarize_profile_sections(sections: list[ProfileSection]) -> dict[str, Any]: - op_total_ms: dict[str, float] = defaultdict(float) - op_count: dict[str, int] = defaultdict(int) - all_entries: list[ProfileEntry] = [] - total_ms_values: list[float] = [] - - for section in sections: - total_ms_values.append(section.total_ms) - all_entries.extend(section.entries) - for entry in section.entries: - op_total_ms[entry.op] += entry.time_ms - op_count[entry.op] += 1 - - by_op = [] - for op, total_ms in sorted(op_total_ms.items(), key=lambda item: item[1], reverse=True): - count = op_count[op] - by_op.append( - { - "op": op, - "count": count, - "total_ms": round(total_ms, 6), - "mean_ms": round(total_ms / max(count, 1), 6), - } - ) - - slowest_nodes = [ - { - "op": entry.op, - "time_ms": round(entry.time_ms, 6), - "shape": entry.shape, - "raw_line": entry.raw_line, - } - for entry in sorted(all_entries, key=lambda item: item.time_ms, reverse=True)[:50] - ] - - return { - "runs_profiled": len(sections), - "graph_total_ms": asdict(_numeric_summary(total_ms_values)) if total_ms_values else None, - "entries_per_run": [len(section.entries) for section in sections], - "by_op": by_op, - "slowest_nodes": slowest_nodes, - } - - -def _run_handwritten_once( - *, - model_handle: Any, - audio_file: str, -) -> HandwrittenRun: - options_json = json.dumps( - { - "use_vad": False, - "max_tokens": 4096, - } - ) - start = time.perf_counter() - response = cactus_transcribe( - model_handle, - audio_file, - "", - options_json, - None, - None, - ) - end = time.perf_counter() - - payload = json.loads(response) - return HandwrittenRun( - wall_ms=(end - start) * 1000.0, - total_time_ms=float(payload.get("total_time_ms", 0.0)), - time_to_first_token_ms=float(payload.get("time_to_first_token_ms", 0.0)), - prefill_tps=float(payload.get("prefill_tps", 0.0)), - decode_tps=float(payload.get("decode_tps", 0.0)), - decode_tokens=int(payload.get("decode_tokens", 0)), - prefill_tokens=int(payload.get("prefill_tokens", 0)), - total_tokens=int(payload.get("total_tokens", 0)), - confidence=float(payload.get("confidence", 0.0)), - transcript=str(payload.get("response", "")), - response_json=payload, - ) - - -def _run_transpiled_once( - *, - tg: Any, - processor: object | None, - input_names: tuple[str, ...], - model_config: dict[str, object], - preprocessor_config: dict[str, object], - model: torch.nn.Module, - torch_dtype: torch.dtype, - audio_file: str, - tokenizer: object | None, - blank_token_id: int | None, -) -> TranspiledRun: - start_pre = time.perf_counter() - prepared = _prepare_audio_inputs( - processor, - input_names=input_names, - config=model_config, - preprocessor_config=preprocessor_config, - model=model, - task="ctc_logits", - audio_file=audio_file, - torch_dtype=torch_dtype, - ) - end_pre = time.perf_counter() - - input_arrays = [tensor.detach().cpu().numpy() for tensor in prepared.tensors] - - start_exec = time.perf_counter() - tg.set_inputs(input_arrays) - outputs = tg.execute() - end_exec = time.perf_counter() - - start_decode = time.perf_counter() - logits = outputs[0].numpy().astype(np.float32) - token_ids = _ctc_greedy_decode_token_ids(logits, blank_token_id=blank_token_id) - transcript = _decode_token_ids(tokenizer, token_ids) if tokenizer is not None else None - end_decode = time.perf_counter() - - return TranspiledRun( - preprocess_ms=(end_pre - start_pre) * 1000.0, - execute_ms=(end_exec - start_exec) * 1000.0, - decode_ms=(end_decode - start_decode) * 1000.0, - total_ms=(end_decode - start_pre) * 1000.0, - graph_only_ms=(end_exec - start_exec) * 1000.0, - output_frames=int(logits.shape[1]), - output_vocab=int(logits.shape[2]), - decode_tokens=len(token_ids), - transcript=transcript, - input_shapes={ - name: list(tensor.shape) - for name, tensor in zip(prepared.names, prepared.tensors) - }, - ) - - -def _run_transpiled_graph_only_once( - *, - tg: Any, - cached_inputs: list[np.ndarray], -) -> float: - start = time.perf_counter() - tg.set_inputs(cached_inputs) - tg.execute() - end = time.perf_counter() - return (end - start) * 1000.0 - - -def _warmup(fn, count: int) -> None: - for _ in range(max(count, 0)): - fn() - - -def _roundtrip_jsonable(obj: Any) -> Any: - if isinstance(obj, dict): - return {str(key): _roundtrip_jsonable(value) for key, value in obj.items()} - if isinstance(obj, list): - return [_roundtrip_jsonable(value) for value in obj] - if isinstance(obj, tuple): - return [_roundtrip_jsonable(value) for value in obj] - if isinstance(obj, Path): - return str(obj) - if isinstance(obj, NumericSummary): - return asdict(obj) - if hasattr(obj, "__dataclass_fields__"): - return asdict(obj) - return obj - - -def _build_transpiled_graph( - *, - model_id: str, - audio_file: str, - weights_dir: str | None, - torch_dtype: torch.dtype, - token: str | None, - trust_remote_code: bool, - local_files_only: bool, - fusion_config: FusionConfig, -): - task = _infer_task_from_config(model_id) - if task != "ctc_logits": - raise RuntimeError(f"expected a CTC model, got task={task}") - - model_source, processor, model, model_config = _load_transformers_bundle( - model_id=model_id, - task=task, - torch_dtype=torch_dtype, - token=token, - trust_remote_code=trust_remote_code, - local_files_only=local_files_only, - ) - - preprocessor_config = _load_optional_json(model_source, "preprocessor_config.json") - if not preprocessor_config: - preprocessor_config = _load_optional_json(model_id, "preprocessor_config.json") - - tokenizer = _load_optional_tokenizer( - model_id=model_id, - model_source=model_source, - token=token, - trust_remote_code=trust_remote_code, - local_files_only=local_files_only, - ) - - canonical = canonicalize_model_interface(model, task=task) - prepared = _prepare_audio_inputs( - processor, - input_names=canonical.input_names, - config=model_config, - preprocessor_config=preprocessor_config, - model=model, - task=task, - audio_file=audio_file, - torch_dtype=torch_dtype, - ) - canonical = canonicalize_model_interface(model, task=task, input_names=prepared.names) - - wrapper = TranspileWrapper(canonical.module, weights_dir=weights_dir).eval() - - capture_started = time.perf_counter() - captured = capture_model(wrapper, prepared.tensors) - canonicalize_exported_graph(captured.ir_graph) - optimize_graph(captured.ir_graph, config=fusion_config) - tg = _lower_preoptimized_ir(captured.ir_graph) - capture_finished = time.perf_counter() - - cached_inputs = [tensor.detach().cpu().numpy() for tensor in prepared.tensors] - blank_token_id = getattr(getattr(model, "config", None), "pad_token_id", None) - - return { - "tg": tg, - "model": model, - "processor": processor, - "model_config": model_config, - "preprocessor_config": preprocessor_config, - "tokenizer": tokenizer, - "blank_token_id": int(blank_token_id) if blank_token_id is not None else None, - "compile_time_ms": (capture_finished - capture_started) * 1000.0, - "cached_inputs": cached_inputs, - "input_names": tuple(prepared.names), - "prepared_input_shapes": { - name: list(tensor.shape) - for name, tensor in zip(prepared.names, prepared.tensors) - }, - "model_source": model_source, - "weight_bindings": sum( - 1 - for value in captured.ir_graph.values.values() - if isinstance(value.meta, dict) and isinstance(value.meta.get("path"), str) - ), - } - - -def _benchmark_handwritten( - *, - weights_dir: Path, - audio_file: str, - warmup: int, - repeats: int, - profile_repeats: int, - artifact_dir: Path | None, -) -> dict[str, Any]: - model_handle = cactus_init(str(weights_dir), None, False) - try: - _warmup(lambda: _run_handwritten_once(model_handle=model_handle, audio_file=audio_file), warmup) - - runs: list[HandwrittenRun] = [] - for _ in range(max(repeats, 1)): - runs.append(_run_handwritten_once(model_handle=model_handle, audio_file=audio_file)) - - profile_path = ( - artifact_dir / "handwritten_profile.txt" - if artifact_dir is not None and profile_repeats > 0 - else None - ) - with _profile_env(profile_path): - for _ in range(max(profile_repeats, 0)): - _run_handwritten_once(model_handle=model_handle, audio_file=audio_file) - - profile_sections = _parse_profile_file(profile_path) if profile_path is not None else [] - profile_summary = _summarize_profile_sections(profile_sections) - - return { - "weights_dir": str(weights_dir), - "runs": [_roundtrip_jsonable(run) for run in runs], - "wall_ms": asdict(_numeric_summary([run.wall_ms for run in runs])), - "internal_total_time_ms": asdict(_numeric_summary([run.total_time_ms for run in runs])), - "time_to_first_token_ms": asdict(_numeric_summary([run.time_to_first_token_ms for run in runs])), - "decode_tps": asdict(_numeric_summary([run.decode_tps for run in runs])), - "prefill_tps": asdict(_numeric_summary([run.prefill_tps for run in runs])), - "decode_tokens": asdict(_numeric_summary([float(run.decode_tokens) for run in runs])), - "confidence": asdict(_numeric_summary([run.confidence for run in runs])), - "latest_transcript": runs[-1].transcript if runs else "", - "profile_file": str(profile_path) if profile_path is not None else None, - "profile_summary": profile_summary, - } - finally: - cactus_destroy(model_handle) - - -def _benchmark_transpiled( - *, - model_id: str, - audio_file: str, - weights_dir: str | None, - warmup: int, - repeats: int, - profile_repeats: int, - torch_dtype: torch.dtype, - token: str | None, - trust_remote_code: bool, - local_files_only: bool, - artifact_dir: Path | None, - fusion_config: FusionConfig, -) -> dict[str, Any]: - built = _build_transpiled_graph( - model_id=model_id, - audio_file=audio_file, - weights_dir=weights_dir, - torch_dtype=torch_dtype, - token=token, - trust_remote_code=trust_remote_code, - local_files_only=local_files_only, - fusion_config=fusion_config, - ) - - tg = built["tg"] - model = built["model"] - processor = built["processor"] - input_names = built["input_names"] - model_config = built["model_config"] - preprocessor_config = built["preprocessor_config"] - tokenizer = built["tokenizer"] - blank_token_id = built["blank_token_id"] - cached_inputs = built["cached_inputs"] - - _warmup( - lambda: _run_transpiled_once( - tg=tg, - processor=processor, - input_names=input_names, - model_config=model_config, - preprocessor_config=preprocessor_config, - model=model, - torch_dtype=torch_dtype, - audio_file=audio_file, - tokenizer=tokenizer, - blank_token_id=blank_token_id, - ), - warmup, - ) - _warmup(lambda: _run_transpiled_graph_only_once(tg=tg, cached_inputs=cached_inputs), warmup) - - runs: list[TranspiledRun] = [] - graph_only_ms: list[float] = [] - for _ in range(max(repeats, 1)): - runs.append( - _run_transpiled_once( - tg=tg, - processor=processor, - input_names=input_names, - model_config=model_config, - preprocessor_config=preprocessor_config, - model=model, - torch_dtype=torch_dtype, - audio_file=audio_file, - tokenizer=tokenizer, - blank_token_id=blank_token_id, - ) - ) - graph_only_ms.append(_run_transpiled_graph_only_once(tg=tg, cached_inputs=cached_inputs)) - - profile_path = ( - artifact_dir / "transpiled_profile.txt" - if artifact_dir is not None and profile_repeats > 0 - else None - ) - with _profile_env(profile_path): - for _ in range(max(profile_repeats, 0)): - _run_transpiled_graph_only_once(tg=tg, cached_inputs=cached_inputs) - - profile_sections = _parse_profile_file(profile_path) if profile_path is not None else [] - profile_summary = _summarize_profile_sections(profile_sections) - - return { - "model_source": built["model_source"], - "compile_time_ms": built["compile_time_ms"], - "weight_bindings": built["weight_bindings"], - "prepared_input_shapes": built["prepared_input_shapes"], - "runs": [_roundtrip_jsonable(run) for run in runs], - "preprocess_ms": asdict(_numeric_summary([run.preprocess_ms for run in runs])), - "execute_ms": asdict(_numeric_summary([run.execute_ms for run in runs])), - "decode_ms": asdict(_numeric_summary([run.decode_ms for run in runs])), - "total_ms": asdict(_numeric_summary([run.total_ms for run in runs])), - "graph_only_ms": asdict(_numeric_summary(graph_only_ms)), - "decode_tokens": asdict(_numeric_summary([float(run.decode_tokens) for run in runs])), - "output_frames": asdict(_numeric_summary([float(run.output_frames) for run in runs])), - "latest_transcript": runs[-1].transcript if runs else "", - "profile_file": str(profile_path) if profile_path is not None else None, - "profile_summary": profile_summary, - } - - -def _print_summary( - *, - audio_duration_sec: float, - handwritten: dict[str, Any] | None, - transpiled: dict[str, Any] | None, -) -> None: - print() - print(f"audio_duration_sec={audio_duration_sec:.3f}") - - if handwritten is not None: - hw_total = handwritten["internal_total_time_ms"]["mean"] - hw_graph = None - profile_total = handwritten.get("profile_summary", {}).get("graph_total_ms") - if isinstance(profile_total, dict): - hw_graph = profile_total.get("mean") - hw_tok = handwritten["decode_tokens"]["mean"] - hw_rtf = hw_total / max(audio_duration_sec * 1000.0, 1e-9) - hw_xrt = (audio_duration_sec * 1000.0) / max(hw_total, 1e-9) - hw_tok_s = hw_tok * 1000.0 / max(hw_total, 1e-9) - print("handwritten:") - print(f" total_ms mean={hw_total:.3f}") - print(f" graph_profile_ms mean={hw_graph:.3f}" if hw_graph is not None else " graph_profile_ms unavailable") - print(f" decode_tokens mean={hw_tok:.3f}") - print(f" output_tokens_per_sec={hw_tok_s:.3f}") - print(f" real_time_factor={hw_rtf:.3f}") - print(f" audio_x_realtime={hw_xrt:.3f}") - - if transpiled is not None: - tr_total = transpiled["total_ms"]["mean"] - tr_graph = transpiled["graph_only_ms"]["mean"] - tr_tok = transpiled["decode_tokens"]["mean"] - tr_rtf = tr_total / max(audio_duration_sec * 1000.0, 1e-9) - tr_xrt = (audio_duration_sec * 1000.0) / max(tr_total, 1e-9) - tr_tok_s = tr_tok * 1000.0 / max(tr_total, 1e-9) - tr_graph_tok_s = tr_tok * 1000.0 / max(tr_graph, 1e-9) - print("transpiled:") - print(f" total_ms mean={tr_total:.3f}") - print(f" graph_only_ms mean={tr_graph:.3f}") - print(f" decode_tokens mean={tr_tok:.3f}") - print(f" output_tokens_per_sec={tr_tok_s:.3f}") - print(f" graph_only_tokens_per_sec={tr_graph_tok_s:.3f}") - print(f" real_time_factor={tr_rtf:.3f}") - print(f" audio_x_realtime={tr_xrt:.3f}") - - -def main() -> int: - parser = argparse.ArgumentParser( - description=( - "Compare handwritten Cactus Parakeet performance against the transpiled " - "Parakeet graph on the same WAV input, including per-op graph profiles." - ) - ) - parser.add_argument("--model-id", default="nvidia/parakeet-ctc-1.1b") - parser.add_argument("--audio-file", required=True) - parser.add_argument( - "--weights-dir", - default="", - help="Converted Cactus weights directory. Reused for handwritten runtime and transpiled mmap bindings.", - ) - parser.add_argument("--artifact-dir", default="", help="Optional directory for JSON summaries and raw profiles.") - parser.add_argument("--warmup", type=int, default=1) - parser.add_argument("--repeats", type=int, default=5) - parser.add_argument("--profile-repeats", type=int, default=1) - parser.add_argument("--torch-dtype", default="float16") - parser.add_argument("--token", default=os.environ.get("HF_TOKEN")) - parser.add_argument("--trust-remote-code", action="store_true") - parser.add_argument("--local-files-only", action="store_true") - parser.add_argument("--skip-handwritten", action="store_true") - parser.add_argument("--skip-transpiled", action="store_true") - parser.add_argument("--no-fuse-gated-deltanet", action="store_true") - parser.add_argument("--no-fuse-rms-norm", action="store_true") - parser.add_argument("--no-fuse-rope", action="store_true") - parser.add_argument("--no-fuse-attention", action="store_true") - parser.add_argument("--no-fuse-attention-block", action="store_true") - parser.add_argument("--no-fuse-add-clipped", action="store_true") - args = parser.parse_args() - - audio_file = str(Path(args.audio_file).resolve()) - if not Path(audio_file).exists(): - raise FileNotFoundError(f"audio file does not exist: {audio_file}") - - torch_dtype = _parse_dtype(args.torch_dtype) - artifact_dir = Path(args.artifact_dir).resolve() if args.artifact_dir else None - if artifact_dir is not None: - artifact_dir.mkdir(parents=True, exist_ok=True) - - weights_dir_path = None - if args.weights_dir.strip(): - weights_dir_path = _validate_weights_dir(args.weights_dir.strip(), model_id=args.model_id) - else: - weights_dir_path = _resolve_default_weights_dir(args.model_id) - - if not args.skip_handwritten and weights_dir_path is None: - raise RuntimeError( - "handwritten comparison requires a converted weights directory.\n" - "\n" - f"Create one with:\n" - f" cactus convert {args.model_id} {PROJECT_ROOT / 'weights' / args.model_id.split('/')[-1].lower()}\n" - "\n" - "Or rerun with --skip-handwritten." - ) - - fusion_config = FusionConfig( - enable_gated_deltanet=not args.no_fuse_gated_deltanet, - enable_rms_norm=not args.no_fuse_rms_norm, - enable_rope=not args.no_fuse_rope, - enable_attention=not args.no_fuse_attention, - enable_attention_block=not args.no_fuse_attention_block, - enable_add_clipped=not args.no_fuse_add_clipped, - ) - - summary: dict[str, Any] = { - "model_id": args.model_id, - "audio_file": audio_file, - "audio_duration_sec": _audio_duration_seconds(audio_file), - "weights_dir": str(weights_dir_path) if weights_dir_path is not None else None, - "warmup": args.warmup, - "repeats": args.repeats, - "profile_repeats": args.profile_repeats, - "torch_dtype": args.torch_dtype, - "trust_remote_code": bool(args.trust_remote_code), - "local_files_only": bool(args.local_files_only), - } - - handwritten_summary = None - if not args.skip_handwritten and weights_dir_path is not None: - handwritten_summary = _benchmark_handwritten( - weights_dir=weights_dir_path, - audio_file=audio_file, - warmup=args.warmup, - repeats=args.repeats, - profile_repeats=args.profile_repeats, - artifact_dir=artifact_dir, - ) - summary["handwritten"] = handwritten_summary - - transpiled_summary = None - if not args.skip_transpiled: - transpiled_summary = _benchmark_transpiled( - model_id=args.model_id, - audio_file=audio_file, - weights_dir=str(weights_dir_path) if weights_dir_path is not None else None, - warmup=args.warmup, - repeats=args.repeats, - profile_repeats=args.profile_repeats, - torch_dtype=torch_dtype, - token=args.token, - trust_remote_code=args.trust_remote_code, - local_files_only=args.local_files_only, - artifact_dir=artifact_dir, - fusion_config=fusion_config, - ) - summary["transpiled"] = transpiled_summary - - if handwritten_summary is not None and transpiled_summary is not None: - hw_total = handwritten_summary["internal_total_time_ms"]["mean"] - tr_total = transpiled_summary["total_ms"]["mean"] - hw_graph = handwritten_summary.get("profile_summary", {}).get("graph_total_ms") - tr_graph = transpiled_summary["graph_only_ms"]["mean"] - summary["comparison"] = { - "end_to_end_speedup_vs_handwritten": hw_total / max(tr_total, 1e-9), - "graph_only_speedup_vs_handwritten": ( - float(hw_graph["mean"]) / max(tr_graph, 1e-9) - if isinstance(hw_graph, dict) and "mean" in hw_graph - else None - ), - } - - if artifact_dir is not None: - summary_path = artifact_dir / "summary.json" - summary_path.write_text(json.dumps(_roundtrip_jsonable(summary), indent=2, sort_keys=True) + "\n") - print(f"saved_summary={summary_path}") - if handwritten_summary is not None and handwritten_summary.get("profile_file"): - print(f"saved_handwritten_profile={handwritten_summary['profile_file']}") - if transpiled_summary is not None and transpiled_summary.get("profile_file"): - print(f"saved_transpiled_profile={transpiled_summary['profile_file']}") - - _print_summary( - audio_duration_sec=float(summary["audio_duration_sec"]), - handwritten=handwritten_summary, - transpiled=transpiled_summary, - ) - - return 0 - - -if __name__ == "__main__": - raise SystemExit(main()) diff --git a/python/tests/transpile/tools/inspect_gemma4_cleanup.py b/python/tests/transpile/tools/inspect_gemma4_cleanup.py deleted file mode 100644 index 835311f2e..000000000 --- a/python/tests/transpile/tools/inspect_gemma4_cleanup.py +++ /dev/null @@ -1,246 +0,0 @@ -from __future__ import annotations - -import argparse -import os -import sys -from pathlib import Path - -import torch -from transformers import AutoModelForCausalLM - -PYTHON_ROOT = Path(__file__).resolve().parents[3] -if str(PYTHON_ROOT) not in sys.path: - sys.path.insert(0, str(PYTHON_ROOT)) - -from cactus.transpile.capture_pytorch import capture_model -from cactus.transpile.canonicalize.cleanup import canonicalize_exported_graph -from cactus.transpile.graph_ir import IRGraph -from cactus.transpile.model_adapters import canonicalize_model_interface - - -class Gemma4FullModelWrapper(torch.nn.Module): - def __init__(self, model: torch.nn.Module): - super().__init__() - self.model = canonicalize_model_interface(model, task="causal_lm_logits").module - - def forward(self, input_ids: torch.Tensor) -> torch.Tensor: - return self.model(input_ids) - - -class Gemma4FirstBlockCheckpointWrapper(torch.nn.Module): - def __init__(self, model: torch.nn.Module, checkpoint_name: str): - super().__init__() - adapter = canonicalize_model_interface(model, task="causal_lm_logits") - if adapter.family != "gemma4": - raise ValueError(f"expected gemma4 adapter, got family={adapter.family}") - if not hasattr(adapter.module, "debug_first_block"): - raise ValueError("Gemma4 adapter does not expose debug_first_block()") - self.adapter = adapter.module - self.checkpoint_name = checkpoint_name - - def forward(self, input_ids: torch.Tensor) -> torch.Tensor: - checkpoints = self.adapter.debug_first_block(input_ids) - if self.checkpoint_name not in checkpoints: - available = ", ".join(sorted(checkpoints.keys())) - raise KeyError(f"unknown checkpoint {self.checkpoint_name!r}; available: {available}") - return checkpoints[self.checkpoint_name] - - -def format_ir_graph(graph: IRGraph, *, max_nodes: int | None = None) -> str: - lines: list[str] = [] - lines.append("IR Summary") - lines.append(f" inputs={graph.inputs}") - lines.append(f" outputs={graph.outputs}") - lines.append(f" nodes={len(graph.order)}") - if graph.meta: - lines.append(f" meta={graph.meta}") - lines.append("") - lines.append("IR Nodes") - - node_ids = graph.order if max_nodes is None else graph.order[:max_nodes] - for index, node_id in enumerate(node_ids): - node = graph.nodes[node_id] - lines.append(f"[{index}] {node.id} op={node.op}") - lines.append(f" inputs={node.inputs}") - lines.append(f" outputs={node.outputs}") - if node.attrs: - lines.append(f" attrs={node.attrs}") - if node.kind != "generic": - lines.append(f" kind={node.kind}") - if node.meta: - lines.append(f" meta={node.meta}") - for output_id in node.outputs: - value = graph.values.get(output_id) - if value is None: - continue - lines.append( - f" value[{output_id}] shape={value.shape} dtype={value.dtype} users={value.users}" - ) - - if max_nodes is not None and len(graph.order) > max_nodes: - lines.append(f"... ({len(graph.order) - max_nodes} more nodes omitted)") - - if graph.constants: - lines.append("") - lines.append("IR Constants") - constant_ids = sorted(graph.constants.keys()) - preview_ids = constant_ids if max_nodes is None else constant_ids[: max(8, min(len(constant_ids), max_nodes))] - for value_id in preview_ids: - value = graph.constants[value_id] - value_meta = graph.values.get(value_id) - shape = None if value_meta is None else value_meta.shape - dtype = None if value_meta is None else value_meta.dtype - lines.append(f"- {value_id}: type={type(value).__name__} shape={shape} dtype={dtype}") - if len(preview_ids) < len(constant_ids): - lines.append(f"... ({len(constant_ids) - len(preview_ids)} more constants omitted)") - - return "\n".join(lines) - - -def format_unsupported_ops(graph: IRGraph) -> str: - counts = graph.meta.get("canonical_unsupported_op_counts", {}) - if not isinstance(counts, dict): - counts = {} - - examples: dict[str, list[str]] = {} - for node_id in graph.order: - node = graph.nodes[node_id] - print(f"node {node_id} op={node.op}") - if node.op not in counts: - continue - bucket = examples.setdefault(node.op, []) - if len(bucket) < 5: - bucket.append(node.id) - - total = sum(int(count) for count in counts.values()) - lines = ["Unsupported Ops", f" unique={len(counts)}", f" total_nodes={total}"] - if not counts: - lines.append(" none") - return "\n".join(lines) - - for op_name, count in counts.items(): - sample_nodes = ", ".join(examples.get(op_name, ())) - lines.append(f"- {op_name}: {count}") - if sample_nodes: - lines.append(f" examples={sample_nodes}") - return "\n".join(lines) - - -def resolve_local_model_path(model_id: str) -> str: - model_path = Path(model_id) - if model_path.exists(): - return str(model_path) - - cache_root = Path.home() / ".cache" / "huggingface" / "hub" / ("models--" + model_id.replace("/", "--")) - snapshots_dir = cache_root / "snapshots" - if not snapshots_dir.exists(): - raise FileNotFoundError( - f"could not find local snapshot for {model_id!r} under {snapshots_dir}" - ) - snapshots = sorted(path for path in snapshots_dir.iterdir() if path.is_dir()) - if not snapshots: - raise FileNotFoundError(f"no snapshots found for {model_id!r} under {snapshots_dir}") - return str(snapshots[-1]) - - -def load_model(model_id: str) -> torch.nn.Module: - token = os.environ.get("HF_TOKEN") - common_kwargs: dict[str, object] = { - "local_files_only": True, - } - if token: - common_kwargs["token"] = token - - local_model_path = resolve_local_model_path(model_id) - model = AutoModelForCausalLM.from_pretrained( - local_model_path, - torch_dtype=torch.float16, - device_map=None, - low_cpu_mem_usage=True, - **common_kwargs, - ).eval() - return model - - -def build_module( - model: torch.nn.Module, - *, - mode: str, - checkpoint: str, -) -> torch.nn.Module: - if mode == "full": - return Gemma4FullModelWrapper(model).eval() - if mode == "first_block": - return Gemma4FirstBlockCheckpointWrapper(model, checkpoint).eval() - raise ValueError(f"unknown mode: {mode}") - - -def main() -> None: - parser = argparse.ArgumentParser( - description="Capture Gemma4, run canonicalize/cleanup.py, and report unsupported ops.", - ) - parser.add_argument( - "--model-id", - default=os.environ.get("CACTUS_GEMMA_HF_MODEL_ID", "google/gemma-4-E2B"), - help="Hugging Face model id. Uses local cache only.", - ) - parser.add_argument( - "--mode", - choices=("first_block", "full"), - default="first_block", - help="Capture a first-block checkpoint or the full Gemma4 adapter graph.", - ) - parser.add_argument( - "--checkpoint", - default="after_ffn_residual", - help="First-block checkpoint name when --mode=first_block.", - ) - parser.add_argument( - "--input-ids", - default="2,818,5279,529,7001,563", - help="Comma-separated integer token ids to capture.", - ) - parser.add_argument( - "--max-nodes", - type=int, - default=120, - help="Maximum IR nodes to print when --print-ir is enabled.", - ) - parser.add_argument( - "--print-ir", - action="store_true", - help="Also print the IR before and after canonical cleanup.", - ) - args = parser.parse_args() - - model = load_model(args.model_id) - module = build_module(model, mode=args.mode, checkpoint=args.checkpoint) - token_ids = [int(token.strip()) for token in args.input_ids.split(",") if token.strip()] - input_ids = torch.tensor([token_ids], dtype=torch.long) - - print(f"Model: {args.model_id}") - print(f"Mode: {args.mode}") - print(f"Checkpoint: {args.checkpoint if args.mode == 'first_block' else '(full model)'}") - print(f"Input IDs shape: {tuple(input_ids.shape)}") - print("") - - captured = capture_model(module, (input_ids,)) - - if args.print_ir: - print("=== IR Before Canonical Cleanup ===") - print(format_ir_graph(captured.ir_graph, max_nodes=args.max_nodes)) - print("") - - canonicalize_exported_graph(captured.ir_graph) - - print("=== Unsupported Ops After Canonical Cleanup ===") - print(format_unsupported_ops(captured.ir_graph)) - - if args.print_ir: - print("") - print("=== IR After Canonical Cleanup ===") - print(format_ir_graph(captured.ir_graph, max_nodes=args.max_nodes)) - - -if __name__ == "__main__": - main() diff --git a/python/tests/transpile/tools/inspect_transpile_cleanup.py b/python/tests/transpile/tools/inspect_transpile_cleanup.py deleted file mode 100644 index b5b2f3d40..000000000 --- a/python/tests/transpile/tools/inspect_transpile_cleanup.py +++ /dev/null @@ -1,144 +0,0 @@ -from __future__ import annotations - -import argparse -import sys -from pathlib import Path - -import torch -import torch.nn as nn - -PYTHON_ROOT = Path(__file__).resolve().parents[3] -if str(PYTHON_ROOT) not in sys.path: - sys.path.insert(0, str(PYTHON_ROOT)) - -from cactus.transpile.capture_pytorch import capture_model -from cactus.transpile.capture_pytorch import dump_graph -from cactus.transpile.canonicalize.cleanup import canonicalize_exported_graph -from cactus.transpile.graph_ir import IRGraph - - -class Toy(nn.Module): - def forward(self, x, y): - z = x + y - z = z * 0.5 - z = torch.nn.functional.gelu(z) - z = torch.softmax(z, dim=-1) - return z - - -class AttentionBlockToy(nn.Module): - def __init__(self, hidden_size: int = 8): - super().__init__() - self.out_proj = nn.Linear(hidden_size, hidden_size, bias=False, dtype=torch.float16) - - def forward(self, q, k, v, gate): - attn = torch.nn.functional.scaled_dot_product_attention(q, k, v, is_causal=True, scale=1.0) - attn = attn.transpose(1, 2).reshape(q.shape[0], q.shape[2], -1) - attn = attn * torch.sigmoid(gate) - attn = attn.to(self.out_proj.weight.dtype) - return self.out_proj(attn) - - -def build_case(name: str) -> tuple[nn.Module, tuple[torch.Tensor, ...]]: - if name == "toy": - model = Toy().eval() - args = ( - torch.randn(2, 4, dtype=torch.float16), - torch.randn(2, 4, dtype=torch.float16), - ) - return model, args - - if name == "attention_block": - model = AttentionBlockToy().eval() - args = ( - torch.randn(1, 2, 3, 4, dtype=torch.float16), - torch.randn(1, 2, 3, 4, dtype=torch.float16), - torch.randn(1, 2, 3, 4, dtype=torch.float16), - torch.randn(1, 3, 8, dtype=torch.float16), - ) - return model, args - - raise ValueError(f"unknown case: {name}") - - -def format_ir_graph(graph: IRGraph) -> str: - lines: list[str] = [] - lines.append("IR Summary") - lines.append(f" inputs={graph.inputs}") - lines.append(f" outputs={graph.outputs}") - lines.append(f" nodes={len(graph.order)}") - if graph.meta: - lines.append(f" meta={graph.meta}") - lines.append("") - lines.append("IR Nodes") - - for index, node_id in enumerate(graph.order): - node = graph.nodes[node_id] - lines.append(f"[{index}] {node.id} op={node.op}") - lines.append(f" inputs={node.inputs}") - lines.append(f" outputs={node.outputs}") - if node.attrs: - lines.append(f" attrs={node.attrs}") - if node.kind != "generic": - lines.append(f" kind={node.kind}") - if node.meta: - lines.append(f" meta={node.meta}") - for output_id in node.outputs: - value = graph.values.get(output_id) - if value is None: - continue - lines.append( - f" value[{output_id}] shape={value.shape} dtype={value.dtype} users={value.users}" - ) - - if graph.constants: - lines.append("") - lines.append("IR Constants") - for value_id, value in graph.constants.items(): - value_meta = graph.values.get(value_id) - shape = None if value_meta is None else value_meta.shape - dtype = None if value_meta is None else value_meta.dtype - typename = type(value).__name__ - lines.append(f"- {value_id}: type={typename} shape={shape} dtype={dtype}") - - return "\n".join(lines) - - -def main() -> None: - parser = argparse.ArgumentParser( - description="Capture a small PyTorch model, dump the exported graph, and print IR before/after cleanup.", - ) - parser.add_argument( - "--case", - choices=("toy", "attention_block"), - default="toy", - help="Which built-in model to inspect.", - ) - parser.add_argument( - "--no-meta", - action="store_true", - help="Hide exported graph metadata to keep the dump shorter.", - ) - args = parser.parse_args() - - model, example_args = build_case(args.case) - captured = capture_model(model, example_args) - - print(f"Case: {args.case}") - print(f"Model: {type(model).__name__}") - print("") - print("=== Exported Graph ===") - print(dump_graph(captured, include_meta=not args.no_meta)) - print("") - print("=== IR Before Cleanup ===") - print(format_ir_graph(captured.ir_graph)) - print("") - - canonicalize_exported_graph(captured.ir_graph) - - print("=== IR After Cleanup ===") - print(format_ir_graph(captured.ir_graph)) - - -if __name__ == "__main__": - main() diff --git a/python/tests/transpile/tools/torch_capture_smoke.py b/python/tests/transpile/tools/torch_capture_smoke.py deleted file mode 100644 index f3be0f51e..000000000 --- a/python/tests/transpile/tools/torch_capture_smoke.py +++ /dev/null @@ -1,79 +0,0 @@ -from __future__ import annotations - -import sys -from pathlib import Path - -import torch -import torch.nn as nn - -PYTHON_ROOT = Path(__file__).resolve().parents[3] -if str(PYTHON_ROOT) not in sys.path: - sys.path.insert(0, str(PYTHON_ROOT)) - -from cactus.transpile.capture_pytorch import ( - capture_model_with_fallback, - dump_graph, - get_dtype, - get_shape, -) -from cactus.transpile.lower import transpile_captured - - -class TinyBlock(nn.Module): - def __init__(self): - super().__init__() - self.fc1 = nn.Linear(8, 16) - self.fc2 = nn.Linear(16, 4) - - def forward(self, x): - h = self.fc1(x) - h = torch.relu(h) - y = self.fc2(h) - return y - - -def main(): - model = TinyBlock().eval() - x = torch.randn(2, 8) - - captured = capture_model_with_fallback(model, args=(x,)) - - print("strict:", captured.strict) - print("graph_module type:", type(captured.graph_module).__name__) - print("state_dict keys:", list(captured.state_dict.keys())) - print() - - print("=== FX Graph ===") - print(captured.graph) - print() - - print("=== Node Walk ===") - for i, node in enumerate(captured.graph.nodes): - print(f"[{i}] op={node.op} name={node.name} target={node.target}") - print(f" args={node.args}") - print(f" kwargs={node.kwargs}") - print(f" shape={get_shape(node)} dtype={get_dtype(node)}") - print() - - print("=== Full Dump ===") - print(dump_graph(captured)) - print() - - print("=== Lower To Cactus ===") - try: - transpiled = transpile_captured(captured) - print("lowering succeeded") - print("lowered graph type:", type(transpiled.graph).__name__) - print("runtime inputs:", transpiled.runtime_inputs) - print("bound constants:", transpiled.bound_constants) - print("outputs:", transpiled.outputs) - print() - print("binding runtime input 0 from example tensor") - transpiled.set_input(0, x) - except NotImplementedError as exc: - print("lowering stopped on unimplemented case:") - print(f" {exc}") - - -if __name__ == "__main__": - main() diff --git a/python/tests/transpile/tools/transpile_hf_causal_lm.py b/python/tests/transpile/tools/transpile_hf_causal_lm.py deleted file mode 100644 index a1982228e..000000000 --- a/python/tests/transpile/tools/transpile_hf_causal_lm.py +++ /dev/null @@ -1,20 +0,0 @@ -from __future__ import annotations - -import sys -from pathlib import Path - -PYTHON_ROOT = Path(__file__).resolve().parents[3] -if str(PYTHON_ROOT) not in sys.path: - sys.path.insert(0, str(PYTHON_ROOT)) - -from cactus.transpile.hf_model import main as transpile_main - - -def main() -> int: - if "--task" not in sys.argv[1:]: - sys.argv[1:1] = ["--task", "causal_lm_logits"] - return transpile_main() - - -if __name__ == "__main__": - raise SystemExit(main()) From a9fcf3fb465cf64581cf45fc135641eb5703f4ee Mon Sep 17 00:00:00 2001 From: Parkirat Date: Wed, 24 Jun 2026 16:40:42 -0400 Subject: [PATCH 02/84] Initial converter progress --- .../cactus/transpiler/converter/constants.py | 1 + python/cactus/transpiler/converter/convert.py | 32 ++ .../converter/model_preprocessing.py | 13 + .../transpiler/converter/model_utils.py | 107 +++++ python/cactus/transpiler/converter/models.py | 60 +++ python/cactus/transpiler/converter/temp.py | 395 ++++++++++++++++++ 6 files changed, 608 insertions(+) create mode 100644 python/cactus/transpiler/converter/constants.py create mode 100644 python/cactus/transpiler/converter/convert.py create mode 100644 python/cactus/transpiler/converter/model_preprocessing.py create mode 100644 python/cactus/transpiler/converter/model_utils.py create mode 100644 python/cactus/transpiler/converter/models.py create mode 100644 python/cactus/transpiler/converter/temp.py diff --git a/python/cactus/transpiler/converter/constants.py b/python/cactus/transpiler/converter/constants.py new file mode 100644 index 000000000..915ae1fc0 --- /dev/null +++ b/python/cactus/transpiler/converter/constants.py @@ -0,0 +1 @@ +DEFAULT_PROMPT = "What is the capital of France?" \ No newline at end of file diff --git a/python/cactus/transpiler/converter/convert.py b/python/cactus/transpiler/converter/convert.py new file mode 100644 index 000000000..87a15b41e --- /dev/null +++ b/python/cactus/transpiler/converter/convert.py @@ -0,0 +1,32 @@ +import torch +import json +from transformers import AutoTokenizer, AutoModelForCausalLM + +from model_preprocessing import load_model, load_input + + + +def model_ops_to_json(id:str, output_path:str): + model = load_model(id) + input = load_input() + + model.config.return_dict = False if hasattr(model, "config") else None + model.config.use_cache = False if hasattr(model, "use_cache") else None + + with torch.no_grad(): + exported_model = torch.export.export(model, args=(), kwargs=dict(input), strict=False).run_decompositions() + + #Iterates across every node in the graph and stores it as a layer object + for i, record in enumerate(exported_model.graph_module.graph.nodes): + + + + + + + + +if __name__ == "main": + load_model(input("Enter model ID: ")) + + diff --git a/python/cactus/transpiler/converter/model_preprocessing.py b/python/cactus/transpiler/converter/model_preprocessing.py new file mode 100644 index 000000000..bc3fd3b14 --- /dev/null +++ b/python/cactus/transpiler/converter/model_preprocessing.py @@ -0,0 +1,13 @@ +import constants +from transformers import AutoTokenizer, AutoModelForCausalLM + +def load_model(id:str) -> object: + model = AutoModelForCausalLM.from_pretrained(id, attn_implementation = "eager") + model.eval() + return model + +#Currently just handles single input type (strings), but will change to make input a list to allow mutli-modal models to run multiple different inputs +def load_input() -> tuple: + tokenizer = AutoTokenizer.from_pretrained(id) + input = (tokenizer(constants.DEFAULT_PROMPT, return_tensors="pt")["input_ids"],) + return input diff --git a/python/cactus/transpiler/converter/model_utils.py b/python/cactus/transpiler/converter/model_utils.py new file mode 100644 index 000000000..1bc5ebb2c --- /dev/null +++ b/python/cactus/transpiler/converter/model_utils.py @@ -0,0 +1,107 @@ +from typing import Any, Optional +import torch +from models import TensorInstance, Slice + +def jsonable(x: Any) -> Any: + """ + Convert graph/export objects into JSON-safe values. + """ + if isinstance(x, torch.fx.Node): + return {"node": x.name} + + if isinstance(x, torch.Tensor): + return TensorInstance(x) + + if isinstance(x, torch.Size): + return list(x) + + if isinstance(x, torch.dtype): + return str(x) + + if isinstance(x, torch.device): + return str(x) + + if isinstance(x, slice): + return Slice(x) + + if isinstance(x, range): + return list(x) + + if isinstance(x, (list, tuple)): + return [jsonable(v) for v in x] + + if isinstance(x, dict): + return {str(k): jsonable(v) for k, v in x.items()} + + if isinstance(x, (str, int, float, bool)) or x is None: + return x + + return repr(x) + + +def aten_name(target: Any) -> str: + """ + Convert default naming to simpler name. + For non-ATen targets, falls back to str(target). + """ + schema = getattr(target, "_schema", None) + + if schema is not None: + if "::" in schema.name: + namespace, op = schema.name.split("::", 1) + overload = schema.overload_name if schema.overload_name else "default" + return f"{namespace}.{op}.{overload}" + + if hasattr(target, "name"): + try: + return target.name() + except Exception: + pass + + return str(target) + + +def extract_tensor_meta(node: torch.fx.Node) -> Optional[Any]: + """ + torch.export usually stores fake tensor output metadata in node.meta["val"]. + """ + if "val" not in node.meta: + return None + + return jsonable(node.meta["val"]) + + +#This function needs to be further optimized and cleaned up +#CLEANUP: Create object for output dictionary (potentially one outter object as well); Condense down for loop +def extract_module_stack(node: torch.fx.Node) -> Optional[Any]: + """ + Sometimes torch.export nodes contain original module context here. + This can help later when mapping ops back to model.layers.X.self_attn, mlp, etc. + """ + stack = node.meta.get("nn_module_stack", None) + + if stack is None: + return None + + out = [] + + for key, value in stack.items(): + if isinstance(value, tuple) and len(value) >= 2: + module_path = value[0] + module_type = value[1] + out.append( + { + "key": str(key), + "module_path": str(module_path), + "module_type": getattr(module_type, "__name__", str(module_type)), + } + ) + else: + out.append( + { + "key": str(key), + "value": repr(value), + } + ) + + return out diff --git a/python/cactus/transpiler/converter/models.py b/python/cactus/transpiler/converter/models.py new file mode 100644 index 000000000..4dc8910a8 --- /dev/null +++ b/python/cactus/transpiler/converter/models.py @@ -0,0 +1,60 @@ +from pydantic import BaseModel +from typing import Any +import model_utils +import torch + + +class InputtedModel(BaseModel): + name: str + input_modalities: list[str] +# output_modalities: list[str] -> Possible need for this, but not sure yet + +class LayerRecord(BaseModel): + index: str + name: str + node_type: str + target: str + schema: str + args: Any + kwargs: Any + users: list[str] + tensor_output_meta: Any | None + module_stack: Any | None + + @classmethod + def from_node(cls, i:int, node: torch.Node) -> "LayerRecord": + return cls( + index = i, + name = str(node.name), + node_type = str(node.op), + target = str(node.target), + args = model_utils.jsonable(node.args), + kwargs = model_utils.jsonable(node.kwargs), + users = [user.name for user in node.users], + tensor_output_meta = model_utils.extract_tensor_meta(node), + module_stack = model_utils.extract_module_stack(node) + ) + +class TensorInstance(BaseModel): + shape: list + dtype: str + device: str + requires_grad: bool + + @classmethod + def from_tensor(cls, x: torch.Tensor) -> "TensorInstance": + return cls(shape=list(x.shape), dtype=str(x.dtype), device=str(x.device), requires_grad = bool(x.requires_grad)) + +class Slice(BaseModel): + start: Any + stop: Any + step: Any + + @classmethod + def from_slice(cls, x: slice) -> "Slice": + return cls(start = model_utils.jsonable(x.start), stop = model_utils.jsonable(x.stop), step = model_utils.jsonable(x.step)) + + +class JSONResult(BaseModel): + +class LayerMap(BaseModel) \ No newline at end of file diff --git a/python/cactus/transpiler/converter/temp.py b/python/cactus/transpiler/converter/temp.py new file mode 100644 index 000000000..2d997cb96 --- /dev/null +++ b/python/cactus/transpiler/converter/temp.py @@ -0,0 +1,395 @@ +# import argparse +# import json +# from collections import Counter +# from typing import Any, Dict, Optional + +# import torch +# from transformers import ( +# AutoModel, +# AutoModelForCausalLM, +# AutoModelForSequenceClassification, +# AutoTokenizer, +# ) + + +# def jsonable(x: Any) -> Any: +# """ +# Convert graph/export objects into JSON-safe values. +# """ +# if isinstance(x, torch.fx.Node): +# return {"node": x.name} + +# if isinstance(x, torch.Tensor): +# return { +# "shape": list(x.shape), +# "dtype": str(x.dtype), +# "device": str(x.device), +# "requires_grad": bool(x.requires_grad), +# } + +# if isinstance(x, torch.Size): +# return list(x) + +# if isinstance(x, torch.dtype): +# return str(x) + +# if isinstance(x, torch.device): +# return str(x) + +# if isinstance(x, slice): +# return { +# "slice": { +# "start": jsonable(x.start), +# "stop": jsonable(x.stop), +# "step": jsonable(x.step), +# } +# } + +# if isinstance(x, range): +# return list(x) + +# if isinstance(x, (list, tuple)): +# return [jsonable(v) for v in x] + +# if isinstance(x, dict): +# return {str(k): jsonable(v) for k, v in x.items()} + +# if isinstance(x, (str, int, float, bool)) or x is None: +# return x + +# return repr(x) + + +# def aten_name(target: Any) -> str: +# """ +# Convert torch.ops.aten.linear.default into: +# aten.linear.default + +# For non-ATen targets, falls back to str(target). +# """ +# schema = getattr(target, "_schema", None) + +# if schema is not None: +# # Example: +# # schema.name = "aten::linear" +# # schema.overload_name = "" or "Tensor" or "int" +# if "::" in schema.name: +# namespace, op = schema.name.split("::", 1) +# overload = schema.overload_name if schema.overload_name else "default" +# return f"{namespace}.{op}.{overload}" + +# if hasattr(target, "name"): +# try: +# return target.name() +# except Exception: +# pass + +# return str(target) + + +# def schema_string(target: Any) -> Optional[str]: +# schema = getattr(target, "_schema", None) +# return str(schema) if schema is not None else None + + +# def extract_tensor_meta(node: torch.fx.Node) -> Optional[Any]: +# """ +# torch.export usually stores fake tensor output metadata in node.meta["val"]. +# """ +# if "val" not in node.meta: +# return None + +# return jsonable(node.meta["val"]) + + +# def extract_module_stack(node: torch.fx.Node) -> Optional[Any]: +# """ +# Sometimes torch.export nodes contain original module context here. +# This can help later when mapping ops back to model.layers.X.self_attn, mlp, etc. +# """ +# stack = node.meta.get("nn_module_stack", None) + +# if stack is None: +# return None + +# out = [] + +# for key, value in stack.items(): +# if isinstance(value, tuple) and len(value) >= 2: +# module_path = value[0] +# module_type = value[1] +# out.append( +# { +# "key": str(key), +# "module_path": str(module_path), +# "module_type": getattr(module_type, "__name__", str(module_type)), +# } +# ) +# else: +# out.append( +# { +# "key": str(key), +# "value": repr(value), +# } +# ) + +# return out + + +# def make_graph_json( +# exported: torch.export.ExportedProgram, +# model_name: str, +# task: str, +# attn_implementation: str, +# ran_decompositions: bool, +# ) -> Dict[str, Any]: +# graph_module = exported.graph_module +# nodes_json = [] + +# for i, node in enumerate(graph_module.graph.nodes): +# target = aten_name(node.target) + +# node_json = { +# "index": i, +# "name": node.name, +# "node_type": node.op, +# "target": target, +# "schema": schema_string(node.target), +# "args": jsonable(node.args), +# "kwargs": jsonable(node.kwargs), +# "users": [user.name for user in node.users], +# } + +# tensor_meta = extract_tensor_meta(node) +# if tensor_meta is not None: +# node_json["tensor_meta"] = tensor_meta + +# module_stack = extract_module_stack(node) +# if module_stack is not None: +# node_json["module_stack"] = module_stack + +# nodes_json.append(node_json) + +# op_counts = Counter( +# node["target"] +# for node in nodes_json +# if node["node_type"] == "call_function" +# ) + +# return { +# "model_name": model_name, +# "task": task, +# "attn_implementation": attn_implementation, +# "ran_decompositions": ran_decompositions, +# "num_nodes": len(nodes_json), +# "num_call_function_nodes": sum(op_counts.values()), +# "op_counts": dict(op_counts.most_common()), +# "graph_signature": repr(exported.graph_signature), +# "range_constraints": repr(exported.range_constraints), +# "nodes": nodes_json, +# } + + +# def load_model(model_name: str, task: str, attn_implementation: str): +# """ +# Load HF model with eager attention when supported. +# """ +# model_cls_map = { +# "base": AutoModel, +# "causal-lm": AutoModelForCausalLM, +# "sequence-classification": AutoModelForSequenceClassification, +# } + +# if task not in model_cls_map: +# raise ValueError(f"Unsupported task: {task}") + +# model_cls = model_cls_map[task] + +# try: +# model = model_cls.from_pretrained( +# model_name, +# attn_implementation=attn_implementation, +# ) +# except TypeError: +# print( +# f"Warning: this model class did not accept " +# f'attn_implementation="{attn_implementation}". ' +# f"Retrying without it." +# ) +# model = model_cls.from_pretrained(model_name) + +# model.eval() + +# # Hugging Face outputs are often ModelOutput dataclasses. +# # Tuples are easier for torch.export to serialize. +# if hasattr(model, "config"): +# model.config.return_dict = False + +# # Important for GPT/LLaMA/Mistral/etc. +# # Disables KV-cache outputs and cache-specific forward paths. +# if hasattr(model.config, "use_cache"): +# model.config.use_cache = False + +# return model + + +# def export_hf_model_to_json( +# model_name: str, +# output_path: str, +# task: str = "causal-lm", +# sample_text: str = "Hello world", +# max_length: int = 32, +# attn_implementation: str = "eager", +# decompose: bool = True, +# strict: bool = False, +# ): +# tokenizer = AutoTokenizer.from_pretrained(model_name) + +# model = load_model( +# model_name=model_name, +# task=task, +# attn_implementation=attn_implementation, +# ) + +# inputs = tokenizer( +# sample_text, +# return_tensors="pt", +# truncation=True, +# max_length=max_length, +# ) + +# # Some decoder-only tokenizers do not define pad_token. +# # Not needed for a single example, but useful if you later batch examples. +# if tokenizer.pad_token is None and tokenizer.eos_token is not None: +# tokenizer.pad_token = tokenizer.eos_token + +# with torch.no_grad(): +# exported = torch.export.export( +# model, +# args=(), +# kwargs=dict(inputs), +# strict=strict, +# ) + +# ran_decompositions = False + +# if decompose: +# try: +# exported = exported.run_decompositions() +# ran_decompositions = True +# except Exception as e: +# print(f"Warning: run_decompositions() failed: {repr(e)}") +# print("Continuing with the original exported graph.") + +# graph_json = make_graph_json( +# exported=exported, +# model_name=model_name, +# task=task, +# attn_implementation=attn_implementation, +# ran_decompositions=ran_decompositions, +# ) + +# with open(output_path, "w") as f: +# json.dump(graph_json, f, indent=2) + +# print(f"Saved graph JSON to: {output_path}") +# print(f"Total nodes: {graph_json['num_nodes']}") +# print(f"call_function nodes: {graph_json['num_call_function_nodes']}") +# print() +# print("Top ops:") +# for op, count in list(graph_json["op_counts"].items())[:30]: +# print(f"{count:5d} {op}") + +# attention_ops = [ +# op for op in graph_json["op_counts"] +# if "attention" in op.lower() or "scaled_dot_product" in op.lower() +# ] + +# if attention_ops: +# print() +# print("Attention-like ops still present:") +# for op in attention_ops: +# print(f" {op}") +# print() +# print( +# "If you still see aten.scaled_dot_product_attention.default, " +# "attention was not fully broken into matmul/softmax/matmul." +# ) + + +# # def parse_args(): +# # parser = argparse.ArgumentParser( +# # description="Export a Hugging Face PyTorch model to lower-level ATen JSON." +# # ) + +# # parser.add_argument( +# # "--model", +# # type=str, +# # required=True, +# # help="HF model name, e.g. bert-base-uncased or gpt2.", +# # ) + +# # parser.add_argument( +# # "--output", +# # type=str, +# # default="model_aten_graph.json", +# # help="Output JSON path.", +# # ) + +# # parser.add_argument( +# # "--task", +# # type=str, +# # default="causal-lm", +# # choices=["base", "causal-lm", "sequence-classification"], +# # help="Which AutoModel class to use.", +# # ) + +# # parser.add_argument( +# # "--text", +# # type=str, +# # default="Hello world", +# # help="Sample text used to trace/export the model.", +# # ) + +# # parser.add_argument( +# # "--max-length", +# # type=int, +# # default=32, +# # help="Max token length for the sample input.", +# # ) + +# # parser.add_argument( +# # "--attn-implementation", +# # type=str, +# # default="eager", +# # help='Attention backend. For graph inspection, use "eager".', +# # ) + +# # parser.add_argument( +# # "--no-decompose", +# # action="store_true", +# # help="Disable exported.run_decompositions().", +# # ) + +# # parser.add_argument( +# # "--strict", +# # action="store_true", +# # help="Use strict=True in torch.export.export.", +# # ) + +# # return parser.parse_args() + + +# if __name__ == "__main__": +# args = parse_args() + +# export_hf_model_to_json( +# model_name=args.model, +# output_path=args.output, +# task=args.task, +# sample_text=args.text, +# max_length=args.max_length, +# attn_implementation=args.attn_implementation, +# decompose=not args.no_decompose, +# strict=args.strict, +# ) \ No newline at end of file From cfaa17605710555fee02b4fdc5345d7cde6112e9 Mon Sep 17 00:00:00 2001 From: Parkirat Date: Thu, 25 Jun 2026 11:44:39 -0400 Subject: [PATCH 03/84] Initial converter complete: convers text-to-text model graphs to JSON --- .../cactus/transpiler/converter/constants.py | 3 +- python/cactus/transpiler/converter/convert.py | 30 ++-- .../converter/model_preprocessing.py | 2 +- .../transpiler/converter/model_utils.py | 104 +----------- python/cactus/transpiler/converter/models.py | 150 +++++++++++++++--- 5 files changed, 152 insertions(+), 137 deletions(-) diff --git a/python/cactus/transpiler/converter/constants.py b/python/cactus/transpiler/converter/constants.py index 915ae1fc0..9060d96d5 100644 --- a/python/cactus/transpiler/converter/constants.py +++ b/python/cactus/transpiler/converter/constants.py @@ -1 +1,2 @@ -DEFAULT_PROMPT = "What is the capital of France?" \ No newline at end of file +DEFAULT_PROMPT = "What is the capital of France?" +MODEL_TASK = "Causal_LLM" \ No newline at end of file diff --git a/python/cactus/transpiler/converter/convert.py b/python/cactus/transpiler/converter/convert.py index 87a15b41e..30b4542be 100644 --- a/python/cactus/transpiler/converter/convert.py +++ b/python/cactus/transpiler/converter/convert.py @@ -1,32 +1,34 @@ import torch import json -from transformers import AutoTokenizer, AutoModelForCausalLM - +from models import LayerRecord, LayerMap from model_preprocessing import load_model, load_input +import constants def model_ops_to_json(id:str, output_path:str): model = load_model(id) - input = load_input() + input = load_input(id) - model.config.return_dict = False if hasattr(model, "config") else None - model.config.use_cache = False if hasattr(model, "use_cache") else None + model.config.return_dict = False if hasattr(model.config, "return_dict") else None + model.config.use_cache = False if hasattr(model.config, "use_cache") else None with torch.no_grad(): - exported_model = torch.export.export(model, args=(), kwargs=dict(input), strict=False).run_decompositions() + exported_model = torch.export.export(model, args=input, kwargs={"use_cache":False}, strict=False).run_decompositions() #Iterates across every node in the graph and stores it as a layer object + nodes = [] for i, record in enumerate(exported_model.graph_module.graph.nodes): - - - - - - + nodes.append(LayerRecord.from_node(num = i, x = record)) + + json_ = LayerMap.from_data(x=exported_model, name=id, model_task = constants.MODEL_TASK, nodes_list=nodes).model_dump_json(indent=4) + + with open(output_path+"output.json", "w", encoding="utf-8") as file: + file.write(json_) -if __name__ == "main": - load_model(input("Enter model ID: ")) +if __name__ == "__main__": + model_ops_to_json("google/gemma-4-E2B", "/Users/sandhup/Documents/personal/cactus/python/cactus/transpiler/converter/jsons/") + print("Done") diff --git a/python/cactus/transpiler/converter/model_preprocessing.py b/python/cactus/transpiler/converter/model_preprocessing.py index bc3fd3b14..5d8f4fc71 100644 --- a/python/cactus/transpiler/converter/model_preprocessing.py +++ b/python/cactus/transpiler/converter/model_preprocessing.py @@ -7,7 +7,7 @@ def load_model(id:str) -> object: return model #Currently just handles single input type (strings), but will change to make input a list to allow mutli-modal models to run multiple different inputs -def load_input() -> tuple: +def load_input(id:str) -> tuple: tokenizer = AutoTokenizer.from_pretrained(id) input = (tokenizer(constants.DEFAULT_PROMPT, return_tensors="pt")["input_ids"],) return input diff --git a/python/cactus/transpiler/converter/model_utils.py b/python/cactus/transpiler/converter/model_utils.py index 1bc5ebb2c..2c5824019 100644 --- a/python/cactus/transpiler/converter/model_utils.py +++ b/python/cactus/transpiler/converter/model_utils.py @@ -1,107 +1,5 @@ from typing import Any, Optional import torch -from models import TensorInstance, Slice +# from models import TensorInstance, Slice -def jsonable(x: Any) -> Any: - """ - Convert graph/export objects into JSON-safe values. - """ - if isinstance(x, torch.fx.Node): - return {"node": x.name} - if isinstance(x, torch.Tensor): - return TensorInstance(x) - - if isinstance(x, torch.Size): - return list(x) - - if isinstance(x, torch.dtype): - return str(x) - - if isinstance(x, torch.device): - return str(x) - - if isinstance(x, slice): - return Slice(x) - - if isinstance(x, range): - return list(x) - - if isinstance(x, (list, tuple)): - return [jsonable(v) for v in x] - - if isinstance(x, dict): - return {str(k): jsonable(v) for k, v in x.items()} - - if isinstance(x, (str, int, float, bool)) or x is None: - return x - - return repr(x) - - -def aten_name(target: Any) -> str: - """ - Convert default naming to simpler name. - For non-ATen targets, falls back to str(target). - """ - schema = getattr(target, "_schema", None) - - if schema is not None: - if "::" in schema.name: - namespace, op = schema.name.split("::", 1) - overload = schema.overload_name if schema.overload_name else "default" - return f"{namespace}.{op}.{overload}" - - if hasattr(target, "name"): - try: - return target.name() - except Exception: - pass - - return str(target) - - -def extract_tensor_meta(node: torch.fx.Node) -> Optional[Any]: - """ - torch.export usually stores fake tensor output metadata in node.meta["val"]. - """ - if "val" not in node.meta: - return None - - return jsonable(node.meta["val"]) - - -#This function needs to be further optimized and cleaned up -#CLEANUP: Create object for output dictionary (potentially one outter object as well); Condense down for loop -def extract_module_stack(node: torch.fx.Node) -> Optional[Any]: - """ - Sometimes torch.export nodes contain original module context here. - This can help later when mapping ops back to model.layers.X.self_attn, mlp, etc. - """ - stack = node.meta.get("nn_module_stack", None) - - if stack is None: - return None - - out = [] - - for key, value in stack.items(): - if isinstance(value, tuple) and len(value) >= 2: - module_path = value[0] - module_type = value[1] - out.append( - { - "key": str(key), - "module_path": str(module_path), - "module_type": getattr(module_type, "__name__", str(module_type)), - } - ) - else: - out.append( - { - "key": str(key), - "value": repr(value), - } - ) - - return out diff --git a/python/cactus/transpiler/converter/models.py b/python/cactus/transpiler/converter/models.py index 4dc8910a8..e7f7f2dca 100644 --- a/python/cactus/transpiler/converter/models.py +++ b/python/cactus/transpiler/converter/models.py @@ -1,20 +1,19 @@ from pydantic import BaseModel -from typing import Any -import model_utils +from typing import Any, Optional import torch class InputtedModel(BaseModel): name: str input_modalities: list[str] -# output_modalities: list[str] -> Possible need for this, but not sure yet +# output_modalities: list[str] -> Possible future need for this, but not yet class LayerRecord(BaseModel): - index: str + index: int name: str node_type: str target: str - schema: str + #model_schema: str | None TODO: Need to add this field back and include model_schema during conversion args: Any kwargs: Any users: list[str] @@ -22,17 +21,17 @@ class LayerRecord(BaseModel): module_stack: Any | None @classmethod - def from_node(cls, i:int, node: torch.Node) -> "LayerRecord": + def from_node(cls, num:int, x: torch.Node) -> "LayerRecord": return cls( - index = i, - name = str(node.name), - node_type = str(node.op), - target = str(node.target), - args = model_utils.jsonable(node.args), - kwargs = model_utils.jsonable(node.kwargs), - users = [user.name for user in node.users], - tensor_output_meta = model_utils.extract_tensor_meta(node), - module_stack = model_utils.extract_module_stack(node) + index = num, + name = str(x.name), + node_type = str(x.op), + target = str(x.target), + args = jsonable(x.args), + kwargs = jsonable(x.kwargs), + users = [user.name for user in x.users], + tensor_output_meta = extract_tensor_meta(x), + module_stack = extract_module_stack(x) ) class TensorInstance(BaseModel): @@ -52,9 +51,124 @@ class Slice(BaseModel): @classmethod def from_slice(cls, x: slice) -> "Slice": - return cls(start = model_utils.jsonable(x.start), stop = model_utils.jsonable(x.stop), step = model_utils.jsonable(x.step)) + return cls(start = jsonable(x.start), stop = jsonable(x.stop), step = jsonable(x.step)) -class JSONResult(BaseModel): +class LayerMap(BaseModel): + model_name: str + task: str + graph_signature: str + range_constants: str + nodes: list[LayerRecord] -class LayerMap(BaseModel) \ No newline at end of file + @classmethod + def from_data(cls, x:torch.export.ExportedProgram, name:str, model_task:str, nodes_list:list[LayerRecord]) -> "LayerMap": + return cls(model_name = name, task = model_task, graph_signature = repr(x.graph_signature), range_constants = repr(x.range_constraints), nodes = nodes_list) + + + +"""###################################### MODEL UTILS!!!!!!! ######################################""" + +def jsonable(x: Any) -> Any: + """ + Convert graph/export objects into JSON-safe values. + """ + if isinstance(x, torch.fx.Node): + return {"node": x.name} + + if isinstance(x, torch.Tensor): + return TensorInstance.from_tensor(x) + + if isinstance(x, torch.Size): + return list(x) + + if isinstance(x, torch.dtype): + return str(x) + + if isinstance(x, torch.device): + return str(x) + + if isinstance(x, slice): + return Slice.from_slice(x) + + if isinstance(x, range): + return list(x) + + if isinstance(x, (list, tuple)): + return [jsonable(v) for v in x] + + if isinstance(x, dict): + return {str(k): jsonable(v) for k, v in x.items()} + + if isinstance(x, (str, int, float, bool)) or x is None: + return x + + return repr(x) + + +def aten_name(target: Any) -> str: + """ + Convert default naming to simpler name. + For non-ATen targets, falls back to str(target). + """ + schema = getattr(target, "_schema", None) + + if schema is not None: + if "::" in schema.name: + namespace, op = schema.name.split("::", 1) + overload = schema.overload_name if schema.overload_name else "default" + return f"{namespace}.{op}.{overload}" + + if hasattr(target, "name"): + try: + return target.name() + except Exception: + pass + + return str(target) + + +def extract_tensor_meta(node: torch.fx.Node) -> Optional[Any]: + """ + torch.export usually stores fake tensor output metadata in node.meta["val"]. + """ + if "val" not in node.meta: + return None + + return jsonable(node.meta["val"]) + + +#This function needs to be further optimized and cleaned up +#CLEANUP: Create object for output dictionary (potentially one outter object as well); Condense down for loop +def extract_module_stack(node: torch.fx.Node) -> Optional[Any]: + """ + Sometimes torch.export nodes contain original module context here. + This can help later when mapping ops back to model.layers.X.self_attn, mlp, etc. + """ + stack = node.meta.get("nn_module_stack", None) + + if stack is None: + return None + + out = [] + + for key, value in stack.items(): + if isinstance(value, tuple) and len(value) >= 2: + module_path = value[0] + module_type = value[1] + out.append( + { + "key": str(key), + "module_path": str(module_path), + "module_type": getattr(module_type, "__name__", str(module_type)), + } + ) + else: + out.append( + { + "key": str(key), + "value": repr(value), + } + ) + + return out \ No newline at end of file From 09cb668aafe02b784c45d3f1ad589140c0ba7d28 Mon Sep 17 00:00:00 2001 From: Parkirat Date: Fri, 26 Jun 2026 09:44:41 -0400 Subject: [PATCH 04/84] Inital IR simplification logic --- python/cactus/transpiler/converter/convert.py | 2 +- python/cactus/transpiler/ir/constants.py | 0 python/cactus/transpiler/ir/models.py | 53 +++++++++++++++++++ python/cactus/transpiler/ir/simplify_ir.py | 12 +++++ python/cactus/transpiler/main.py | 6 +++ 5 files changed, 72 insertions(+), 1 deletion(-) create mode 100644 python/cactus/transpiler/ir/constants.py create mode 100644 python/cactus/transpiler/ir/models.py create mode 100644 python/cactus/transpiler/ir/simplify_ir.py create mode 100644 python/cactus/transpiler/main.py diff --git a/python/cactus/transpiler/converter/convert.py b/python/cactus/transpiler/converter/convert.py index 30b4542be..02bdb2fdb 100644 --- a/python/cactus/transpiler/converter/convert.py +++ b/python/cactus/transpiler/converter/convert.py @@ -6,7 +6,7 @@ -def model_ops_to_json(id:str, output_path:str): +def model_ops_to_json(id:str, output_path:str) -> LayerMap: model = load_model(id) input = load_input(id) diff --git a/python/cactus/transpiler/ir/constants.py b/python/cactus/transpiler/ir/constants.py new file mode 100644 index 000000000..e69de29bb diff --git a/python/cactus/transpiler/ir/models.py b/python/cactus/transpiler/ir/models.py new file mode 100644 index 000000000..5545c69e9 --- /dev/null +++ b/python/cactus/transpiler/ir/models.py @@ -0,0 +1,53 @@ +from pydantic import BaseModel +from typing import Any, Optional +import torch +import converter.models as CVModels + +class Value(BaseModel): + id:str + type:str + shape:list[int] | None + + @classmethod + def from_placeholder(x:CVModels.LayerRecord) -> "Value": + pass + + +class Operation(BaseModel): + id:str + op:str + inputs:list["Operation"] + attrs:dict[Any, Any] + module_path:str + + @classmethod + def from_op(x:CVModels.LayerRecord) -> "Operation": + pass + +class Output(BaseModel): + id:str + inputs:list[Operation] + attrs:dict[Any, Any] + module_path:str + + @classmethod + def from_output(x:CVModels.LayerRecord): + pass + + +class SimplifiedModel(BaseModel): + operations:list[Operation] + values:list[Value] + + @classmethod + def from_(x:CVModels) -> "SimplifiedModel": + pass + + + + + + + +"""###################################### MODEL UTILS!!!!!!! ######################################""" + diff --git a/python/cactus/transpiler/ir/simplify_ir.py b/python/cactus/transpiler/ir/simplify_ir.py new file mode 100644 index 000000000..b33360b21 --- /dev/null +++ b/python/cactus/transpiler/ir/simplify_ir.py @@ -0,0 +1,12 @@ +import models, constants +import converter.models as CVModels +import models + +function_map = { + "placeholder" : models.Value.from_placeholder, + "call_function" : models.Operation.from_op, + "output" : models.Output.from_output, +} + +def simplify(messy_ir:CVModels.LayerMap): + for value in messy_ir.nodes: diff --git a/python/cactus/transpiler/main.py b/python/cactus/transpiler/main.py new file mode 100644 index 000000000..6e2408627 --- /dev/null +++ b/python/cactus/transpiler/main.py @@ -0,0 +1,6 @@ +import converter.convert as cv +import converter.models as CVModels + +if __name__ == "__main__": + model_map:CVModels = cv.model_ops_to_json("google/gemma-4-E2B", "/Users/sandhup/Documents/personal/cactus/python/cactus/transpiler/converter/jsons/") + \ No newline at end of file From f8b7ae3d52f8755eb06f644b8757648422149e0f Mon Sep 17 00:00:00 2001 From: Parkirat Date: Tue, 30 Jun 2026 12:20:56 -0400 Subject: [PATCH 05/84] Completed computation graph object --- python/cactus/transpiler/ir/constants.py | 41 ++++++++ python/cactus/transpiler/ir/models.py | 110 ++++++++++++++++----- python/cactus/transpiler/ir/simplify_ir.py | 19 ++-- 3 files changed, 140 insertions(+), 30 deletions(-) diff --git a/python/cactus/transpiler/ir/constants.py b/python/cactus/transpiler/ir/constants.py index e69de29bb..b3502b0f8 100644 --- a/python/cactus/transpiler/ir/constants.py +++ b/python/cactus/transpiler/ir/constants.py @@ -0,0 +1,41 @@ +OP_ATTRS_SCHEMAS = { + "aten._assert_tensor_metadata.default": ["dtype"], + "aten._softmax.default": ["dim", "half_to_float"], + "aten._to_copy.default": [], + "aten.add.Tensor": ["scalar"], + "aten.arange.start_step": ["start", "end", "step"], + "aten.bitwise_and.Tensor": [], + "aten.bitwise_or.Tensor": [], + "aten.bmm.default": [], + "aten.cat.default": ["dim"], + "aten.clone.default": [], + "aten.cos.default": [], + "aten.cumsum.default": ["dim"], + "aten.div.Tensor": ["scalar"], + "aten.embedding.default": ["padding_idx"], + "aten.eq.Scalar": ["scalar"], + "aten.eq.Tensor": [], + "aten.expand.default": ["shape"], + "aten.full.default": ["shape", "fill_value"], + "aten.gelu.default": [], + "aten.gt.Tensor": [], + "aten.index.Tensor": [], + "aten.index_put.default": [], + "aten.le.Tensor": [], + "aten.mean.dim": ["dim", "keepdim"], + "aten.mm.default": [], + "aten.mul.Tensor": ["scalar"], + "aten.ne.Scalar": ["scalar"], + "aten.neg.default": [], + "aten.permute.default": ["dims"], + "aten.pow.Tensor_Scalar": ["exponent"], + "aten.scalar_tensor.default": ["value"], + "aten.select.int": ["dim", "index"], + "aten.sin.default": [], + "aten.slice.Tensor": ["dim", "start", "end", "step"], + "aten.sub.Tensor": ["scalar"], + "aten.tanh.default": [], + "aten.unsqueeze.default": ["dim"], + "aten.view.default": ["shape"], + "aten.where.self": [], +} \ No newline at end of file diff --git a/python/cactus/transpiler/ir/models.py b/python/cactus/transpiler/ir/models.py index 5545c69e9..db191347a 100644 --- a/python/cactus/transpiler/ir/models.py +++ b/python/cactus/transpiler/ir/models.py @@ -1,53 +1,117 @@ from pydantic import BaseModel -from typing import Any, Optional -import torch +from typing import Any +from collections import deque import converter.models as CVModels +import constants -class Value(BaseModel): + +class ValueNode(BaseModel): id:str type:str - shape:list[int] | None + # shape:list[int] | None TODO: May need shape (value or output) later on, but may not be necessary since graph is known to be complete/correct @classmethod - def from_placeholder(x:CVModels.LayerRecord) -> "Value": - pass + def from_placeholder(cls, x:CVModels.LayerRecord) -> "ValueNode": + return cls(id = x.name, type = x.node_type) -class Operation(BaseModel): +class OperationNode(BaseModel): id:str op:str - inputs:list["Operation"] + inputs_id:list[str] + users:list[str] attrs:dict[Any, Any] - module_path:str + #Constructor @classmethod - def from_op(x:CVModels.LayerRecord) -> "Operation": - pass - -class Output(BaseModel): + def from_op(cls, x:CVModels.LayerRecord) -> "OperationNode": + attrs_list:list[str] = constants.OP_ATTRS_SCHEMAS[x.target] + attrs_temp:dict[str,Any] = {} + inputs:list[str] = [] + count = 0 + + for item in x.args: + if isinstance(item, dict) and "node" in item: + inputs.append(item["node"]) + else: + if count < len(attrs_list): + attrs_temp[attrs_list[count]] = item + count += 1 + else: + raise ValueError(f"{x.name}: Layer has more attributes than expected!") + + + return cls(id = x.name, op = x.target, inputs_id = inputs, users = x.users, attrs = attrs_temp) + +class OutputNode(BaseModel): id:str - inputs:list[Operation] - attrs:dict[Any, Any] - module_path:str + inputs_id:list[str] + #Constructor @classmethod - def from_output(x:CVModels.LayerRecord): - pass + def from_output(cls, x: CVModels.LayerRecord) -> "OutputNode": + inputs:list[str] = [] + count = 0 + + for item in x.args: + if isinstance(item, dict) and "node" in item: + inputs.append(item["node"]) + + return cls(id = x.name, inputs_id = inputs) + -class SimplifiedModel(BaseModel): - operations:list[Operation] - values:list[Value] +class CompGraph(BaseModel): + nodes:dict[str, Any] # Turn this into dict of (name: node) + source_nodes:deque[OperationNode, ValueNode] + #Constructor @classmethod - def from_(x:CVModels) -> "SimplifiedModel": + def from_graph(cls, graph:CVModels.LayerMap): + return generate_comp_graph_(graph) + + + def perform_fusions(): pass +"""###################################### MODEL UTILS!!!!!!! ######################################""" +function_map = { + "placeholder" : ValueNode.from_placeholder, + "call_function" : OperationNode.from_op, + "output" : OutputNode.from_output, +} +def generate_comp_graph_(graph: CVModels.LayerMap): + nodes_dict = {} + output_nodes = deque() -"""###################################### MODEL UTILS!!!!!!! ######################################""" + for node in graph.nodes: + new_node = function[node.node_type](node) + nodes_dict[node.name] = new_node + + if isinstance(new_node, OutputNode): + output_nodes.append(new_node) + + return CompGraph(nodes = nodes_dict, source_nodes = output_nodes) + + +# def is_source_node(value: list) -> bool: +# #TODO: Horrifically ugly logical statement here, try to shorten it down. Just checking if input to this node is source node +# if len(value) > 0 and isinstance(value, dict) and "node" in value[0] and value[0]["node"] == "input_ids": +# return True +# else: +# False + +# def add_node_(self:CompGraph, x: CVModels.LayerRecord) -> None: +# node = function_map[x.node_type](x) +# self.nodes.append(node) + +# if(is_source_node(x.args)): +# self.source_nodes.append() + + diff --git a/python/cactus/transpiler/ir/simplify_ir.py b/python/cactus/transpiler/ir/simplify_ir.py index b33360b21..97b43d6af 100644 --- a/python/cactus/transpiler/ir/simplify_ir.py +++ b/python/cactus/transpiler/ir/simplify_ir.py @@ -1,12 +1,17 @@ +""" +The code in this directory is responsible for taking in the raw components lists from the converter folder and transform +it into a valid computation DAG (consisting of both operation nodes and value nodes). Topological sort will then be performed +on the DAG to generate an in-order execution list. +""" + import models, constants import converter.models as CVModels import models -function_map = { - "placeholder" : models.Value.from_placeholder, - "call_function" : models.Operation.from_op, - "output" : models.Output.from_output, -} - def simplify(messy_ir:CVModels.LayerMap): - for value in messy_ir.nodes: + graph:models.CompGraph = models.CompGraph() + + for node in messy_ir.nodes: + function_map[node.node_type](node) + + From 7a5c7508caaf738dafdc52de11886e45087dd051 Mon Sep 17 00:00:00 2001 From: Parkirat Date: Tue, 30 Jun 2026 13:19:11 -0400 Subject: [PATCH 06/84] Added a new field to ValueNodes to keep track of value users --- python/cactus/transpiler/ir/models.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/python/cactus/transpiler/ir/models.py b/python/cactus/transpiler/ir/models.py index db191347a..9d0ad23fb 100644 --- a/python/cactus/transpiler/ir/models.py +++ b/python/cactus/transpiler/ir/models.py @@ -8,11 +8,12 @@ class ValueNode(BaseModel): id:str type:str + users:list[str] # shape:list[int] | None TODO: May need shape (value or output) later on, but may not be necessary since graph is known to be complete/correct @classmethod def from_placeholder(cls, x:CVModels.LayerRecord) -> "ValueNode": - return cls(id = x.name, type = x.node_type) + return cls(id = x.name, type = x.node_type, users = x.users) class OperationNode(BaseModel): From c23a00c87450a9ad4a77c1418e24b03957e0e569 Mon Sep 17 00:00:00 2001 From: Parkirat Date: Tue, 7 Jul 2026 16:39:55 -0400 Subject: [PATCH 07/84] improving ir --- python/cactus/transpiler/ir/constants.py | 41 ------- python/cactus/transpiler/ir/fusions.py | 98 ++++++++++++++++ python/cactus/transpiler/ir/models.py | 123 +++++---------------- python/cactus/transpiler/ir/simplify_ir.py | 5 +- 4 files changed, 127 insertions(+), 140 deletions(-) create mode 100644 python/cactus/transpiler/ir/fusions.py diff --git a/python/cactus/transpiler/ir/constants.py b/python/cactus/transpiler/ir/constants.py index b3502b0f8..e69de29bb 100644 --- a/python/cactus/transpiler/ir/constants.py +++ b/python/cactus/transpiler/ir/constants.py @@ -1,41 +0,0 @@ -OP_ATTRS_SCHEMAS = { - "aten._assert_tensor_metadata.default": ["dtype"], - "aten._softmax.default": ["dim", "half_to_float"], - "aten._to_copy.default": [], - "aten.add.Tensor": ["scalar"], - "aten.arange.start_step": ["start", "end", "step"], - "aten.bitwise_and.Tensor": [], - "aten.bitwise_or.Tensor": [], - "aten.bmm.default": [], - "aten.cat.default": ["dim"], - "aten.clone.default": [], - "aten.cos.default": [], - "aten.cumsum.default": ["dim"], - "aten.div.Tensor": ["scalar"], - "aten.embedding.default": ["padding_idx"], - "aten.eq.Scalar": ["scalar"], - "aten.eq.Tensor": [], - "aten.expand.default": ["shape"], - "aten.full.default": ["shape", "fill_value"], - "aten.gelu.default": [], - "aten.gt.Tensor": [], - "aten.index.Tensor": [], - "aten.index_put.default": [], - "aten.le.Tensor": [], - "aten.mean.dim": ["dim", "keepdim"], - "aten.mm.default": [], - "aten.mul.Tensor": ["scalar"], - "aten.ne.Scalar": ["scalar"], - "aten.neg.default": [], - "aten.permute.default": ["dims"], - "aten.pow.Tensor_Scalar": ["exponent"], - "aten.scalar_tensor.default": ["value"], - "aten.select.int": ["dim", "index"], - "aten.sin.default": [], - "aten.slice.Tensor": ["dim", "start", "end", "step"], - "aten.sub.Tensor": ["scalar"], - "aten.tanh.default": [], - "aten.unsqueeze.default": ["dim"], - "aten.view.default": ["shape"], - "aten.where.self": [], -} \ No newline at end of file diff --git a/python/cactus/transpiler/ir/fusions.py b/python/cactus/transpiler/ir/fusions.py new file mode 100644 index 000000000..bc21b138c --- /dev/null +++ b/python/cactus/transpiler/ir/fusions.py @@ -0,0 +1,98 @@ +from dataclasses import dataclass, field +from typing import Any + +@dataclass(slots=True) +class FusionPatter: + target:str + ops:tuple[str] + path:tuple[int] + + + + + +# @dataclass(frozen=True) +# class FusionPattern: +# """Description of an operation sequence that can become one fused op.""" + +# target: str +# operations: tuple[str, ...] +# path_inputs: tuple[int, ...] +# required_attrs: dict[int, dict[str, Any]] = field(default_factory=dict) +# same_inputs: tuple[tuple[InputRef, InputRef], ...] = () +# input_order: tuple[InputRef, ...] = () +# output_attrs: dict[str, AttrRef] = field(default_factory=dict) + +# @property +# def sequence(self) -> str: +# return "->".join(self.operations) + +# def __post_init__(self) -> None: +# if not self.operations: +# raise ValueError("operations cannot be empty") + +# if len(self.path_inputs) != len(self.operations) - 1: +# raise ValueError("path_inputs must contain one entry between each operation") + +# for node_index in self.required_attrs: +# if node_index < 0 or node_index >= len(self.operations): +# raise ValueError(f"required_attrs contains invalid node index {node_index}") + +# input_refs = [input_ref for input_pair in self.same_inputs for input_ref in input_pair] +# input_refs.extend(self.input_order) + +# for node_index, input_index in input_refs: +# if node_index < 0 or node_index >= len(self.operations): +# raise ValueError(f"input reference contains invalid node index {node_index}") +# if input_index < 0: +# raise ValueError("input reference indexes cannot be negative") + +# for node_index, _ in self.output_attrs.values(): +# if node_index < 0 or node_index >= len(self.operations): +# raise ValueError(f"output_attrs contains invalid node index {node_index}") + + +# @dataclass(frozen=True) +# class FusionMatch: +# """Concrete graph nodes and boundary values matched by a FusionPattern.""" + +# root_id: str +# target: str +# matched_ids: tuple[str, ...] +# external_input_ids: tuple[str, ...] +# attrs: dict[str, Any] = field(default_factory=dict) + + +# RMSNORM_FUSION = FusionPattern( +# target="rmsnorm", +# operations=( +# "aten.mul.Tensor", +# "aten.mul.Tensor", +# "aten.pow.Tensor_Scalar", +# "aten.add.Tensor", +# "aten.mean.dim", +# "aten.pow.Tensor_Scalar", +# ), +# path_inputs=(0, 1, 0, 0, 0), +# required_attrs={ +# 2: {"exponent": -0.5}, +# 3: {"scalar": 1e-06}, +# 4: {"dim": [-1], "keepdim": True}, +# 5: {"exponent": 2}, +# }, +# same_inputs=(((1, 0), (5, 0)),), +# input_order=((1, 0), (0, 1)), +# output_attrs={"epsilon": (3, "scalar")}, +# ) + + +# LINEAR_RESHAPE_FUSION = FusionPattern( +# target="linear_reshape", +# operations=( +# "aten.view.default", +# "aten.mm.default", +# ), +# path_inputs=(0,), +# input_order=((1, 0), (1, 1)), +# output_attrs={"shape": (0, "shape")}, +# ) diff --git a/python/cactus/transpiler/ir/models.py b/python/cactus/transpiler/ir/models.py index 9d0ad23fb..0c8a34b1a 100644 --- a/python/cactus/transpiler/ir/models.py +++ b/python/cactus/transpiler/ir/models.py @@ -1,118 +1,49 @@ -from pydantic import BaseModel +from dataclasses import dataclass from typing import Any from collections import deque import converter.models as CVModels import constants +@dataclass(slots=True) +class Node: + layer: CVModels.LayerRecord + parents: list["Node"] + children: list["Node"] -class ValueNode(BaseModel): - id:str - type:str - users:list[str] - # shape:list[int] | None TODO: May need shape (value or output) later on, but may not be necessary since graph is known to be complete/correct - @classmethod - def from_placeholder(cls, x:CVModels.LayerRecord) -> "ValueNode": - return cls(id = x.name, type = x.node_type, users = x.users) - - -class OperationNode(BaseModel): - id:str - op:str - inputs_id:list[str] - users:list[str] - attrs:dict[Any, Any] +@dataclass(slots=True) +class Graph: + name_map:dict[str, CVModels.LayerRecord] + source_nodes: deque[Node] + consumed_ids: set[str] - #Constructor @classmethod - def from_op(cls, x:CVModels.LayerRecord) -> "OperationNode": - attrs_list:list[str] = constants.OP_ATTRS_SCHEMAS[x.target] - attrs_temp:dict[str,Any] = {} - inputs:list[str] = [] - count = 0 - - for item in x.args: - if isinstance(item, dict) and "node" in item: - inputs.append(item["node"]) - else: - if count < len(attrs_list): - attrs_temp[attrs_list[count]] = item - count += 1 - else: - raise ValueError(f"{x.name}: Layer has more attributes than expected!") - - - return cls(id = x.name, op = x.target, inputs_id = inputs, users = x.users, attrs = attrs_temp) + def from_layer_map(cls, map:CVModels.LayerMap): + return generate_graph(map) -class OutputNode(BaseModel): - id:str - inputs_id:list[str] - - #Constructor - @classmethod - def from_output(cls, x: CVModels.LayerRecord) -> "OutputNode": - inputs:list[str] = [] - count = 0 - - for item in x.args: - if isinstance(item, dict) and "node" in item: - inputs.append(item["node"]) - - return cls(id = x.name, inputs_id = inputs) - - - -class CompGraph(BaseModel): - nodes:dict[str, Any] # Turn this into dict of (name: node) - source_nodes:deque[OperationNode, ValueNode] - - #Constructor - @classmethod - def from_graph(cls, graph:CVModels.LayerMap): - return generate_comp_graph_(graph) - def perform_fusions(): - pass - - - - """###################################### MODEL UTILS!!!!!!! ######################################""" -function_map = { - "placeholder" : ValueNode.from_placeholder, - "call_function" : OperationNode.from_op, - "output" : OutputNode.from_output, -} +def generate_graph(map: CVModels.LayerMap): + temp_map:dict[str, Node] = {} + temp_source = deque() -def generate_comp_graph_(graph: CVModels.LayerMap): - nodes_dict = {} - output_nodes = deque() + for layer_ in map.nodes: + temp = Node(layer=layer_, parents=[], children=[]) + temp_map[layer_.name] = temp + temp_source.append(temp) if layer_.node_type == "output" else None - for node in graph.nodes: - new_node = function[node.node_type](node) - nodes_dict[node.name] = new_node + for layer_ in map.nodes: + for name in layer_.users: + temp_map[layer_.name].children.append(temp_map[name]) - if isinstance(new_node, OutputNode): - output_nodes.append(new_node) - - return CompGraph(nodes = nodes_dict, source_nodes = output_nodes) - + for item in layer_.args: + if isinstance(item, dict) and "node" in item: + temp_map[layer_.name].parents.append(temp_map[item["node"]]) -# def is_source_node(value: list) -> bool: -# #TODO: Horrifically ugly logical statement here, try to shorten it down. Just checking if input to this node is source node -# if len(value) > 0 and isinstance(value, dict) and "node" in value[0] and value[0]["node"] == "input_ids": -# return True -# else: -# False + return Graph(name_map=temp_map, source_nodes=temp_source, consumed_ids={}) -# def add_node_(self:CompGraph, x: CVModels.LayerRecord) -> None: -# node = function_map[x.node_type](x) -# self.nodes.append(node) -# if(is_source_node(x.args)): -# self.source_nodes.append() - diff --git a/python/cactus/transpiler/ir/simplify_ir.py b/python/cactus/transpiler/ir/simplify_ir.py index 97b43d6af..86854d9a7 100644 --- a/python/cactus/transpiler/ir/simplify_ir.py +++ b/python/cactus/transpiler/ir/simplify_ir.py @@ -9,9 +9,8 @@ import models def simplify(messy_ir:CVModels.LayerMap): - graph:models.CompGraph = models.CompGraph() + graph: models.Graph = models.Graph(messy_ir) - for node in messy_ir.nodes: - function_map[node.node_type](node) + From 88d944648bbe275588de5ca0689ff210ad8e27e0 Mon Sep 17 00:00:00 2001 From: Parkirat Date: Tue, 7 Jul 2026 21:54:38 -0400 Subject: [PATCH 08/84] Updates to fusions base object --- python/cactus/transpiler/ir/fusions.py | 7 +++++++ 1 file changed, 7 insertions(+) diff --git a/python/cactus/transpiler/ir/fusions.py b/python/cactus/transpiler/ir/fusions.py index bc21b138c..544213a13 100644 --- a/python/cactus/transpiler/ir/fusions.py +++ b/python/cactus/transpiler/ir/fusions.py @@ -6,6 +6,13 @@ class FusionPatter: target:str ops:tuple[str] path:tuple[int] + required_attrs:dict[str, Any] + input_refs:tuple[tuple[int]] + shared_input_refs:tuple[tuple[int]] + + def match(self): + + From ecdfdaf5f083c936e4c7ded6f88c42d95e60c294 Mon Sep 17 00:00:00 2001 From: Parkirat Date: Tue, 7 Jul 2026 23:08:45 -0400 Subject: [PATCH 09/84] Fusion matching logic + attributions cleaning --- python/cactus/transpiler/ir/constants.py | 41 +++++++++ python/cactus/transpiler/ir/fusions.py | 106 +++++++++++++++++++++-- python/cactus/transpiler/ir/models.py | 32 ++++++- 3 files changed, 169 insertions(+), 10 deletions(-) diff --git a/python/cactus/transpiler/ir/constants.py b/python/cactus/transpiler/ir/constants.py index e69de29bb..dd120da66 100644 --- a/python/cactus/transpiler/ir/constants.py +++ b/python/cactus/transpiler/ir/constants.py @@ -0,0 +1,41 @@ +LAYER_ATTRS_MAP = { + "aten._assert_tensor_metadata.default": ["size", "stride", "dtype"], + "aten._softmax.default": ["dim", "half_to_float"], + "aten._to_copy.default": [], + "aten.add.Tensor": ["scalar"], + "aten.arange.start_step": ["start", "end", "step"], + "aten.bitwise_and.Tensor": [], + "aten.bitwise_or.Tensor": [], + "aten.bmm.default": [], + "aten.cat.default": ["dim"], + "aten.clone.default": [], + "aten.cos.default": [], + "aten.cumsum.default": ["dim"], + "aten.div.Tensor": ["scalar"], + "aten.embedding.default": ["padding_idx"], + "aten.eq.Scalar": ["scalar"], + "aten.eq.Tensor": [], + "aten.expand.default": ["shape"], + "aten.full.default": ["shape", "fill_value"], + "aten.gelu.default": [], + "aten.gt.Tensor": [], + "aten.index.Tensor": [], + "aten.index_put.default": [], + "aten.le.Tensor": [], + "aten.mean.dim": ["dim", "keepdim"], + "aten.mm.default": [], + "aten.mul.Tensor": ["scalar"], + "aten.ne.Scalar": ["scalar"], + "aten.neg.default": [], + "aten.permute.default": ["dims"], + "aten.pow.Tensor_Scalar": ["exponent"], + "aten.scalar_tensor.default": ["value"], + "aten.select.int": ["dim", "index"], + "aten.sin.default": [], + "aten.slice.Tensor": ["dim", "start", "end", "step"], + "aten.sub.Tensor": ["scalar"], + "aten.tanh.default": [], + "aten.unsqueeze.default": ["dim"], + "aten.view.default": ["shape"], + "aten.where.self": [], +} diff --git a/python/cactus/transpiler/ir/fusions.py b/python/cactus/transpiler/ir/fusions.py index 544213a13..099616431 100644 --- a/python/cactus/transpiler/ir/fusions.py +++ b/python/cactus/transpiler/ir/fusions.py @@ -1,16 +1,108 @@ from dataclasses import dataclass, field from typing import Any +import constants +import models + @dataclass(slots=True) -class FusionPatter: +class FusionPattern: target:str - ops:tuple[str] - path:tuple[int] - required_attrs:dict[str, Any] - input_refs:tuple[tuple[int]] - shared_input_refs:tuple[tuple[int]] + ops:tuple[str, ...] + path:tuple[int, ...] + required_attrs:dict[int, dict[str, Any]] + input_refs:tuple[tuple[int, int],...] + shared_input_refs:tuple[tuple[tuple[int, int], tuple[int, int]], ...] + + +OPS_MAP: dict[str, list[FusionPattern]] = {} + +def len_match(fusion: FusionPattern, nodes: list[models.Node]) -> bool: + return len(fusion.ops) == len(nodes) + +def match_ops(fusion: FusionPattern, nodes: list[models.Node]) -> bool: + for i in range(len(fusion.ops)): + if(fusion.ops[i] != nodes[i].layer.target): + return False + return True + + +def match_path(fusion: FusionPattern, nodes: list[models.Node]) -> bool: + for i, parent_index in enumerate(fusion.path): + if parent_index < 0 or parent_index >= len(nodes[i].parents): + return False + + if nodes[i].parents[parent_index].layer.name != nodes[i + 1].layer.name: + return False + + return True + +def match_attrs(fusion: FusionPattern, nodes: list[models.Node]) -> bool: + for node_index, required_attrs in fusion.required_attrs.items(): + if node_index < 0 or node_index >= len(nodes): + return False + + node_attrs = nodes[node_index].normalized_attrs + for attr_name, expected_value in required_attrs.items(): + if attr_name not in node_attrs: + return False + + if node_attrs[attr_name] != expected_value: + return False + + return True + + +def match_input_refs(fusion: FusionPattern, nodes: list[models.Node]) -> bool: + for node_index, parent_index in fusion.input_refs: + if node_index < 0 or node_index >= len(nodes): + return False + + if parent_index < 0 or parent_index >= len(nodes[node_index].parents): + return False + + return True + + +def match_shared_input_refs(fusion: FusionPattern, nodes: list[models.Node]) -> bool: + for left_ref, right_ref in fusion.shared_input_refs: + left_node_index, left_parent_index = left_ref + right_node_index, right_parent_index = right_ref + + if left_node_index < 0 or left_node_index >= len(nodes): + return False + + if right_node_index < 0 or right_node_index >= len(nodes): + return False + + if left_parent_index < 0 or left_parent_index >= len(nodes[left_node_index].parents): + return False + + if right_parent_index < 0 or right_parent_index >= len(nodes[right_node_index].parents): + return False + + left_parent = nodes[left_node_index].parents[left_parent_index] + right_parent = nodes[right_node_index].parents[right_parent_index] + + if left_parent.layer.name != right_parent.layer.name: + return False + + return True + + + +def fusion_match(ops:str, nodes: list[models.Node]) -> FusionPattern | None: + if ops in OPS_MAP: + for fusion in OPS_MAP[ops]: + if match(fusion, nodes): + return fusion + + return None + + + + + - def match(self): diff --git a/python/cactus/transpiler/ir/models.py b/python/cactus/transpiler/ir/models.py index 0c8a34b1a..921c35015 100644 --- a/python/cactus/transpiler/ir/models.py +++ b/python/cactus/transpiler/ir/models.py @@ -7,13 +7,14 @@ @dataclass(slots=True) class Node: layer: CVModels.LayerRecord + normalized_attrs: dict[str, Any] parents: list["Node"] children: list["Node"] @dataclass(slots=True) class Graph: - name_map:dict[str, CVModels.LayerRecord] + name_map:dict[str, Node] source_nodes: deque[Node] consumed_ids: set[str] @@ -25,12 +26,38 @@ def from_layer_map(cls, map:CVModels.LayerMap): """###################################### MODEL UTILS!!!!!!! ######################################""" +def extract_attrs(layer_: CVModels.LayerRecord) -> dict[str, Any]: + def contains_node_ref(value: Any) -> bool: + if isinstance(value, dict): + if "node" in value: + return True + return any(contains_node_ref(item) for item in value.values()) + + if isinstance(value, list): + return any(contains_node_ref(item) for item in value) + + return False + + attrs = {} + + if isinstance(layer_.kwargs, dict): + attrs.update(layer_.kwargs) + + attr_names = constants.LAYER_ATTRS_MAP.get(layer_.target, []) + positional_attrs = [item for item in layer_.args if not contains_node_ref(item)] + + for i, attr_name in enumerate(attr_names): + if i < len(positional_attrs) and attr_name not in attrs: + attrs[attr_name] = positional_attrs[i] + + return attrs + def generate_graph(map: CVModels.LayerMap): temp_map:dict[str, Node] = {} temp_source = deque() for layer_ in map.nodes: - temp = Node(layer=layer_, parents=[], children=[]) + temp = Node(layer=layer_, normalized_attrs=extract_attrs(layer_), parents=[], children=[]) temp_map[layer_.name] = temp temp_source.append(temp) if layer_.node_type == "output" else None @@ -46,4 +73,3 @@ def generate_graph(map: CVModels.LayerMap): - From 5d66cddbc77e3237ac11b38c59f8b10d59268eaa Mon Sep 17 00:00:00 2001 From: Parkirat Date: Wed, 8 Jul 2026 11:34:35 -0400 Subject: [PATCH 10/84] Updated fusion matching logic --- python/cactus/transpiler/ir/fusions.py | 32 ++++++++++++++++++++-- python/cactus/transpiler/ir/models.py | 1 + python/cactus/transpiler/ir/simplify_ir.py | 2 ++ 3 files changed, 33 insertions(+), 2 deletions(-) diff --git a/python/cactus/transpiler/ir/fusions.py b/python/cactus/transpiler/ir/fusions.py index 099616431..0b89e2a0a 100644 --- a/python/cactus/transpiler/ir/fusions.py +++ b/python/cactus/transpiler/ir/fusions.py @@ -4,6 +4,7 @@ import models + @dataclass(slots=True) class FusionPattern: target:str @@ -13,7 +14,7 @@ class FusionPattern: input_refs:tuple[tuple[int, int],...] shared_input_refs:tuple[tuple[tuple[int, int], tuple[int, int]], ...] - +#TODO: Update with str paths and their respective fusion object OPS_MAP: dict[str, list[FusionPattern]] = {} def len_match(fusion: FusionPattern, nodes: list[models.Node]) -> bool: @@ -27,6 +28,9 @@ def match_ops(fusion: FusionPattern, nodes: list[models.Node]) -> bool: def match_path(fusion: FusionPattern, nodes: list[models.Node]) -> bool: + if len(fusion.path) != len(nodes) - 1: + return False + for i, parent_index in enumerate(fusion.path): if parent_index < 0 or parent_index >= len(nodes[i].parents): return False @@ -89,11 +93,35 @@ def match_shared_input_refs(fusion: FusionPattern, nodes: list[models.Node]) -> return True +def match_no_external_internal_children(fusion: FusionPattern, nodes: list[models.Node]) -> bool: + fused_node_names = {node.layer.name for node in nodes} + + for node in nodes[1:]: + for child in node.children: + if child.layer.name not in fused_node_names: + return False + + return True + + +def match(fusion: FusionPattern, nodes: list[models.Node]) -> bool: + return len_match(fusion, nodes) and match_attrs(fusion, nodes) and match_input_refs(fusion, nodes) and match_ops(fusion, nodes) and match_path(fusion, nodes) and match_shared_input_refs(fusion, nodes) and match_no_external_internal_children(fusion, nodes) + +#Update this with any additional matching functions if any are added +MATCHERS = ( + len_match, + match_attrs, + match_input_refs, + match_ops, + match_path, + match_shared_input_refs, + match_no_external_internal_children +) def fusion_match(ops:str, nodes: list[models.Node]) -> FusionPattern | None: if ops in OPS_MAP: for fusion in OPS_MAP[ops]: - if match(fusion, nodes): + if all(matcher(fusion, nodes) for matcher in MATCHERS): return fusion return None diff --git a/python/cactus/transpiler/ir/models.py b/python/cactus/transpiler/ir/models.py index 921c35015..d5d32cd7c 100644 --- a/python/cactus/transpiler/ir/models.py +++ b/python/cactus/transpiler/ir/models.py @@ -27,6 +27,7 @@ def from_layer_map(cls, map:CVModels.LayerMap): """###################################### MODEL UTILS!!!!!!! ######################################""" def extract_attrs(layer_: CVModels.LayerRecord) -> dict[str, Any]: + def contains_node_ref(value: Any) -> bool: if isinstance(value, dict): if "node" in value: diff --git a/python/cactus/transpiler/ir/simplify_ir.py b/python/cactus/transpiler/ir/simplify_ir.py index 86854d9a7..47a197599 100644 --- a/python/cactus/transpiler/ir/simplify_ir.py +++ b/python/cactus/transpiler/ir/simplify_ir.py @@ -14,3 +14,5 @@ def simplify(messy_ir:CVModels.LayerMap): + + From e73b35c82ad1a0ddafc960b2bdc2c4ed84d1c6c8 Mon Sep 17 00:00:00 2001 From: Parkirat Date: Wed, 8 Jul 2026 14:10:31 -0400 Subject: [PATCH 11/84] Added FusionResult object for fusions --- python/cactus/transpiler/ir/constants.py | 3 + python/cactus/transpiler/ir/fusions.py | 225 ------------------- python/cactus/transpiler/ir/fusions_utils.py | 117 ++++++++++ python/cactus/transpiler/ir/models.py | 62 +++++ python/cactus/transpiler/ir/simplify_ir.py | 123 +++++++++- 5 files changed, 301 insertions(+), 229 deletions(-) delete mode 100644 python/cactus/transpiler/ir/fusions.py create mode 100644 python/cactus/transpiler/ir/fusions_utils.py diff --git a/python/cactus/transpiler/ir/constants.py b/python/cactus/transpiler/ir/constants.py index dd120da66..1f8c9c692 100644 --- a/python/cactus/transpiler/ir/constants.py +++ b/python/cactus/transpiler/ir/constants.py @@ -39,3 +39,6 @@ "aten.view.default": ["shape"], "aten.where.self": [], } + +#Dictates maximum DFS depth during fusion search +DFS_DEPTH = 10 diff --git a/python/cactus/transpiler/ir/fusions.py b/python/cactus/transpiler/ir/fusions.py deleted file mode 100644 index 0b89e2a0a..000000000 --- a/python/cactus/transpiler/ir/fusions.py +++ /dev/null @@ -1,225 +0,0 @@ -from dataclasses import dataclass, field -from typing import Any -import constants -import models - - - -@dataclass(slots=True) -class FusionPattern: - target:str - ops:tuple[str, ...] - path:tuple[int, ...] - required_attrs:dict[int, dict[str, Any]] - input_refs:tuple[tuple[int, int],...] - shared_input_refs:tuple[tuple[tuple[int, int], tuple[int, int]], ...] - -#TODO: Update with str paths and their respective fusion object -OPS_MAP: dict[str, list[FusionPattern]] = {} - -def len_match(fusion: FusionPattern, nodes: list[models.Node]) -> bool: - return len(fusion.ops) == len(nodes) - -def match_ops(fusion: FusionPattern, nodes: list[models.Node]) -> bool: - for i in range(len(fusion.ops)): - if(fusion.ops[i] != nodes[i].layer.target): - return False - return True - - -def match_path(fusion: FusionPattern, nodes: list[models.Node]) -> bool: - if len(fusion.path) != len(nodes) - 1: - return False - - for i, parent_index in enumerate(fusion.path): - if parent_index < 0 or parent_index >= len(nodes[i].parents): - return False - - if nodes[i].parents[parent_index].layer.name != nodes[i + 1].layer.name: - return False - - return True - -def match_attrs(fusion: FusionPattern, nodes: list[models.Node]) -> bool: - for node_index, required_attrs in fusion.required_attrs.items(): - if node_index < 0 or node_index >= len(nodes): - return False - - node_attrs = nodes[node_index].normalized_attrs - for attr_name, expected_value in required_attrs.items(): - if attr_name not in node_attrs: - return False - - if node_attrs[attr_name] != expected_value: - return False - - return True - - -def match_input_refs(fusion: FusionPattern, nodes: list[models.Node]) -> bool: - for node_index, parent_index in fusion.input_refs: - if node_index < 0 or node_index >= len(nodes): - return False - - if parent_index < 0 or parent_index >= len(nodes[node_index].parents): - return False - - return True - - -def match_shared_input_refs(fusion: FusionPattern, nodes: list[models.Node]) -> bool: - for left_ref, right_ref in fusion.shared_input_refs: - left_node_index, left_parent_index = left_ref - right_node_index, right_parent_index = right_ref - - if left_node_index < 0 or left_node_index >= len(nodes): - return False - - if right_node_index < 0 or right_node_index >= len(nodes): - return False - - if left_parent_index < 0 or left_parent_index >= len(nodes[left_node_index].parents): - return False - - if right_parent_index < 0 or right_parent_index >= len(nodes[right_node_index].parents): - return False - - left_parent = nodes[left_node_index].parents[left_parent_index] - right_parent = nodes[right_node_index].parents[right_parent_index] - - if left_parent.layer.name != right_parent.layer.name: - return False - - return True - - -def match_no_external_internal_children(fusion: FusionPattern, nodes: list[models.Node]) -> bool: - fused_node_names = {node.layer.name for node in nodes} - - for node in nodes[1:]: - for child in node.children: - if child.layer.name not in fused_node_names: - return False - - return True - - -def match(fusion: FusionPattern, nodes: list[models.Node]) -> bool: - return len_match(fusion, nodes) and match_attrs(fusion, nodes) and match_input_refs(fusion, nodes) and match_ops(fusion, nodes) and match_path(fusion, nodes) and match_shared_input_refs(fusion, nodes) and match_no_external_internal_children(fusion, nodes) - -#Update this with any additional matching functions if any are added -MATCHERS = ( - len_match, - match_attrs, - match_input_refs, - match_ops, - match_path, - match_shared_input_refs, - match_no_external_internal_children -) - -def fusion_match(ops:str, nodes: list[models.Node]) -> FusionPattern | None: - if ops in OPS_MAP: - for fusion in OPS_MAP[ops]: - if all(matcher(fusion, nodes) for matcher in MATCHERS): - return fusion - - return None - - - - - - - - - - - - - -# @dataclass(frozen=True) -# class FusionPattern: -# """Description of an operation sequence that can become one fused op.""" - -# target: str -# operations: tuple[str, ...] -# path_inputs: tuple[int, ...] -# required_attrs: dict[int, dict[str, Any]] = field(default_factory=dict) -# same_inputs: tuple[tuple[InputRef, InputRef], ...] = () -# input_order: tuple[InputRef, ...] = () -# output_attrs: dict[str, AttrRef] = field(default_factory=dict) - -# @property -# def sequence(self) -> str: -# return "->".join(self.operations) - -# def __post_init__(self) -> None: -# if not self.operations: -# raise ValueError("operations cannot be empty") - -# if len(self.path_inputs) != len(self.operations) - 1: -# raise ValueError("path_inputs must contain one entry between each operation") - -# for node_index in self.required_attrs: -# if node_index < 0 or node_index >= len(self.operations): -# raise ValueError(f"required_attrs contains invalid node index {node_index}") - -# input_refs = [input_ref for input_pair in self.same_inputs for input_ref in input_pair] -# input_refs.extend(self.input_order) - -# for node_index, input_index in input_refs: -# if node_index < 0 or node_index >= len(self.operations): -# raise ValueError(f"input reference contains invalid node index {node_index}") -# if input_index < 0: -# raise ValueError("input reference indexes cannot be negative") - -# for node_index, _ in self.output_attrs.values(): -# if node_index < 0 or node_index >= len(self.operations): -# raise ValueError(f"output_attrs contains invalid node index {node_index}") - - -# @dataclass(frozen=True) -# class FusionMatch: -# """Concrete graph nodes and boundary values matched by a FusionPattern.""" - -# root_id: str -# target: str -# matched_ids: tuple[str, ...] -# external_input_ids: tuple[str, ...] -# attrs: dict[str, Any] = field(default_factory=dict) - - -# RMSNORM_FUSION = FusionPattern( -# target="rmsnorm", -# operations=( -# "aten.mul.Tensor", -# "aten.mul.Tensor", -# "aten.pow.Tensor_Scalar", -# "aten.add.Tensor", -# "aten.mean.dim", -# "aten.pow.Tensor_Scalar", -# ), -# path_inputs=(0, 1, 0, 0, 0), -# required_attrs={ -# 2: {"exponent": -0.5}, -# 3: {"scalar": 1e-06}, -# 4: {"dim": [-1], "keepdim": True}, -# 5: {"exponent": 2}, -# }, -# same_inputs=(((1, 0), (5, 0)),), -# input_order=((1, 0), (0, 1)), -# output_attrs={"epsilon": (3, "scalar")}, -# ) - - -# LINEAR_RESHAPE_FUSION = FusionPattern( -# target="linear_reshape", -# operations=( -# "aten.view.default", -# "aten.mm.default", -# ), -# path_inputs=(0,), -# input_order=((1, 0), (1, 1)), -# output_attrs={"shape": (0, "shape")}, -# ) diff --git a/python/cactus/transpiler/ir/fusions_utils.py b/python/cactus/transpiler/ir/fusions_utils.py new file mode 100644 index 000000000..6b6af71b7 --- /dev/null +++ b/python/cactus/transpiler/ir/fusions_utils.py @@ -0,0 +1,117 @@ +import models +#TODO: Update with str paths and their respective fusion object +OPS_MAP: dict[str, list[models.FusionPattern]] = {} + + + +"""###################################### FUSION UTILS!!!!!!! ######################################""" + + +def len_match(fusion: models.FusionPattern, nodes: list[models.Node]) -> bool: + return len(fusion.ops) == len(nodes) + +def match_ops(fusion: models.FusionPattern, nodes: list[models.Node]) -> bool: + for i in range(len(fusion.ops)): + if(fusion.ops[i] != nodes[i].layer.target): + return False + return True + + +def match_path(fusion: models.FusionPattern, nodes: list[models.Node]) -> bool: + if len(fusion.path) != len(nodes) - 1: + return False + + for i, parent_index in enumerate(fusion.path): + if parent_index < 0 or parent_index >= len(nodes[i].parents): + return False + + if nodes[i].parents[parent_index].layer.name != nodes[i + 1].layer.name: + return False + + return True + +def match_attrs(fusion: models.FusionPattern, nodes: list[models.Node]) -> bool: + for node_index, required_attrs in fusion.required_attrs.items(): + if node_index < 0 or node_index >= len(nodes): + return False + + node_attrs = nodes[node_index].normalized_attrs + for attr_name, expected_value in required_attrs.items(): + if attr_name not in node_attrs: + return False + + if node_attrs[attr_name] != expected_value: + return False + + return True + + +def match_input_refs(fusion: models.FusionPattern, nodes: list[models.Node]) -> bool: + for node_index, parent_index in fusion.input_refs: + if node_index < 0 or node_index >= len(nodes): + return False + + if parent_index < 0 or parent_index >= len(nodes[node_index].parents): + return False + + return True + + +def match_shared_input_refs(fusion: models.FusionPattern, nodes: list[models.Node]) -> bool: + for left_ref, right_ref in fusion.shared_input_refs: + left_node_index, left_parent_index = left_ref + right_node_index, right_parent_index = right_ref + + if left_node_index < 0 or left_node_index >= len(nodes): + return False + + if right_node_index < 0 or right_node_index >= len(nodes): + return False + + if left_parent_index < 0 or left_parent_index >= len(nodes[left_node_index].parents): + return False + + if right_parent_index < 0 or right_parent_index >= len(nodes[right_node_index].parents): + return False + + left_parent = nodes[left_node_index].parents[left_parent_index] + right_parent = nodes[right_node_index].parents[right_parent_index] + + if left_parent.layer.name != right_parent.layer.name: + return False + + return True + + +def match_no_external_internal_children(fusion: models.FusionPattern, nodes: list[models.Node]) -> bool: + fused_node_names = {node.layer.name for node in nodes} + + for node in nodes[1:]: + for child in node.children: + if child.layer.name not in fused_node_names: + return False + + return True + + +def match(fusion: models.FusionPattern, nodes: list[models.Node]) -> bool: + return len_match(fusion, nodes) and match_attrs(fusion, nodes) and match_input_refs(fusion, nodes) and match_ops(fusion, nodes) and match_path(fusion, nodes) and match_shared_input_refs(fusion, nodes) and match_no_external_internal_children(fusion, nodes) + +#Update this with any additional matching functions if any are added +MATCHERS = ( + len_match, + match_attrs, + match_input_refs, + match_ops, + match_path, + match_shared_input_refs, + match_no_external_internal_children +) + +def fusion_match(ops:str, nodes: list[models.Node]) -> models.FusionPattern | None: + if ops in OPS_MAP: + for fusion in OPS_MAP[ops]: + if all(matcher(fusion, nodes) for matcher in MATCHERS): + return fusion + + return None \ No newline at end of file diff --git a/python/cactus/transpiler/ir/models.py b/python/cactus/transpiler/ir/models.py index d5d32cd7c..cf3440061 100644 --- a/python/cactus/transpiler/ir/models.py +++ b/python/cactus/transpiler/ir/models.py @@ -22,10 +22,72 @@ class Graph: def from_layer_map(cls, map:CVModels.LayerMap): return generate_graph(map) +@dataclass(slots=True) +class FusionPattern: + #Input requirements + target:str + ops:tuple[str, ...] + path:tuple[int, ...] + required_attrs:dict[int, dict[str, Any]] + input_refs:tuple[tuple[int, int],...] + shared_input_refs:tuple[tuple[tuple[int, int], tuple[int, int]], ...] + + #Output requirements + output_attrs: dict[str, Any] + + +@dataclass(slots=True) +class FusionResult: + fusion: FusionPattern + matched_nodes: list[Node] + external_input_nodes: list[Node] + node: Node + + @property + def matched_node_ids(self) -> list[str]: + return [node.layer.name for node in self.matched_nodes] + + @classmethod + def from_match(cls, fusion: FusionPattern, nodes: list[Node]) -> "FusionResult": + return generate_fusion(fusion, nodes) + """###################################### MODEL UTILS!!!!!!! ######################################""" +def generate_fusion(fusion: FusionPattern, nodes: list[Node]) -> FusionResult: + external_input_nodes = [nodes[node_index].parents[parent_index] for node_index, parent_index in fusion.input_refs] + + fused_name = f"{fusion.target}_{nodes[0].layer.name}".replace(".", "_") + + fused_layer = CVModels.LayerRecord( + index=nodes[0].layer.index, + name=fused_name, + node_type="call_function", + target=fusion.target, + args=[{"node": node.layer.name} for node in external_input_nodes], + kwargs={}, + users=nodes[0].layer.users, + tensor_output_meta=nodes[0].layer.tensor_output_meta, + module_stack=nodes[0].layer.module_stack, + ) + + fused_node = Node( + layer=fused_layer, + normalized_attrs= fusion.output_attrs, + parents=external_input_nodes, + children=nodes[0].children + ) + + + return FusionResult( + fusion=fusion, + matched_nodes=nodes, + external_input_nodes=external_input_nodes, + node=fused_node + ) + + def extract_attrs(layer_: CVModels.LayerRecord) -> dict[str, Any]: def contains_node_ref(value: Any) -> bool: diff --git a/python/cactus/transpiler/ir/simplify_ir.py b/python/cactus/transpiler/ir/simplify_ir.py index 47a197599..7b4b23adf 100644 --- a/python/cactus/transpiler/ir/simplify_ir.py +++ b/python/cactus/transpiler/ir/simplify_ir.py @@ -4,15 +4,130 @@ on the DAG to generate an in-order execution list. """ -import models, constants +from dataclasses import dataclass +import models, python.cactus.transpiler.ir.fusions_utils as fusions_utils import converter.models as CVModels -import models -def simplify(messy_ir:CVModels.LayerMap): - graph: models.Graph = models.Graph(messy_ir) +@dataclass(slots=True) +class FusionResult: + fusion: fusions_utils.FusionPattern + matched_nodes: list[models.Node] + external_input_nodes: list[models.Node] + fused_node: models.Node + + @property + def matched_node_ids(self) -> list[str]: + return [node.layer.name for node in self.matched_nodes] + + @classmethod + def from_match(cls, fusion: fusions_utils.FusionPattern, nodes: list[models.Node]) -> "FusionResult": + external_input_nodes = [nodes[node_index].parents[parent_index] for node_index, parent_index in fusion.input_refs] + + fused_name = f"{fusion.target}_{nodes[0].layer.name}".replace(".", "_") + fused_layer = CVModels.LayerRecord( + index=nodes[0].layer.index, + name=fused_name, + node_type="call_function", + target=fusion.target, + args=[{"node": node.layer.name} for node in external_input_nodes], + kwargs={}, + users=nodes[0].layer.users, + tensor_output_meta=nodes[0].layer.tensor_output_meta, + module_stack=nodes[0].layer.module_stack, + ) + + return cls( + fusion=fusion, + matched_nodes=nodes, + external_input_nodes=external_input_nodes, + fused_node=models.Node( + layer=fused_layer, + normalized_attrs={}, + parents=external_input_nodes, + children=nodes[0].children, + ), + ) + + +def try_match_from_node(node: models.Node, graph: models.Graph, max_depth: int = 10) -> FusionResult | None: + if node.layer.name in graph.consumed_ids: + return None + + candidate_paths: list[list[models.Node]] = [] + + def collect_paths(path: list[models.Node]) -> None: + candidate_paths.append(path) + + if len(path) >= max_depth: + return + + for parent in path[-1].parents: + if parent.layer.name in graph.consumed_ids: + continue + + if any(existing.layer.name == parent.layer.name for existing in path): + continue + + collect_paths([*path, parent]) + + collect_paths([node]) + + for candidate in sorted(candidate_paths, key=len, reverse=True): + ops = "->".join(candidate_node.layer.target for candidate_node in candidate) + if ops not in fusions_utils.OPS_MAP: + continue + + fusion = fusions_utils.fusion_match(ops, candidate) + if fusion is not None: + return FusionResult.from_match(fusion, candidate) + + return None + + +def rev_top_sort(graph:models.Graph) -> models.Graph: + """Option 1: + Take the source node + Do DFS style traversal of that source node its parents + Max depth of dfs = 10 + Check if path potential fusion (does string exist as keys in dict) + If fused, add the og layers to set of fused nodes + Layers in set of fused nodes will not be proposed in fusions again + """ + worklist = list(graph.source_nodes) + visited = set() + fused = [] + + while worklist: + node = worklist.pop() + + if node.layer.name in visited: + continue + visited.add(node.layer.name) + + if node.layer.name in graph.consumed_ids: + continue + + match = try_match_from_node(node, graph) + + if match: + fused.append(match) + graph.consumed_ids.update(match.matched_node_ids) + worklist.extend(match.external_input_nodes) + else: + worklist.extend(node.parents) + + return graph + + + + +def simplify(messy_ir:CVModels.LayerMap): + graph: models.Graph = models.Graph.from_layer_map(messy_ir) + + From 6b2069bea764bf667a4864f3d93e7f97ab77346b Mon Sep 17 00:00:00 2001 From: Parkirat Date: Wed, 8 Jul 2026 16:04:16 -0400 Subject: [PATCH 12/84] Added final fused graph clss --- python/cactus/transpiler/ir/fusions_utils.py | 40 +++++- python/cactus/transpiler/ir/models.py | 24 +++- python/cactus/transpiler/ir/simplify_ir.py | 140 ++++++++----------- 3 files changed, 108 insertions(+), 96 deletions(-) diff --git a/python/cactus/transpiler/ir/fusions_utils.py b/python/cactus/transpiler/ir/fusions_utils.py index 6b6af71b7..317fe52c7 100644 --- a/python/cactus/transpiler/ir/fusions_utils.py +++ b/python/cactus/transpiler/ir/fusions_utils.py @@ -1,4 +1,4 @@ -import models +import models, constants #TODO: Update with str paths and their respective fusion object OPS_MAP: dict[str, list[models.FusionPattern]] = {} @@ -93,9 +93,17 @@ def match_no_external_internal_children(fusion: models.FusionPattern, nodes: lis return True +def match_external_inputs(fusion: models.FusionPattern, nodes: list[models.Node]): + proposed_fusion_node_names = {node.layer.name for node in nodes} -def match(fusion: models.FusionPattern, nodes: list[models.Node]) -> bool: - return len_match(fusion, nodes) and match_attrs(fusion, nodes) and match_input_refs(fusion, nodes) and match_ops(fusion, nodes) and match_path(fusion, nodes) and match_shared_input_refs(fusion, nodes) and match_no_external_internal_children(fusion, nodes) + declared_inputs = {nodes[node_index].parents[parent_index].layer.name for node_index, parent_index in fusion.input_refs} + + for node in nodes: + for parent in node.parents: + if parent.layer.name not in proposed_fusion_node_names and parent.layer.name not in declared_inputs: + return False + + return True #Update this with any additional matching functions if any are added MATCHERS = ( @@ -105,7 +113,8 @@ def match(fusion: models.FusionPattern, nodes: list[models.Node]) -> bool: match_ops, match_path, match_shared_input_refs, - match_no_external_internal_children + match_no_external_internal_children, + match_external_inputs, ) def fusion_match(ops:str, nodes: list[models.Node]) -> models.FusionPattern | None: @@ -114,4 +123,25 @@ def fusion_match(ops:str, nodes: list[models.Node]) -> models.FusionPattern | No if all(matcher(fusion, nodes) for matcher in MATCHERS): return fusion - return None \ No newline at end of file + return None + +"""###################################### FUSION GRAPH TRAVERSAL UTILS!!!!!!! ######################################""" +def dfs_path_gen(path:list[models.Node], graph:models.Graph, max_depth: int = constants.DFS_DEPTH): + if len(path) >= max_depth or not path[-1].parents: + yield path + return + + added_to_path = False + + for parent in path[-1].parents: + if parent.layer.name in graph.consumed_ids: + continue + + if any(existing.layer.name == parent.layer.name for existing in path): + continue + + added_to_path = True + yield from dfs_path_gen([*path, parent], graph, max_depth) + + if not added_to_path: + yield path diff --git a/python/cactus/transpiler/ir/models.py b/python/cactus/transpiler/ir/models.py index cf3440061..1d92baa42 100644 --- a/python/cactus/transpiler/ir/models.py +++ b/python/cactus/transpiler/ir/models.py @@ -15,13 +15,14 @@ class Node: @dataclass(slots=True) class Graph: name_map:dict[str, Node] - source_nodes: deque[Node] + source_nodes: list[Node] consumed_ids: set[str] @classmethod - def from_layer_map(cls, map:CVModels.LayerMap): + def from_layer_map(cls, map:CVModels.LayerMap) -> "Graph": return generate_graph(map) - + + @dataclass(slots=True) class FusionPattern: #Input requirements @@ -50,7 +51,19 @@ def matched_node_ids(self) -> list[str]: @classmethod def from_match(cls, fusion: FusionPattern, nodes: list[Node]) -> "FusionResult": return generate_fusion(fusion, nodes) - + + +@dataclass(slots=True) +class SimplifiedGraph: + raw_graph: Graph + name_map: dict[str, Node] + source_nodes: list[Node] + replacement_map: dict[str, str] + fused_results: list[FusionResult] + + @classmethod + def from_graph(cls, graph: Graph) -> "SimplifiedGraph": + return cls(raw_graph=graph, name_map={}, source_nodes=[], replacement_map={}, fused_results=[]) """###################################### MODEL UTILS!!!!!!! ######################################""" @@ -117,7 +130,7 @@ def contains_node_ref(value: Any) -> bool: def generate_graph(map: CVModels.LayerMap): temp_map:dict[str, Node] = {} - temp_source = deque() + temp_source = [] for layer_ in map.nodes: temp = Node(layer=layer_, normalized_attrs=extract_attrs(layer_), parents=[], children=[]) @@ -135,4 +148,3 @@ def generate_graph(map: CVModels.LayerMap): return Graph(name_map=temp_map, source_nodes=temp_source, consumed_ids={}) - diff --git a/python/cactus/transpiler/ir/simplify_ir.py b/python/cactus/transpiler/ir/simplify_ir.py index 7b4b23adf..00dc75c64 100644 --- a/python/cactus/transpiler/ir/simplify_ir.py +++ b/python/cactus/transpiler/ir/simplify_ir.py @@ -4,120 +4,90 @@ on the DAG to generate an in-order execution list. """ -from dataclasses import dataclass -import models, python.cactus.transpiler.ir.fusions_utils as fusions_utils +import models, fusions_utils import converter.models as CVModels -@dataclass(slots=True) -class FusionResult: - fusion: fusions_utils.FusionPattern - matched_nodes: list[models.Node] - external_input_nodes: list[models.Node] - fused_node: models.Node - - @property - def matched_node_ids(self) -> list[str]: - return [node.layer.name for node in self.matched_nodes] - - @classmethod - def from_match(cls, fusion: fusions_utils.FusionPattern, nodes: list[models.Node]) -> "FusionResult": - external_input_nodes = [nodes[node_index].parents[parent_index] for node_index, parent_index in fusion.input_refs] - - fused_name = f"{fusion.target}_{nodes[0].layer.name}".replace(".", "_") - fused_layer = CVModels.LayerRecord( - index=nodes[0].layer.index, - name=fused_name, - node_type="call_function", - target=fusion.target, - args=[{"node": node.layer.name} for node in external_input_nodes], - kwargs={}, - users=nodes[0].layer.users, - tensor_output_meta=nodes[0].layer.tensor_output_meta, - module_stack=nodes[0].layer.module_stack, - ) - - return cls( - fusion=fusion, - matched_nodes=nodes, - external_input_nodes=external_input_nodes, - fused_node=models.Node( - layer=fused_layer, - normalized_attrs={}, - parents=external_input_nodes, - children=nodes[0].children, - ), - ) - - -def try_match_from_node(node: models.Node, graph: models.Graph, max_depth: int = 10) -> FusionResult | None: +def try_match_from_node(node: models.Node, graph: models.Graph) -> models.FusionResult | None: if node.layer.name in graph.consumed_ids: return None + + best_match = None + best_match_len = 0 + checked_candidates = set() - candidate_paths: list[list[models.Node]] = [] - - def collect_paths(path: list[models.Node]) -> None: - candidate_paths.append(path) - - if len(path) >= max_depth: - return + for path in fusions_utils.dfs_path_gen([node], graph): + for curr_len in range(len(path), 0, -1): + candidate_path = path[:curr_len] + candidate_ids = tuple(candidate_node.layer.name for candidate_node in candidate_path) - for parent in path[-1].parents: - if parent.layer.name in graph.consumed_ids: + if candidate_ids in checked_candidates: continue + checked_candidates.add(candidate_ids) - if any(existing.layer.name == parent.layer.name for existing in path): + ops_key = "->".join(candidate_node.layer.target for candidate_node in candidate_path) + if ops_key not in fusions_utils.OPS_MAP: continue - collect_paths([*path, parent]) - - collect_paths([node]) - - for candidate in sorted(candidate_paths, key=len, reverse=True): - ops = "->".join(candidate_node.layer.target for candidate_node in candidate) - if ops not in fusions_utils.OPS_MAP: - continue + fusion = fusions_utils.fusion_match(ops_key, candidate_path) + if fusion is None: + continue - fusion = fusions_utils.fusion_match(ops, candidate) - if fusion is not None: - return FusionResult.from_match(fusion, candidate) + if curr_len > best_match_len: + best_match = models.FusionResult.from_match(fusion, candidate_path) + best_match_len = curr_len - return None + if best_match_len == len(path): + return best_match + + return best_match -def rev_top_sort(graph:models.Graph) -> models.Graph: - """Option 1: - Take the source node - Do DFS style traversal of that source node its parents - Max depth of dfs = 10 - Check if path potential fusion (does string exist as keys in dict) - If fused, add the og layers to set of fused nodes - Layers in set of fused nodes will not be proposed in fusions again - """ +def rev_top_sort(graph:models.Graph) -> models.SimplifiedGraph: worklist = list(graph.source_nodes) visited = set() - fused = [] + fused_results = [] while worklist: node = worklist.pop() + node_id = node.layer.name - if node.layer.name in visited: + if node_id in visited: continue - visited.add(node.layer.name) + visited.add(node_id) - if node.layer.name in graph.consumed_ids: + if node_id in graph.consumed_ids: continue match = try_match_from_node(node, graph) - if match: - fused.append(match) + if match is not None: + fused_results.append(match) graph.consumed_ids.update(match.matched_node_ids) - worklist.extend(match.external_input_nodes) - else: - worklist.extend(node.parents) + worklist.extend(input_node for input_node in match.external_input_nodes if input_node.layer.name not in graph.consumed_ids) + continue + + worklist.extend(parent for parent in node.parents if parent.layer.name not in graph.consumed_ids) + + simplified_graph = models.SimplifiedGraph.from_graph(graph) + simplified_graph.fused_results = fused_results + + for match in fused_results: + simplified_graph.name_map[match.node.layer.name] = match.node + for node_id in match.matched_node_ids: + simplified_graph.replacement_map[node_id] = match.node.layer.name + + for node_id, node in graph.name_map.items(): + if node_id not in graph.consumed_ids: + simplified_graph.name_map[node_id] = node + + simplified_graph.source_nodes = [ + simplified_graph.name_map[simplified_graph.replacement_map.get(node.layer.name, node.layer.name)] + for node in graph.source_nodes + if simplified_graph.replacement_map.get(node.layer.name, node.layer.name) in simplified_graph.name_map + ] - return graph + return simplified_graph From ea0207afed4fabbc786289f99180925e8db86352 Mon Sep 17 00:00:00 2001 From: Parkirat Date: Wed, 8 Jul 2026 16:06:59 -0400 Subject: [PATCH 13/84] Added simplified graph representation + topo sort on simplified graph --- python/cactus/transpiler/ir/simplify_ir.py | 133 ++++++++++++++++++++- 1 file changed, 131 insertions(+), 2 deletions(-) diff --git a/python/cactus/transpiler/ir/simplify_ir.py b/python/cactus/transpiler/ir/simplify_ir.py index 00dc75c64..287693372 100644 --- a/python/cactus/transpiler/ir/simplify_ir.py +++ b/python/cactus/transpiler/ir/simplify_ir.py @@ -90,11 +90,140 @@ def rev_top_sort(graph:models.Graph) -> models.SimplifiedGraph: return simplified_graph +def resolve_name(name: str, replacement_map: dict[str, str]) -> str: + seen = set() + while name in replacement_map and name not in seen: + seen.add(name) + name = replacement_map[name] -def simplify(messy_ir:CVModels.LayerMap): + return name + + +def rewrite_node_refs(value, replacement_map: dict[str, str]): + if isinstance(value, dict): + rewritten = {key: rewrite_node_refs(item, replacement_map) for key, item in value.items()} + if "node" in rewritten: + rewritten["node"] = resolve_name(rewritten["node"], replacement_map) + return rewritten + + if isinstance(value, list): + return [rewrite_node_refs(item, replacement_map) for item in value] + + return value + + +def extract_node_refs(value) -> list[str]: + if isinstance(value, dict): + refs = [] + if "node" in value: + refs.append(value["node"]) + + for item in value.values(): + refs.extend(extract_node_refs(item)) + + return refs + + if isinstance(value, list): + refs = [] + for item in value: + refs.extend(extract_node_refs(item)) + return refs + + return [] + + +def topo_sort(graph: models.SimplifiedGraph, args_map: dict[str, object], kwargs_map: dict[str, object]) -> list[str]: + deps = {} + users = {node_id: set() for node_id in graph.name_map} + + for node_id in graph.name_map: + node_refs = extract_node_refs(args_map[node_id]) + extract_node_refs(kwargs_map[node_id]) + deps[node_id] = { + ref + for ref in node_refs + if ref in graph.name_map and ref != node_id + } + + for dep in deps[node_id]: + users[dep].add(node_id) + + ready = sorted( + [node_id for node_id, node_deps in deps.items() if not node_deps], + key=lambda node_id: graph.name_map[node_id].layer.index, + ) + ordered = [] + + while ready: + node_id = ready.pop(0) + ordered.append(node_id) + + for user_id in sorted(users[node_id], key=lambda item: graph.name_map[item].layer.index): + deps[user_id].remove(node_id) + if not deps[user_id]: + ready.append(user_id) + + ready.sort(key=lambda item: graph.name_map[item].layer.index) + + if len(ordered) != len(graph.name_map): + raise ValueError("Simplified graph contains a cycle or unresolved dependency") + + return ordered + + +def layer_map_from_simplified_graph(graph: models.SimplifiedGraph, original: CVModels.LayerMap) -> CVModels.LayerMap: + fused_node_ids = {result.node.layer.name for result in graph.fused_results} + args_map = {} + kwargs_map = {} + + for node_id, node in graph.name_map.items(): + args_map[node_id] = rewrite_node_refs(node.layer.args, graph.replacement_map) + + if node_id in fused_node_ids: + kwargs_map[node_id] = rewrite_node_refs(node.normalized_attrs, graph.replacement_map) + else: + kwargs_map[node_id] = rewrite_node_refs(node.layer.kwargs, graph.replacement_map) + + ordered_node_ids = topo_sort(graph, args_map, kwargs_map) + users_map = {node_id: [] for node_id in graph.name_map} + + for node_id in ordered_node_ids: + node_refs = extract_node_refs(args_map[node_id]) + extract_node_refs(kwargs_map[node_id]) + for ref in node_refs: + if ref in users_map and node_id not in users_map[ref]: + users_map[ref].append(node_id) + + nodes = [] + for index, node_id in enumerate(ordered_node_ids): + node = graph.name_map[node_id] + nodes.append( + CVModels.LayerRecord( + index=index, + name=node.layer.name, + node_type=node.layer.node_type, + target=node.layer.target, + args=args_map[node_id], + kwargs=kwargs_map[node_id], + users=users_map[node_id], + tensor_output_meta=node.layer.tensor_output_meta, + module_stack=node.layer.module_stack, + ) + ) + + + return CVModels.LayerMap( + model_name=original.model_name, + task=original.task, + graph_signature=original.graph_signature, + range_constants=original.range_constants, + nodes=nodes, + ) + + +def simplify(messy_ir:CVModels.LayerMap) -> CVModels.LayerMap: graph: models.Graph = models.Graph.from_layer_map(messy_ir) - + simplified_graph = rev_top_sort(graph) + return layer_map_from_simplified_graph(simplified_graph, messy_ir) From c9876a81e521609490ca0dffcd2bbb0d34e42f28 Mon Sep 17 00:00:00 2001 From: Parkirat Date: Thu, 9 Jul 2026 09:55:40 -0400 Subject: [PATCH 14/84] DAG based fusion: Initial models refactoring --- python/cactus/transpiler/ir/fusions.py | 26 +++ python/cactus/transpiler/ir/fusions_utils.py | 137 ++++++------ python/cactus/transpiler/ir/models.py | 122 ++++------- python/cactus/transpiler/ir/nodes.py | 2 + python/cactus/transpiler/ir/simplify_ir.py | 219 ++++++++++--------- 5 files changed, 256 insertions(+), 250 deletions(-) create mode 100644 python/cactus/transpiler/ir/fusions.py create mode 100644 python/cactus/transpiler/ir/nodes.py diff --git a/python/cactus/transpiler/ir/fusions.py b/python/cactus/transpiler/ir/fusions.py new file mode 100644 index 000000000..09df5428a --- /dev/null +++ b/python/cactus/transpiler/ir/fusions.py @@ -0,0 +1,26 @@ +import models, nodes + + +RMS_NORM = models.FusionGraph( + target="rms_norm", + root="weighted_output", + edges=( + models.Edge("weighted_output", 0, "normalized_x"), + models.Edge("normalized_x", 1, "inverse_rms"), + models.Edge("inverse_rms", 0, "variance_with_eps"), + models.Edge("variance_with_eps", 0, "variance"), + models.Edge("variance", 0, "squared_x"), + ), +) + + +ALL_FUSIONS = ( + RMS_NORM, +) + + +ROOT_TARGET_MAP = { + "aten.mul.Tensor": ( + RMS_NORM, + ), +} diff --git a/python/cactus/transpiler/ir/fusions_utils.py b/python/cactus/transpiler/ir/fusions_utils.py index 317fe52c7..0449d807d 100644 --- a/python/cactus/transpiler/ir/fusions_utils.py +++ b/python/cactus/transpiler/ir/fusions_utils.py @@ -7,103 +7,106 @@ """###################################### FUSION UTILS!!!!!!! ######################################""" -def len_match(fusion: models.FusionPattern, nodes: list[models.Node]) -> bool: - return len(fusion.ops) == len(nodes) -def match_ops(fusion: models.FusionPattern, nodes: list[models.Node]) -> bool: - for i in range(len(fusion.ops)): - if(fusion.ops[i] != nodes[i].layer.target): - return False - return True +# def len_match(fusion: models.FusionPattern, nodes: list[models.Node]) -> bool: +# return len(fusion.ops) == len(nodes) +# def match_ops(fusion: models.FusionPattern, nodes: list[models.Node]) -> bool: +# for i in range(len(fusion.ops)): +# if(fusion.ops[i] != nodes[i].layer.target): +# return False +# return True -def match_path(fusion: models.FusionPattern, nodes: list[models.Node]) -> bool: - if len(fusion.path) != len(nodes) - 1: - return False - for i, parent_index in enumerate(fusion.path): - if parent_index < 0 or parent_index >= len(nodes[i].parents): - return False +# def match_path(fusion: models.FusionPattern, nodes: list[models.Node]) -> bool: +# if len(fusion.path) != len(nodes) - 1: +# return False - if nodes[i].parents[parent_index].layer.name != nodes[i + 1].layer.name: - return False +# for i, parent_index in enumerate(fusion.path): +# if parent_index < 0 or parent_index >= len(nodes[i].parents): +# return False - return True +# if nodes[i].parents[parent_index].layer.name != nodes[i + 1].layer.name: +# return False -def match_attrs(fusion: models.FusionPattern, nodes: list[models.Node]) -> bool: - for node_index, required_attrs in fusion.required_attrs.items(): - if node_index < 0 or node_index >= len(nodes): - return False +# return True - node_attrs = nodes[node_index].normalized_attrs - for attr_name, expected_value in required_attrs.items(): - if attr_name not in node_attrs: - return False +# def match_attrs(fusion: models.FusionPattern, nodes: list[models.Node]) -> bool: +# for node_index, required_attrs in fusion.required_attrs.items(): +# if node_index < 0 or node_index >= len(nodes): +# return False - if node_attrs[attr_name] != expected_value: - return False +# node_attrs = nodes[node_index].normalized_attrs +# for attr_name, expected_value in required_attrs.items(): +# if attr_name not in node_attrs: +# return False - return True +# if node_attrs[attr_name] != expected_value: +# return False +# return True -def match_input_refs(fusion: models.FusionPattern, nodes: list[models.Node]) -> bool: - for node_index, parent_index in fusion.input_refs: - if node_index < 0 or node_index >= len(nodes): - return False - if parent_index < 0 or parent_index >= len(nodes[node_index].parents): - return False +# def match_input_refs(fusion: models.FusionPattern, nodes: list[models.Node]) -> bool: +# for node_index, parent_index in fusion.input_refs: +# if node_index < 0 or node_index >= len(nodes): +# return False - return True +# if parent_index < 0 or parent_index >= len(nodes[node_index].parents): +# return False +# return True -def match_shared_input_refs(fusion: models.FusionPattern, nodes: list[models.Node]) -> bool: - for left_ref, right_ref in fusion.shared_input_refs: - left_node_index, left_parent_index = left_ref - right_node_index, right_parent_index = right_ref - if left_node_index < 0 or left_node_index >= len(nodes): - return False +# def match_shared_input_refs(fusion: models.FusionPattern, nodes: list[models.Node]) -> bool: +# for left_ref, right_ref in fusion.shared_input_refs: +# left_node_index, left_parent_index = left_ref +# right_node_index, right_parent_index = right_ref - if right_node_index < 0 or right_node_index >= len(nodes): - return False +# if left_node_index < 0 or left_node_index >= len(nodes): +# return False - if left_parent_index < 0 or left_parent_index >= len(nodes[left_node_index].parents): - return False +# if right_node_index < 0 or right_node_index >= len(nodes): +# return False - if right_parent_index < 0 or right_parent_index >= len(nodes[right_node_index].parents): - return False +# if left_parent_index < 0 or left_parent_index >= len(nodes[left_node_index].parents): +# return False - left_parent = nodes[left_node_index].parents[left_parent_index] - right_parent = nodes[right_node_index].parents[right_parent_index] +# if right_parent_index < 0 or right_parent_index >= len(nodes[right_node_index].parents): +# return False - if left_parent.layer.name != right_parent.layer.name: - return False +# left_parent = nodes[left_node_index].parents[left_parent_index] +# right_parent = nodes[right_node_index].parents[right_parent_index] - return True +# if left_parent.layer.name != right_parent.layer.name: +# return False +# return True -def match_no_external_internal_children(fusion: models.FusionPattern, nodes: list[models.Node]) -> bool: - fused_node_names = {node.layer.name for node in nodes} - for node in nodes[1:]: - for child in node.children: - if child.layer.name not in fused_node_names: - return False +# def match_no_external_internal_children(fusion: models.FusionPattern, nodes: list[models.Node]) -> bool: +# fused_node_names = {node.layer.name for node in nodes} - return True +# for node in nodes[1:]: +# for child in node.children: +# if child.layer.name not in fused_node_names: +# return False -def match_external_inputs(fusion: models.FusionPattern, nodes: list[models.Node]): - proposed_fusion_node_names = {node.layer.name for node in nodes} +# return True - declared_inputs = {nodes[node_index].parents[parent_index].layer.name for node_index, parent_index in fusion.input_refs} +# def match_external_inputs(fusion: models.FusionPattern, nodes: list[models.Node]): +# proposed_fusion_node_names = {node.layer.name for node in nodes} - for node in nodes: - for parent in node.parents: - if parent.layer.name not in proposed_fusion_node_names and parent.layer.name not in declared_inputs: - return False +# declared_inputs = {nodes[node_index].parents[parent_index].layer.name for node_index, parent_index in fusion.input_refs} + +# for node in nodes: +# for parent in node.parents: +# if parent.layer.name not in proposed_fusion_node_names and parent.layer.name not in declared_inputs: +# return False - return True +# return True + + #Update this with any additional matching functions if any are added MATCHERS = ( @@ -125,6 +128,8 @@ def fusion_match(ops:str, nodes: list[models.Node]) -> models.FusionPattern | No return None + + """###################################### FUSION GRAPH TRAVERSAL UTILS!!!!!!! ######################################""" def dfs_path_gen(path:list[models.Node], graph:models.Graph, max_depth: int = constants.DFS_DEPTH): if len(path) >= max_depth or not path[-1].parents: diff --git a/python/cactus/transpiler/ir/models.py b/python/cactus/transpiler/ir/models.py index 1d92baa42..7b2710f94 100644 --- a/python/cactus/transpiler/ir/models.py +++ b/python/cactus/transpiler/ir/models.py @@ -1,12 +1,12 @@ from dataclasses import dataclass -from typing import Any -from collections import deque +from typing import Any, Optional import converter.models as CVModels import constants @dataclass(slots=True) class Node: - layer: CVModels.LayerRecord + layer: Optional[CVModels.LayerRecord] + underlying_op: str normalized_attrs: dict[str, Any] parents: list["Node"] children: list["Node"] @@ -22,84 +22,46 @@ class Graph: def from_layer_map(cls, map:CVModels.LayerMap) -> "Graph": return generate_graph(map) - -@dataclass(slots=True) -class FusionPattern: - #Input requirements - target:str - ops:tuple[str, ...] - path:tuple[int, ...] - required_attrs:dict[int, dict[str, Any]] - input_refs:tuple[tuple[int, int],...] - shared_input_refs:tuple[tuple[tuple[int, int], tuple[int, int]], ...] +@dataclass(slots = True) +class NodeSpec: + underlying_op:str + required_attrs: dict[str, Any] - #Output requirements - output_attrs: dict[str, Any] +@dataclass(slots = True) +class Edge: + from_node:str + parent_index:int + to_node:str - @dataclass(slots=True) -class FusionResult: - fusion: FusionPattern - matched_nodes: list[Node] - external_input_nodes: list[Node] - node: Node +class FusionGraph: + target:str + root:str + edges:tuple[Edge] + - @property - def matched_node_ids(self) -> list[str]: - return [node.layer.name for node in self.matched_nodes] - @classmethod - def from_match(cls, fusion: FusionPattern, nodes: list[Node]) -> "FusionResult": - return generate_fusion(fusion, nodes) +"""###################################### MODEL UTILS!!!!!!! ######################################""" +def extract_node_names(value: Any) -> list[str]: + + if isinstance(value, dict): + refs = [] + if "node" in value: + refs.append(value["node"]) -@dataclass(slots=True) -class SimplifiedGraph: - raw_graph: Graph - name_map: dict[str, Node] - source_nodes: list[Node] - replacement_map: dict[str, str] - fused_results: list[FusionResult] + for item in value.values(): + refs.extend(extract_node_names(item)) - @classmethod - def from_graph(cls, graph: Graph) -> "SimplifiedGraph": - return cls(raw_graph=graph, name_map={}, source_nodes=[], replacement_map={}, fused_results=[]) - + return refs -"""###################################### MODEL UTILS!!!!!!! ######################################""" + if isinstance(value, list): + refs = [] + for item in value: + refs.extend(extract_node_names(item)) + return refs -def generate_fusion(fusion: FusionPattern, nodes: list[Node]) -> FusionResult: - external_input_nodes = [nodes[node_index].parents[parent_index] for node_index, parent_index in fusion.input_refs] - - fused_name = f"{fusion.target}_{nodes[0].layer.name}".replace(".", "_") - - fused_layer = CVModels.LayerRecord( - index=nodes[0].layer.index, - name=fused_name, - node_type="call_function", - target=fusion.target, - args=[{"node": node.layer.name} for node in external_input_nodes], - kwargs={}, - users=nodes[0].layer.users, - tensor_output_meta=nodes[0].layer.tensor_output_meta, - module_stack=nodes[0].layer.module_stack, - ) - - fused_node = Node( - layer=fused_layer, - normalized_attrs= fusion.output_attrs, - parents=external_input_nodes, - children=nodes[0].children - ) - - - return FusionResult( - fusion=fusion, - matched_nodes=nodes, - external_input_nodes=external_input_nodes, - node=fused_node - ) - + return [] def extract_attrs(layer_: CVModels.LayerRecord) -> dict[str, Any]: @@ -133,7 +95,7 @@ def generate_graph(map: CVModels.LayerMap): temp_source = [] for layer_ in map.nodes: - temp = Node(layer=layer_, normalized_attrs=extract_attrs(layer_), parents=[], children=[]) + temp = Node(layer=layer_, normalized_attrs=extract_attrs(layer_), underlying_op= layer_.target, parents=[], children=[]) temp_map[layer_.name] = temp temp_source.append(temp) if layer_.node_type == "output" else None @@ -141,10 +103,18 @@ def generate_graph(map: CVModels.LayerMap): for name in layer_.users: temp_map[layer_.name].children.append(temp_map[name]) - for item in layer_.args: - if isinstance(item, dict) and "node" in item: - temp_map[layer_.name].parents.append(temp_map[item["node"]]) + for name in extract_node_names(layer_.args): + temp_map[layer_.name].parents.append(temp_map[name]) + + return Graph(name_map=temp_map, source_nodes=temp_source, consumed_ids=set()) + + + + + + + + - return Graph(name_map=temp_map, source_nodes=temp_source, consumed_ids={}) diff --git a/python/cactus/transpiler/ir/nodes.py b/python/cactus/transpiler/ir/nodes.py new file mode 100644 index 000000000..538717e8c --- /dev/null +++ b/python/cactus/transpiler/ir/nodes.py @@ -0,0 +1,2 @@ +import models + diff --git a/python/cactus/transpiler/ir/simplify_ir.py b/python/cactus/transpiler/ir/simplify_ir.py index 287693372..e2eaba665 100644 --- a/python/cactus/transpiler/ir/simplify_ir.py +++ b/python/cactus/transpiler/ir/simplify_ir.py @@ -4,7 +4,7 @@ on the DAG to generate an in-order execution list. """ -import models, fusions_utils +import models, fusions_utils, fusions import converter.models as CVModels @@ -90,142 +90,145 @@ def rev_top_sort(graph:models.Graph) -> models.SimplifiedGraph: return simplified_graph -def resolve_name(name: str, replacement_map: dict[str, str]) -> str: - seen = set() +# - while name in replacement_map and name not in seen: - seen.add(name) - name = replacement_map[name] - return name +def simplify(messy_ir:CVModels.LayerMap) -> CVModels.LayerMap: + graph: models.Graph = models.Graph.from_layer_map(messy_ir) + simplified_graph = rev_top_sort(graph) + return layer_map_from_simplified_graph(simplified_graph, messy_ir) -def rewrite_node_refs(value, replacement_map: dict[str, str]): - if isinstance(value, dict): - rewritten = {key: rewrite_node_refs(item, replacement_map) for key, item in value.items()} - if "node" in rewritten: - rewritten["node"] = resolve_name(rewritten["node"], replacement_map) - return rewritten - if isinstance(value, list): - return [rewrite_node_refs(item, replacement_map) for item in value] - return value +# def resolve_name(name: str, replacement_map: dict[str, str]) -> str: +# seen = set() +# while name in replacement_map and name not in seen: +# seen.add(name) +# name = replacement_map[name] -def extract_node_refs(value) -> list[str]: - if isinstance(value, dict): - refs = [] - if "node" in value: - refs.append(value["node"]) +# return name - for item in value.values(): - refs.extend(extract_node_refs(item)) - return refs +# def rewrite_node_refs(value, replacement_map: dict[str, str]): +# if isinstance(value, dict): +# rewritten = {key: rewrite_node_refs(item, replacement_map) for key, item in value.items()} +# if "node" in rewritten: +# rewritten["node"] = resolve_name(rewritten["node"], replacement_map) +# return rewritten - if isinstance(value, list): - refs = [] - for item in value: - refs.extend(extract_node_refs(item)) - return refs +# if isinstance(value, list): +# return [rewrite_node_refs(item, replacement_map) for item in value] - return [] +# return value -def topo_sort(graph: models.SimplifiedGraph, args_map: dict[str, object], kwargs_map: dict[str, object]) -> list[str]: - deps = {} - users = {node_id: set() for node_id in graph.name_map} +# def extract_node_refs(value) -> list[str]: +# if isinstance(value, dict): +# refs = [] +# if "node" in value: +# refs.append(value["node"]) - for node_id in graph.name_map: - node_refs = extract_node_refs(args_map[node_id]) + extract_node_refs(kwargs_map[node_id]) - deps[node_id] = { - ref - for ref in node_refs - if ref in graph.name_map and ref != node_id - } +# for item in value.values(): +# refs.extend(extract_node_refs(item)) - for dep in deps[node_id]: - users[dep].add(node_id) +# return refs - ready = sorted( - [node_id for node_id, node_deps in deps.items() if not node_deps], - key=lambda node_id: graph.name_map[node_id].layer.index, - ) - ordered = [] +# if isinstance(value, list): +# refs = [] +# for item in value: +# refs.extend(extract_node_refs(item)) +# return refs - while ready: - node_id = ready.pop(0) - ordered.append(node_id) +# return [] - for user_id in sorted(users[node_id], key=lambda item: graph.name_map[item].layer.index): - deps[user_id].remove(node_id) - if not deps[user_id]: - ready.append(user_id) - ready.sort(key=lambda item: graph.name_map[item].layer.index) +# def topo_sort(graph: models.SimplifiedGraph, args_map: dict[str, object], kwargs_map: dict[str, object]) -> list[str]: +# deps = {} +# users = {node_id: set() for node_id in graph.name_map} - if len(ordered) != len(graph.name_map): - raise ValueError("Simplified graph contains a cycle or unresolved dependency") +# for node_id in graph.name_map: +# node_refs = extract_node_refs(args_map[node_id]) + extract_node_refs(kwargs_map[node_id]) +# deps[node_id] = { +# ref +# for ref in node_refs +# if ref in graph.name_map and ref != node_id +# } - return ordered +# for dep in deps[node_id]: +# users[dep].add(node_id) +# ready = sorted( +# [node_id for node_id, node_deps in deps.items() if not node_deps], +# key=lambda node_id: graph.name_map[node_id].layer.index, +# ) +# ordered = [] -def layer_map_from_simplified_graph(graph: models.SimplifiedGraph, original: CVModels.LayerMap) -> CVModels.LayerMap: - fused_node_ids = {result.node.layer.name for result in graph.fused_results} - args_map = {} - kwargs_map = {} +# while ready: +# node_id = ready.pop(0) +# ordered.append(node_id) - for node_id, node in graph.name_map.items(): - args_map[node_id] = rewrite_node_refs(node.layer.args, graph.replacement_map) - - if node_id in fused_node_ids: - kwargs_map[node_id] = rewrite_node_refs(node.normalized_attrs, graph.replacement_map) - else: - kwargs_map[node_id] = rewrite_node_refs(node.layer.kwargs, graph.replacement_map) - - ordered_node_ids = topo_sort(graph, args_map, kwargs_map) - users_map = {node_id: [] for node_id in graph.name_map} - - for node_id in ordered_node_ids: - node_refs = extract_node_refs(args_map[node_id]) + extract_node_refs(kwargs_map[node_id]) - for ref in node_refs: - if ref in users_map and node_id not in users_map[ref]: - users_map[ref].append(node_id) - - nodes = [] - for index, node_id in enumerate(ordered_node_ids): - node = graph.name_map[node_id] - nodes.append( - CVModels.LayerRecord( - index=index, - name=node.layer.name, - node_type=node.layer.node_type, - target=node.layer.target, - args=args_map[node_id], - kwargs=kwargs_map[node_id], - users=users_map[node_id], - tensor_output_meta=node.layer.tensor_output_meta, - module_stack=node.layer.module_stack, - ) - ) - - - return CVModels.LayerMap( - model_name=original.model_name, - task=original.task, - graph_signature=original.graph_signature, - range_constants=original.range_constants, - nodes=nodes, - ) +# for user_id in sorted(users[node_id], key=lambda item: graph.name_map[item].layer.index): +# deps[user_id].remove(node_id) +# if not deps[user_id]: +# ready.append(user_id) +# ready.sort(key=lambda item: graph.name_map[item].layer.index) -def simplify(messy_ir:CVModels.LayerMap) -> CVModels.LayerMap: - graph: models.Graph = models.Graph.from_layer_map(messy_ir) - simplified_graph = rev_top_sort(graph) - return layer_map_from_simplified_graph(simplified_graph, messy_ir) +# if len(ordered) != len(graph.name_map): +# raise ValueError("Simplified graph contains a cycle or unresolved dependency") + +# return ordered + + +# def layer_map_from_simplified_graph(graph: models.SimplifiedGraph, original: CVModels.LayerMap) -> CVModels.LayerMap: +# fused_node_ids = {result.node.layer.name for result in graph.fused_results} +# args_map = {} +# kwargs_map = {} + +# for node_id, node in graph.name_map.items(): +# args_map[node_id] = rewrite_node_refs(node.layer.args, graph.replacement_map) + +# if node_id in fused_node_ids: +# kwargs_map[node_id] = rewrite_node_refs(node.normalized_attrs, graph.replacement_map) +# else: +# kwargs_map[node_id] = rewrite_node_refs(node.layer.kwargs, graph.replacement_map) + +# ordered_node_ids = topo_sort(graph, args_map, kwargs_map) +# users_map = {node_id: [] for node_id in graph.name_map} +# for node_id in ordered_node_ids: +# node_refs = extract_node_refs(args_map[node_id]) + extract_node_refs(kwargs_map[node_id]) +# for ref in node_refs: +# if ref in users_map and node_id not in users_map[ref]: +# users_map[ref].append(node_id) +# nodes = [] +# for index, node_id in enumerate(ordered_node_ids): +# node = graph.name_map[node_id] +# nodes.append( +# CVModels.LayerRecord( +# index=index, +# name=node.layer.name, +# node_type=node.layer.node_type, +# target=node.layer.target, +# args=args_map[node_id], +# kwargs=kwargs_map[node_id], +# users=users_map[node_id], +# tensor_output_meta=node.layer.tensor_output_meta, +# module_stack=node.layer.module_stack, +# ) +# ) + + +# return CVModels.LayerMap( +# model_name=original.model_name, +# task=original.task, +# graph_signature=original.graph_signature, +# range_constants=original.range_constants, +# nodes=nodes, +# ) From e6baa928b31bf7ad8c33f37a10b77112bbc260bd Mon Sep 17 00:00:00 2001 From: Parkirat Date: Thu, 9 Jul 2026 11:26:30 -0400 Subject: [PATCH 15/84] Added matchers which will check for synthetic graph and main graph equality (utilized to detect valid fusions) --- python/cactus/transpiler/ir/fusions.py | 16 ++ python/cactus/transpiler/ir/fusions_utils.py | 207 +++++++++---------- python/cactus/transpiler/ir/models.py | 42 ++-- python/cactus/transpiler/ir/nodes.py | 1 + 4 files changed, 143 insertions(+), 123 deletions(-) diff --git a/python/cactus/transpiler/ir/fusions.py b/python/cactus/transpiler/ir/fusions.py index 09df5428a..3559660f1 100644 --- a/python/cactus/transpiler/ir/fusions.py +++ b/python/cactus/transpiler/ir/fusions.py @@ -11,6 +11,19 @@ models.Edge("variance_with_eps", 0, "variance"), models.Edge("variance", 0, "squared_x"), ), + inputs=( + models.InputSpec("normalized_x", 0), + models.InputSpec("weighted_output", 1), + ), + shared_inputs=( + ( + models.InputSpec("normalized_x", 0), + models.InputSpec("squared_x", 0), + ), + ), + output_attrs={ + "eps": models.AttrRef("variance_with_eps", "scalar"), + }, ) @@ -24,3 +37,6 @@ RMS_NORM, ), } + +#TODO: Update with str paths and their respective fusion object +OPS_MAP: dict[str, list[models.FusionPattern]] = {} diff --git a/python/cactus/transpiler/ir/fusions_utils.py b/python/cactus/transpiler/ir/fusions_utils.py index 0449d807d..df02ecbf0 100644 --- a/python/cactus/transpiler/ir/fusions_utils.py +++ b/python/cactus/transpiler/ir/fusions_utils.py @@ -1,152 +1,147 @@ -import models, constants -#TODO: Update with str paths and their respective fusion object -OPS_MAP: dict[str, list[models.FusionPattern]] = {} +import models, nodes, fusions + + """###################################### FUSION UTILS!!!!!!! ######################################""" +def node_match(node:models.Node, synth_node:models.NodeSpec) -> bool: + if node.underlying_op != synth_node.underlying_op: + return False + + for key, expected_value in synth_node.required_attrs.items(): + if key not in node.normalized_attrs: + return False + if node.normalized_attrs[key] != expected_value: + return False -# def len_match(fusion: models.FusionPattern, nodes: list[models.Node]) -> bool: -# return len(fusion.ops) == len(nodes) + return True + +def get_bound_parent(bindings: dict[str, models.Node], node_name: str, parent_index: int) -> models.Node | None: + if node_name not in bindings: + return None -# def match_ops(fusion: models.FusionPattern, nodes: list[models.Node]) -> bool: -# for i in range(len(fusion.ops)): -# if(fusion.ops[i] != nodes[i].layer.target): -# return False -# return True + node = bindings[node_name] + if parent_index < 0 or parent_index >= len(node.parents): + return None + return node.parents[parent_index] -# def match_path(fusion: models.FusionPattern, nodes: list[models.Node]) -> bool: -# if len(fusion.path) != len(nodes) - 1: -# return False +def get_input_parent(bindings: dict[str, models.Node], input_spec: models.InputSpec) -> models.Node | None: + return get_bound_parent(bindings, input_spec.node, input_spec.parent_index) -# for i, parent_index in enumerate(fusion.path): -# if parent_index < 0 or parent_index >= len(nodes[i].parents): -# return False +#Match edges +def edges_match(node: models.Node, synth_graph: models.FusionGraph) -> dict[str, models.Node] | None: + bindings:dict[str,models.Node] = {} -# if nodes[i].parents[parent_index].layer.name != nodes[i + 1].layer.name: -# return False + if not node_match(node, nodes.NODE_MAP[synth_graph.root]): + return None -# return True + bindings[synth_graph.root] = node + + for edge in synth_graph.edges: + parent = get_bound_parent(bindings, edge.from_node, edge.parent_index) + if parent is None: + return None -# def match_attrs(fusion: models.FusionPattern, nodes: list[models.Node]) -> bool: -# for node_index, required_attrs in fusion.required_attrs.items(): -# if node_index < 0 or node_index >= len(nodes): -# return False + if edge.to_node not in nodes.NODE_MAP: + return None -# node_attrs = nodes[node_index].normalized_attrs -# for attr_name, expected_value in required_attrs.items(): -# if attr_name not in node_attrs: -# return False + if not node_match(parent, nodes.NODE_MAP[edge.to_node]): + return None -# if node_attrs[attr_name] != expected_value: -# return False + if edge.to_node in bindings and bindings[edge.to_node] is not parent: + return None -# return True + bindings[edge.to_node] = parent + return bindings -# def match_input_refs(fusion: models.FusionPattern, nodes: list[models.Node]) -> bool: -# for node_index, parent_index in fusion.input_refs: -# if node_index < 0 or node_index >= len(nodes): -# return False -# if parent_index < 0 or parent_index >= len(nodes[node_index].parents): -# return False +def match_inputs_exist(bindings: dict[str, models.Node], synth_graph: models.FusionGraph) -> bool: + for input_spec in synth_graph.inputs: + if get_input_parent(bindings, input_spec) is None: + return False -# return True + return True +def match_external_inputs_dec(bindings: dict[str, models.Node], synth_graph: models.FusionGraph) -> bool: + internal_node_ids = {id(node) for node in bindings.values()} + declared_input_ids = set() -# def match_shared_input_refs(fusion: models.FusionPattern, nodes: list[models.Node]) -> bool: -# for left_ref, right_ref in fusion.shared_input_refs: -# left_node_index, left_parent_index = left_ref -# right_node_index, right_parent_index = right_ref + for input_spec in synth_graph.inputs: + parent = get_input_parent(bindings, input_spec) + if parent is None: + return False -# if left_node_index < 0 or left_node_index >= len(nodes): -# return False + declared_input_ids.add(id(parent)) -# if right_node_index < 0 or right_node_index >= len(nodes): -# return False + for node in bindings.values(): + for parent in node.parents: + if id(parent) not in internal_node_ids and id(parent) not in declared_input_ids: + return False -# if left_parent_index < 0 or left_parent_index >= len(nodes[left_node_index].parents): -# return False + return True -# if right_parent_index < 0 or right_parent_index >= len(nodes[right_node_index].parents): -# return False +def match_shared_inputs(bindings: dict[str, models.Node], synth_graph: models.FusionGraph) -> bool: + for left_input, right_input in synth_graph.shared_inputs: + left_parent = get_input_parent(bindings, left_input) + right_parent = get_input_parent(bindings, right_input) + if left_parent is None or right_parent is None: + return False -# left_parent = nodes[left_node_index].parents[left_parent_index] -# right_parent = nodes[right_node_index].parents[right_parent_index] + if left_parent is not right_parent: + return False -# if left_parent.layer.name != right_parent.layer.name: -# return False + return True -# return True +def match_no_external_internal_children(bindings: dict[str, models.Node], synth_graph: models.FusionGraph) -> bool: + if synth_graph.root not in bindings: + return False + root_node = bindings[synth_graph.root] + internal_node_ids = {id(node) for node in bindings.values()} -# def match_no_external_internal_children(fusion: models.FusionPattern, nodes: list[models.Node]) -> bool: -# fused_node_names = {node.layer.name for node in nodes} + for node in bindings.values(): + if node is root_node: + continue -# for node in nodes[1:]: -# for child in node.children: -# if child.layer.name not in fused_node_names: -# return False + for child in node.children: + if id(child) not in internal_node_ids: + return False -# return True + return True -# def match_external_inputs(fusion: models.FusionPattern, nodes: list[models.Node]): -# proposed_fusion_node_names = {node.layer.name for node in nodes} +def match_output_attrs(bindings: dict[str, models.Node], synth_graph: models.FusionGraph) -> bool: + for attr_value in synth_graph.output_attrs.values(): + if not isinstance(attr_value, models.AttrRef): + continue -# declared_inputs = {nodes[node_index].parents[parent_index].layer.name for node_index, parent_index in fusion.input_refs} + if attr_value.node not in bindings: + return False -# for node in nodes: -# for parent in node.parents: -# if parent.layer.name not in proposed_fusion_node_names and parent.layer.name not in declared_inputs: -# return False - -# return True + node = bindings[attr_value.node] + if attr_value.attr not in node.normalized_attrs: + return False + return True #Update this with any additional matching functions if any are added MATCHERS = ( - len_match, - match_attrs, - match_input_refs, - match_ops, - match_path, - match_shared_input_refs, + match_inputs_exist, + match_external_inputs_dec, + match_shared_inputs, match_no_external_internal_children, - match_external_inputs, + match_output_attrs, ) -def fusion_match(ops:str, nodes: list[models.Node]) -> models.FusionPattern | None: - if ops in OPS_MAP: - for fusion in OPS_MAP[ops]: - if all(matcher(fusion, nodes) for matcher in MATCHERS): - return fusion - - return None - - - -"""###################################### FUSION GRAPH TRAVERSAL UTILS!!!!!!! ######################################""" -def dfs_path_gen(path:list[models.Node], graph:models.Graph, max_depth: int = constants.DFS_DEPTH): - if len(path) >= max_depth or not path[-1].parents: - yield path - return - - added_to_path = False - - for parent in path[-1].parents: - if parent.layer.name in graph.consumed_ids: - continue - - if any(existing.layer.name == parent.layer.name for existing in path): - continue - - added_to_path = True - yield from dfs_path_gen([*path, parent], graph, max_depth) +def fusion_match(node: models.Node) -> models.FusionGraph | None: - if not added_to_path: - yield path + for fusion in fusions.ROOT_TARGET_MAP.get(node.underlying_op, ()): + bindings = edges_match(node, fusion) + if bindings != None and all(matcher(bindings, fusion) for matcher in MATCHERS): + #Will change this to return FusionResult object which will have a list of all nodes being fused + return fusion diff --git a/python/cactus/transpiler/ir/models.py b/python/cactus/transpiler/ir/models.py index 7b2710f94..2dca41c07 100644 --- a/python/cactus/transpiler/ir/models.py +++ b/python/cactus/transpiler/ir/models.py @@ -1,4 +1,4 @@ -from dataclasses import dataclass +from dataclasses import dataclass, field from typing import Any, Optional import converter.models as CVModels import constants @@ -22,24 +22,43 @@ class Graph: def from_layer_map(cls, map:CVModels.LayerMap) -> "Graph": return generate_graph(map) -@dataclass(slots = True) +@dataclass(slots=True) class NodeSpec: underlying_op:str - required_attrs: dict[str, Any] + required_attrs: dict[str, Any] = field(default_factory=dict) -@dataclass(slots = True) +@dataclass(slots=True) class Edge: from_node:str parent_index:int to_node:str + +@dataclass(slots=True) +class InputSpec: + node:str + parent_index:int + + +@dataclass(slots=True) +class AttrRef: + node:str + attr:str + + @dataclass(slots=True) class FusionGraph: target:str root:str - edges:tuple[Edge] - + edges:tuple[Edge, ...] + inputs:tuple[InputSpec, ...] = () + shared_inputs:tuple[tuple[InputSpec, InputSpec], ...] = () + output_attrs:dict[str, Any] = field(default_factory=dict) + +class FusionResult: + node:Node + edges_consumed:list[str] #Can be changed to """###################################### MODEL UTILS!!!!!!! ######################################""" @@ -107,14 +126,3 @@ def generate_graph(map: CVModels.LayerMap): temp_map[layer_.name].parents.append(temp_map[name]) return Graph(name_map=temp_map, source_nodes=temp_source, consumed_ids=set()) - - - - - - - - - - - diff --git a/python/cactus/transpiler/ir/nodes.py b/python/cactus/transpiler/ir/nodes.py index 538717e8c..39561b473 100644 --- a/python/cactus/transpiler/ir/nodes.py +++ b/python/cactus/transpiler/ir/nodes.py @@ -1,2 +1,3 @@ import models +NODE_MAP = {} \ No newline at end of file From 6961ed4bf331c4a90a888db52dc78e94fcea05b5 Mon Sep 17 00:00:00 2001 From: Parkirat Date: Thu, 9 Jul 2026 18:23:38 -0400 Subject: [PATCH 16/84] Reset architecture to support harder fusions --- python/cactus/transpiler/converter/convert.py | 11 +- .../converter/model_preprocessing.py | 2 +- python/cactus/transpiler/ir/constants.py | 32 +- .../transpiler/ir/fusion_graph_catalog.py | 530 +++++++++++ python/cactus/transpiler/ir/fusions.py | 42 - python/cactus/transpiler/ir/fusions_utils.py | 491 ++++++++-- python/cactus/transpiler/ir/models.py | 896 ++++++++++++++++-- python/cactus/transpiler/ir/nodes.py | 3 - python/cactus/transpiler/ir/simplify_ir.py | 224 +---- python/cactus/transpiler/main.py | 3 +- 10 files changed, 1762 insertions(+), 472 deletions(-) create mode 100644 python/cactus/transpiler/ir/fusion_graph_catalog.py delete mode 100644 python/cactus/transpiler/ir/fusions.py delete mode 100644 python/cactus/transpiler/ir/nodes.py diff --git a/python/cactus/transpiler/converter/convert.py b/python/cactus/transpiler/converter/convert.py index 02bdb2fdb..08e4bd77c 100644 --- a/python/cactus/transpiler/converter/convert.py +++ b/python/cactus/transpiler/converter/convert.py @@ -1,8 +1,7 @@ import torch -import json -from models import LayerRecord, LayerMap -from model_preprocessing import load_model, load_input -import constants +from .models import LayerRecord, LayerMap +from .model_preprocessing import load_model, load_input +from . import constants @@ -23,12 +22,12 @@ def model_ops_to_json(id:str, output_path:str) -> LayerMap: json_ = LayerMap.from_data(x=exported_model, name=id, model_task = constants.MODEL_TASK, nodes_list=nodes).model_dump_json(indent=4) - with open(output_path+"output.json", "w", encoding="utf-8") as file: + with open(output_path+"output1.json", "w", encoding="utf-8") as file: file.write(json_) if __name__ == "__main__": - model_ops_to_json("google/gemma-4-E2B", "/Users/sandhup/Documents/personal/cactus/python/cactus/transpiler/converter/jsons/") + model_ops_to_json("LiquidAI/LFM2.5-8B-A1B", "/Users/sandhup/Documents/personal/cactus/python/cactus/transpiler/converter/jsons/") print("Done") diff --git a/python/cactus/transpiler/converter/model_preprocessing.py b/python/cactus/transpiler/converter/model_preprocessing.py index 5d8f4fc71..4d66129f9 100644 --- a/python/cactus/transpiler/converter/model_preprocessing.py +++ b/python/cactus/transpiler/converter/model_preprocessing.py @@ -1,4 +1,4 @@ -import constants +from . import constants from transformers import AutoTokenizer, AutoModelForCausalLM def load_model(id:str) -> object: diff --git a/python/cactus/transpiler/ir/constants.py b/python/cactus/transpiler/ir/constants.py index 1f8c9c692..4ecf94e6f 100644 --- a/python/cactus/transpiler/ir/constants.py +++ b/python/cactus/transpiler/ir/constants.py @@ -1,16 +1,28 @@ LAYER_ATTRS_MAP = { "aten._assert_tensor_metadata.default": ["size", "stride", "dtype"], "aten._softmax.default": ["dim", "half_to_float"], - "aten._to_copy.default": [], + "aten._to_copy.default": [None, "dtype", "layout", "device", "pin_memory", "non_blocking", "memory_format"], + "aten.add.Scalar": ["scalar"], "aten.add.Tensor": ["scalar"], + "aten.addmm.default": [None, None, None, "beta", "alpha"], "aten.arange.start_step": ["start", "end", "step"], + "aten.avg_pool1d.default": [None, "kernel_size", "stride", "padding", "ceil_mode", "count_include_pad"], + "aten.avg_pool2d.default": [None, "kernel_size", "stride", "padding", "ceil_mode", "count_include_pad", "divisor_override"], + "aten.batch_norm.default": [None, None, None, None, None, "training", "momentum", "eps", "cudnn_enabled"], "aten.bitwise_and.Tensor": [], "aten.bitwise_or.Tensor": [], "aten.bmm.default": [], "aten.cat.default": ["dim"], "aten.clone.default": [], + "aten.clamp.default": [None, "min", "max"], + "aten.clamp_max.default": [None, "max"], + "aten.clamp_min.default": [None, "min"], + "aten.convolution.default": [None, None, None, "stride", "padding", "dilation", "transposed", "output_padding", "groups"], + "aten.conv1d.default": [None, None, None, "stride", "padding", "dilation", "groups"], + "aten.conv2d.default": [None, None, None, "stride", "padding", "dilation", "groups"], "aten.cos.default": [], "aten.cumsum.default": ["dim"], + "aten.div.Scalar": ["scalar"], "aten.div.Tensor": ["scalar"], "aten.embedding.default": ["padding_idx"], "aten.eq.Scalar": ["scalar"], @@ -18,24 +30,42 @@ "aten.expand.default": ["shape"], "aten.full.default": ["shape", "fill_value"], "aten.gelu.default": [], + "aten.glu.default": ["dim"], + "aten.group_norm.default": [None, "num_groups", None, None, "eps", "cudnn_enabled"], "aten.gt.Tensor": [], "aten.index.Tensor": [], "aten.index_put.default": [], + "aten.layer_norm.default": [None, "normalized_shape", None, None, "eps", "cudnn_enabled"], "aten.le.Tensor": [], + "aten.linear.default": [], + "aten.max.dim": ["dim", "keepdim"], + "aten.max_pool1d.default": [None, "kernel_size", "stride", "padding", "dilation", "ceil_mode"], + "aten.max_pool2d.default": [None, "kernel_size", "stride", "padding", "dilation", "ceil_mode"], "aten.mean.dim": ["dim", "keepdim"], "aten.mm.default": [], + "aten.mul.Scalar": ["scalar"], "aten.mul.Tensor": ["scalar"], + "aten.native_batch_norm.default": [None, None, None, None, None, "training", "momentum", "eps"], + "aten.native_layer_norm.default": [None, "normalized_shape", None, None, "eps"], "aten.ne.Scalar": ["scalar"], + "aten.ne.Tensor": [], "aten.neg.default": [], "aten.permute.default": ["dims"], "aten.pow.Tensor_Scalar": ["exponent"], + "aten.rms_norm.default": [None, "normalized_shape", None, "eps"], "aten.scalar_tensor.default": ["value"], + "aten.scaled_dot_product_attention.default": [None, None, None, "attn_mask", "dropout_p", "is_causal", "scale", "enable_gqa"], "aten.select.int": ["dim", "index"], "aten.sin.default": [], "aten.slice.Tensor": ["dim", "start", "end", "step"], + "aten.sub.Scalar": ["scalar"], "aten.sub.Tensor": ["scalar"], + "aten.sum.dim_IntList": ["dim", "keepdim"], "aten.tanh.default": [], + "aten.topk.default": [None, "k", "dim", "largest", "sorted"], + "aten.transpose.int": ["dim0", "dim1"], "aten.unsqueeze.default": ["dim"], + "aten.var.dim": ["dim", "unbiased", "keepdim"], "aten.view.default": ["shape"], "aten.where.self": [], } diff --git a/python/cactus/transpiler/ir/fusion_graph_catalog.py b/python/cactus/transpiler/ir/fusion_graph_catalog.py new file mode 100644 index 000000000..4bc13a0f5 --- /dev/null +++ b/python/cactus/transpiler/ir/fusion_graph_catalog.py @@ -0,0 +1,530 @@ +from __future__ import annotations + +from dataclasses import dataclass, field +from typing import Any + + +@dataclass(frozen=True, slots=True) +class FusionNode: + name: str + op: str | tuple[str, ...] + attrs: dict[str, Any] = field(default_factory=dict) + optional: bool = False + transparent: bool = False + repeated: bool = False + notes: str = "" + + +@dataclass(frozen=True, slots=True) +class FusionEdge: + source: str + dest: str + dest_input_index: int + source_output_index: int | None = None + optional: bool = False + repeated: bool = False + + +@dataclass(frozen=True, slots=True) +class FusionInput: + role: str + node: str + parent_index: int | None = None + variadic: bool = False + optional: bool = False + repeated_group: str | None = None + + +@dataclass(frozen=True, slots=True) +class AttrCapture: + name: str + source_node: str | None = None + source_attr: str | None = None + default: Any = None + required: bool = True + notes: str = "" + + +@dataclass(frozen=True, slots=True) +class FusionGraph: + name: str + target: str + cactus_inputs: tuple[str, ...] + cactus_attrs: tuple[str, ...] + nodes: dict[str, FusionNode] = field(default_factory=dict) + edges: tuple[FusionEdge, ...] = () + root: str | None = None + outputs: tuple[str, ...] = () + inputs: tuple[FusionInput, ...] = () + output_attrs: tuple[AttrCapture, ...] = () + constraints: tuple[str, ...] = () + variants: tuple[str, ...] = () + repeated_groups: tuple[str, ...] = () + notes: str = "" + + +def n( + name: str, + op: str | tuple[str, ...], + *, + attrs: dict[str, Any] | None = None, + optional: bool = False, + transparent: bool = False, + repeated: bool = False, + notes: str = "", +) -> FusionNode: + return FusionNode( + name=name, + op=op, + attrs=attrs or {}, + optional=optional, + transparent=transparent, + repeated=repeated, + notes=notes, + ) + + +def e( + source: str, + dest: str, + dest_input_index: int, + *, + source_output_index: int | None = None, + optional: bool = False, + repeated: bool = False, +) -> FusionEdge: + return FusionEdge( + source=source, + dest=dest, + dest_input_index=dest_input_index, + source_output_index=source_output_index, + optional=optional, + repeated=repeated, + ) + + +def inp( + role: str, + node: str, + parent_index: int | None = None, + *, + variadic: bool = False, + optional: bool = False, + repeated_group: str | None = None, +) -> FusionInput: + return FusionInput( + role=role, + node=node, + parent_index=parent_index, + variadic=variadic, + optional=optional, + repeated_group=repeated_group, + ) + + +def cap( + name: str, + source_node: str | None = None, + source_attr: str | None = None, + *, + default: Any = None, + required: bool = True, + notes: str = "", +) -> AttrCapture: + return AttrCapture( + name=name, + source_node=source_node, + source_attr=source_attr, + default=default, + required=required, + notes=notes, + ) + + +def direct( + target: str, + cactus_inputs: tuple[str, ...], + cactus_attrs: tuple[str, ...], + raw_ops: str | tuple[str, ...], + *, + attr_map: dict[str, str] | None = None, + optional_attr_defaults: dict[str, Any] | None = None, + notes: str = "", +) -> FusionGraph: + root = f"{target}_root" + output_attrs = [] + + for cactus_attr in cactus_attrs: + source_attr = (attr_map or {}).get(cactus_attr, cactus_attr) + output_attrs.append( + cap( + cactus_attr, + root, + source_attr, + default=(optional_attr_defaults or {}).get(cactus_attr), + required=cactus_attr not in (optional_attr_defaults or {}), + ) + ) + + return FusionGraph( + name=f"{target}_direct", + target=target, + cactus_inputs=cactus_inputs, + cactus_attrs=cactus_attrs, + nodes={root: n(root, raw_ops)}, + root=root, + outputs=(root,), + inputs=tuple(inp(role, root, i) for i, role in enumerate(cactus_inputs)), + output_attrs=tuple(output_attrs), + notes=notes, + ) + + +def source_op(target: str, cactus_attrs: tuple[str, ...], notes: str) -> FusionGraph: + root = f"{target}_source" + return FusionGraph( + name=f"{target}_source", + target=target, + cactus_inputs=(), + cactus_attrs=cactus_attrs, + nodes={root: n(root, "cactus.source_op", notes=notes)}, + root=root, + outputs=(root,), + output_attrs=tuple(cap(attr, default=None, notes="Caller/runtime supplied attr") for attr in cactus_attrs), + notes=notes, + ) + + +FUSION_GRAPHS: dict[str, FusionGraph] = { + "add": direct("add", ("a", "b"), (), "aten.add.Tensor"), + "add_clipped": direct( + "add_clipped", + ("a", "b"), + (), + "aten.add.Tensor", + notes="Needs post-add clipping semantics from a trace before this can be safely fused.", + ), + "subtract": direct("subtract", ("a", "b"), (), "aten.sub.Tensor"), + "multiply": direct("multiply", ("a", "b"), (), "aten.mul.Tensor"), + "divide": direct("divide", ("a", "b"), (), "aten.div.Tensor"), + "not_equal": direct("not_equal", ("a", "b"), (), "aten.ne.Tensor"), + "abs": direct("abs", ("x",), (), "aten.abs.default"), + "pow": direct("pow", ("x",), ("exponent",), "aten.pow.Tensor_Scalar"), + "precision_cast": direct( + "precision_cast", + ("x",), + ("dtype",), + "aten._to_copy.default", + attr_map={"dtype": "dtype"}, + ), + "quantize_activations": direct( + "quantize_activations", + ("x",), + (), + ("aten.quantize_per_tensor.default", "aten._to_copy.default"), + notes="Needs quantization-specific attrs/scales once a quantized export is available.", + ), + "scalar_add": direct("scalar_add", ("x",), ("value",), "aten.add.Scalar", attr_map={"value": "scalar"}), + "scalar_subtract": direct("scalar_subtract", ("x",), ("value",), "aten.sub.Scalar", attr_map={"value": "scalar"}), + "scalar_multiply": direct("scalar_multiply", ("x",), ("value",), ("aten.mul.Scalar", "aten.neg.default"), attr_map={"value": "scalar"}), + "scalar_divide": direct("scalar_divide", ("x",), ("value",), "aten.div.Scalar", attr_map={"value": "scalar"}), + "scalar_floor_divide": direct("scalar_floor_divide", ("x",), ("value",), "aten.floor_divide.default", attr_map={"value": "scalar"}), + "scalar_not_equal": direct("scalar_not_equal", ("x",), ("value",), "aten.ne.Scalar", attr_map={"value": "scalar"}), + "scalar_exp": direct("scalar_exp", ("x",), (), "aten.exp.default"), + "scalar_sqrt": direct("scalar_sqrt", ("x",), (), "aten.sqrt.default"), + "scalar_cos": direct("scalar_cos", ("x",), (), "aten.cos.default"), + "scalar_sin": direct("scalar_sin", ("x",), (), "aten.sin.default"), + "scalar_log": direct("scalar_log", ("x",), (), "aten.log.default"), + "clamp": direct("clamp", ("x",), ("lo", "hi"), "aten.clamp.default", attr_map={"lo": "min", "hi": "max"}), + "masked_select_prefix": direct( + "masked_select_prefix", + ("x", "mask"), + (), + ("aten.masked_select.default", "aten.index.Tensor"), + notes="Needs prefix semantics check, not just generic masked_select.", + ), + "masked_scatter": direct("masked_scatter", ("x", "mask", "source"), (), "aten.masked_scatter.default"), + "view": direct("view", ("x",), ("shape",), "aten.view.default"), + "reshape": direct("reshape", ("x",), ("shape",), "aten.reshape.default"), + "expand": direct("expand", ("x",), ("shape",), "aten.expand.default"), + "flatten": direct( + "flatten", + ("x",), + ("start_dim", "end_dim"), + "aten.flatten.using_ints", + optional_attr_defaults={"start_dim": 0, "end_dim": -1}, + ), + "slice": direct("slice", ("x",), ("axis", "start", "end", "step"), "aten.slice.Tensor", attr_map={"axis": "dim"}), + "index": direct("index", ("x",), ("index_value", "axis"), "aten.select.int", attr_map={"axis": "dim", "index_value": "index"}), + "transpose": direct( + "transpose", + ("x",), + ("backend",), + "aten.transpose.int", + optional_attr_defaults={"backend": 0}, + notes="Needs dim0/dim1 capture if Cactus transpose grows beyond backend-only attrs.", + ), + "permute": direct( + "permute", + ("x",), + ("permutation", "backend"), + "aten.permute.default", + attr_map={"permutation": "dims"}, + optional_attr_defaults={"backend": 0}, + ), + "matmul": direct( + "matmul", + ("a", "b"), + ("pretransposed_rhs", "backend", "output_dtype"), + ("aten.mm.default", "aten.matmul.default", "aten.bmm.default"), + optional_attr_defaults={"pretransposed_rhs": False, "backend": 0, "output_dtype": None}, + ), + "gather": direct("gather", ("tensor", "indices"), (), "aten.gather.default"), + "embedding_from_tensor": direct("embedding_from_tensor", ("embedding_tensor", "indices"), (), "aten.embedding.default"), + "embedding_from_file": source_op("embedding_from_file", ("filename",), "Runtime/source op. No ATen subgraph unless file loading is represented in IR."), + "mmap_embeddings": source_op("mmap_embeddings", ("filename",), "Runtime/source op for memory-mapped embedding weights."), + "mmap_weights": source_op("mmap_weights", ("filename",), "Runtime/source op for memory-mapped weights."), + "bilinear_interpolation": direct( + "bilinear_interpolation", + ("pos_embeds",), + ("dst_height", "dst_width"), + ("aten.upsample_bilinear2d.vec", "aten._upsample_bilinear2d_aa.vec"), + ), + "concat": direct("concat", ("a", "b"), ("axis",), "aten.cat.default", attr_map={"axis": "dim"}), + "cat": FusionGraph( + name="cat_variadic", + target="cat", + cactus_inputs=("tensors",), + cactus_attrs=("axis",), + nodes={"cat": n("cat", "aten.cat.default")}, + root="cat", + outputs=("cat",), + inputs=(inp("tensors", "cat", variadic=True),), + output_attrs=(cap("axis", "cat", "dim"),), + ), + "groupnorm": direct("groupnorm", ("x", "weight", "bias"), ("num_groups", "eps"), "aten.group_norm.default"), + "layernorm": direct("layernorm", ("x", "weight", "bias"), ("eps",), ("aten.layer_norm.default", "aten.native_layer_norm.default")), + "batchnorm": direct("batchnorm", ("x", "weight", "bias", "running_mean", "running_var"), ("axis", "eps"), ("aten.batch_norm.default", "aten.native_batch_norm.default"), optional_attr_defaults={"axis": 1}), + "rms_norm": FusionGraph( + name="rms_norm_decomposed_or_direct", + target="rms_norm", + cactus_inputs=("x", "weight"), + cactus_attrs=("eps",), + nodes={ + "squared_x": n("squared_x", "aten.pow.Tensor_Scalar", attrs={"exponent": 2}), + "variance": n("variance", "aten.mean.dim"), + "variance_with_eps": n("variance_with_eps", "aten.add.Tensor"), + "inverse_rms": n("inverse_rms", "aten.pow.Tensor_Scalar", attrs={"exponent": -0.5}), + "normalized_x": n("normalized_x", "aten.mul.Tensor"), + "weighted_output": n("weighted_output", "aten.mul.Tensor"), + }, + edges=( + e("squared_x", "variance", 0), + e("variance", "variance_with_eps", 0), + e("variance_with_eps", "inverse_rms", 0), + e("inverse_rms", "normalized_x", 1), + e("normalized_x", "weighted_output", 0), + ), + root="weighted_output", + outputs=("weighted_output",), + inputs=(inp("x", "normalized_x", 0), inp("weight", "weighted_output", 1)), + output_attrs=(cap("eps", "variance_with_eps", "scalar"),), + constraints=("normalized_x input 0 and squared_x input 0 must be the same node", "mean axis must be hidden dimension"), + variants=("aten.rms_norm.default", "pow_mean_add_pow_mul_mul"), + ), + "topk": direct("topk", ("x",), ("k",), "aten.topk.default"), + "rope": FusionGraph( + name="rope_decomposed", + target="rope", + cactus_inputs=("x",), + cactus_attrs=("theta", "position_offset", "backend"), + nodes={ + "x": n("x", "external"), + "sin": n("sin", ("aten.sin.default", "constant_sin_cache")), + "cos": n("cos", ("aten.cos.default", "constant_cos_cache")), + "rotate_half": n("rotate_half", "rope_rotate_half_subgraph"), + "scaled_x": n("scaled_x", "aten.mul.Tensor"), + "scaled_rot": n("scaled_rot", "aten.mul.Tensor"), + "out": n("out", "aten.add.Tensor"), + }, + edges=(e("x", "scaled_x", 0), e("cos", "scaled_x", 1), e("x", "rotate_half", 0), e("rotate_half", "scaled_rot", 0), e("sin", "scaled_rot", 1), e("scaled_x", "out", 0), e("scaled_rot", "out", 1)), + root="out", + outputs=("out",), + inputs=(inp("x", "scaled_x", 0),), + output_attrs=(cap("theta", default=None, required=False), cap("position_offset", default=0, required=False), cap("backend", default=0, required=False)), + constraints=("pairwise rotary dims must match Cactus layout", "sin/cos cache or generated constants must be identified"), + ), + "rope_gptj": FusionGraph( + name="rope_gptj_decomposed", + target="rope_gptj", + cactus_inputs=("x",), + cactus_attrs=("theta", "position_offset", "rot_dim", "backend"), + nodes={"rope": n("rope", "rope_gptj_subgraph")}, + root="rope", + outputs=("rope",), + inputs=(inp("x", "rope", 0),), + output_attrs=(cap("theta", default=None, required=False), cap("position_offset", default=0), cap("rot_dim", default=0), cap("backend", default=0)), + constraints=("GPT-J interleaved rotary layout check",), + ), + "sum": direct("sum", ("x",), ("axis",), "aten.sum.dim_IntList", attr_map={"axis": "dim"}), + "mean": direct("mean", ("x",), ("axis", "keepdim"), "aten.mean.dim", attr_map={"axis": "dim"}), + "variance": direct("variance", ("x",), ("axis", "keepdim"), "aten.var.dim", attr_map={"axis": "dim"}), + "min": direct("min", ("x",), ("axis",), "aten.min.dim", attr_map={"axis": "dim"}), + "max": direct("max", ("x",), ("axis",), "aten.max.dim", attr_map={"axis": "dim"}), + "cumsum": direct("cumsum", ("x",), ("axis",), "aten.cumsum.default", attr_map={"axis": "dim"}), + "softmax": direct("softmax", ("x",), ("axis",), "aten._softmax.default", attr_map={"axis": "dim"}), + "attention": FusionGraph( + name="attention_direct_or_decomposed", + target="attention", + cactus_inputs=("query", "key", "value", "mask"), + cactus_attrs=("scale", "is_causal", "position_offset", "window_size", "backend", "use_mask", "additive_mask"), + nodes={ + "query": n("query", "external"), + "key": n("key", "external"), + "value": n("value", "external"), + "k_transpose": n("k_transpose", ("aten.transpose.int", "aten.permute.default"), optional=True, transparent=True), + "scores": n("scores", ("aten.matmul.default", "aten.bmm.default")), + "scale": n("scale", ("aten.mul.Scalar", "aten.div.Scalar"), optional=True), + "mask_apply": n("mask_apply", ("aten.add.Tensor", "aten.where.self"), optional=True), + "softmax": n("softmax", "aten._softmax.default"), + "context": n("context", ("aten.matmul.default", "aten.bmm.default")), + "direct_sdpa": n("direct_sdpa", "aten.scaled_dot_product_attention.default", optional=True), + }, + edges=(e("key", "k_transpose", 0, optional=True), e("query", "scores", 0), e("k_transpose", "scores", 1, optional=True), e("scores", "scale", 0, optional=True), e("scale", "mask_apply", 0, optional=True), e("mask_apply", "softmax", 0, optional=True), e("scores", "softmax", 0, optional=True), e("softmax", "context", 0), e("value", "context", 1)), + root="context", + outputs=("context",), + inputs=(inp("query", "scores", 0), inp("key", "scores", 1), inp("value", "context", 1), inp("mask", "mask_apply", 1, optional=True)), + output_attrs=(cap("scale", "scale", "scalar", default=None, required=False), cap("is_causal", "direct_sdpa", "is_causal", default=False, required=False), cap("position_offset", default=0), cap("window_size", default=0), cap("backend", default=0), cap("use_mask", default=False), cap("additive_mask", default=False)), + constraints=("q/k head_dim match", "softmax axis is score last dimension", "mask semantics are boolean/additive/causal", "RoPE, if fused separately, must already be applied to q and k"), + variants=("aten.scaled_dot_product_attention.default", "matmul_scale_mask_softmax_matmul", "bmm_scale_mask_softmax_bmm"), + ), + "kv_cache_state": source_op("kv_cache_state", ("max_seq_len", "num_kv_heads", "head_dim", "window_size", "sink_size", "num_slots"), "State allocation op; not discovered from ordinary ATen math subgraph."), + "kv_cache_append": direct("kv_cache_append", ("new_kv", "cache_state"), ("window_size", "sink_size"), "kv_cache_append_subgraph", optional_attr_defaults={"window_size": 0, "sink_size": 0}), + "attention_cached": FusionGraph( + name="attention_cached", + target="attention_cached", + cactus_inputs=("query", "key_new", "value_new", "k_cache_state", "v_cache_state"), + cactus_attrs=("scale", "position_offset", "window_size", "v_head_dim"), + nodes={"attention_cached": n("attention_cached", "cached_attention_subgraph")}, + root="attention_cached", + outputs=("attention_cached",), + inputs=tuple(inp(role, "attention_cached", i) for i, role in enumerate(("query", "key_new", "value_new", "k_cache_state", "v_cache_state"))), + output_attrs=(cap("scale", default=None, required=False), cap("position_offset", default=0), cap("window_size", default=0), cap("v_head_dim", default=None, required=False)), + constraints=("must identify cache state nodes", "must verify append/read ordering"), + ), + "conv_cache_state": source_op("conv_cache_state", ("window_size", "hidden_dim"), "State allocation op for streaming conv."), + "conv_cache_append": direct("conv_cache_append", ("new_data", "cache_state"), (), "conv_cache_append_subgraph"), + "conv_cache_initialize": direct("conv_cache_initialize", ("rows", "cache_state"), (), "conv_cache_initialize_subgraph"), + "recurrent_cache_state": source_op("recurrent_cache_state", ("shape", "dtype"), "State allocation op for recurrent kernels."), + "recurrent_cache_write": direct("recurrent_cache_write", ("new_value", "cache_input"), (), "recurrent_cache_write_subgraph"), + "rel_pos_bias": direct("rel_pos_bias", ("query", "relative_key"), ("scale",), "relative_position_bias_subgraph"), + "attention_int8_hybrid": direct( + "attention_int8_hybrid", + ("query", "key_new", "value_new"), + ("scale", "position_offset", "cache_len", "num_kv_heads", "head_dim", "window_size"), + "attention_int8_hybrid_subgraph", + notes="Requires quantized/cache-aware attention trace.", + ), + "relu": direct("relu", ("x",), (), "aten.relu.default"), + "silu": FusionGraph( + name="silu_direct_or_sigmoid_mul", + target="silu", + cactus_inputs=("x",), + cactus_attrs=(), + nodes={"sigmoid": n("sigmoid", "aten.sigmoid.default", optional=True), "mul": n("mul", "aten.mul.Tensor")}, + edges=(e("sigmoid", "mul", 0, optional=True),), + root="mul", + outputs=("mul",), + inputs=(inp("x", "mul", 1),), + constraints=("mul other input and sigmoid input must be same x",), + variants=("aten.silu.default", "sigmoid_mul"), + ), + "gelu": direct("gelu", ("x",), (), "aten.gelu.default"), + "gelu_erf": direct("gelu_erf", ("x",), (), "gelu_erf_subgraph", notes="Detect erf-based GELU decomposition."), + "sigmoid": direct("sigmoid", ("x",), (), "aten.sigmoid.default"), + "tanh": direct("tanh", ("x",), (), "aten.tanh.default"), + "glu": direct("glu", ("x",), ("axis",), "aten.glu.default", attr_map={"axis": "dim"}), + "conv1d_causal": direct("conv1d_causal", ("x", "weight"), ("kernel_size", "dilation"), "aten.convolution.default", notes="Requires padding/causal layout check."), + "conv1d_k3": direct("conv1d_k3", ("x", "weight"), ("stride",), "aten.convolution.default", optional_attr_defaults={"stride": 1}, notes="Requires weight kernel_size == 3."), + "conv1d_k7s3": direct("conv1d_k7s3", ("x", "weight", "bias"), (), "aten.convolution.default", notes="Requires kernel_size == 7 and stride == 3."), + "conv1d": direct("conv1d", ("x", "weight", "bias"), ("stride",), ("aten.conv1d.default", "aten.convolution.default"), optional_attr_defaults={"stride": 1}), + "conv1d_same_depthwise_k9": direct("conv1d_same_depthwise_k9", ("x", "weight", "bias"), (), "aten.convolution.default", notes="Requires groups == channels, kernel_size == 9, same padding."), + "conv1d_pointwise": direct("conv1d_pointwise", ("x", "weight", "bias"), (), "aten.convolution.default", notes="Requires kernel_size == 1."), + "conv2d_k3s2p1": direct("conv2d_k3s2p1", ("x", "weight", "bias"), (), "aten.convolution.default", notes="Requires kernel_size == 3, stride == 2, padding == 1."), + "conv2d_depthwise_k3s2p1": direct("conv2d_depthwise_k3s2p1", ("x", "weight", "bias"), (), "aten.convolution.default", notes="Requires depthwise groups and k3s2p1."), + "conv2d_pointwise_1x1": direct("conv2d_pointwise_1x1", ("x", "weight", "bias"), (), "aten.convolution.default", notes="Requires kernel_size == 1x1."), + "conv2d_k3s1p1": direct("conv2d_k3s1p1", ("x", "weight", "bias"), (), "aten.convolution.default", notes="Requires kernel_size == 3, stride == 1, padding == 1."), + "conv2d": direct("conv2d", ("x", "weight", "bias"), ("stride", "padding", "dilation", "groups"), ("aten.conv2d.default", "aten.convolution.default"), optional_attr_defaults={"stride": 1, "padding": 0, "dilation": 1, "groups": 1}), + "lstm_cell": direct("lstm_cell", ("input", "h_prev", "c_prev", "weight_ih", "weight_hh", "bias_ih", "bias_hh"), (), "lstm_cell_subgraph", notes="Requires gate split/order checks: input, forget, cell, output."), + "gated_deltanet_decode": direct("gated_deltanet_decode", ("query", "key", "value", "gate_log", "beta", "initial_state"), ("scale",), "gated_deltanet_decode_subgraph"), + "gated_deltanet_prefill": direct("gated_deltanet_prefill", ("query", "key", "value", "gate_log", "beta", "initial_state"), ("chunk_size", "scale"), "gated_deltanet_prefill_subgraph"), + "stft": direct("stft", ("x", "weight"), ("stride", "num_fft_bins"), ("aten.stft.default", "stft_subgraph")), + "altup_predict": direct("altup_predict", ("coefs", "streams"), (), "altup_predict_subgraph"), + "altup_correct": direct("altup_correct", ("coefs", "innovation", "predictions"), (), "altup_correct_subgraph"), + "gaussian_topk": direct("gaussian_topk", ("x",), ("ppf",), "gaussian_topk_subgraph"), + "moe_layer_gated": FusionGraph( + name="moe_layer_gated_topk_experts", + target="moe_layer_gated", + cactus_inputs=("hidden", "routing_probs", "topk_indices", "w1_weights", "w3_weights", "w2_weights"), + cactus_attrs=("num_experts", "num_experts_per_tok", "normalize_routing", "epsilon", "routed_scaling_factor"), + nodes={ + "router_logits": n("router_logits", ("aten.mm.default", "aten.matmul.default", "aten.linear.default")), + "routing_probs": n("routing_probs", ("aten._softmax.default", "aten.sigmoid.default")), + "topk": n("topk", "aten.topk.default"), + "topk_values": n("topk_values", ("operator.getitem", "")), + "topk_indices": n("topk_indices", ("operator.getitem", "")), + "renorm": n("renorm", ("aten.sum.dim_IntList", "aten.div.Tensor"), optional=True), + "expert_gate": n("expert_gate", ("aten.mm.default", "aten.matmul.default", "aten.linear.default"), repeated=True), + "expert_up": n("expert_up", ("aten.mm.default", "aten.matmul.default", "aten.linear.default"), repeated=True), + "expert_act": n("expert_act", ("aten.silu.default", "aten.gelu.default", "aten.relu.default"), repeated=True), + "expert_mul": n("expert_mul", "aten.mul.Tensor", repeated=True), + "expert_down": n("expert_down", ("aten.mm.default", "aten.matmul.default", "aten.linear.default"), repeated=True), + "weighted_expert": n("weighted_expert", "aten.mul.Tensor", repeated=True), + "combine": n("combine", ("aten.add.Tensor", "aten.index_put.default", "aten.scatter_add.default")), + }, + edges=(e("router_logits", "routing_probs", 0), e("routing_probs", "topk", 0), e("topk", "topk_values", 0, source_output_index=0), e("topk", "topk_indices", 0, source_output_index=1), e("expert_gate", "expert_act", 0, repeated=True), e("expert_act", "expert_mul", 0, repeated=True), e("expert_up", "expert_mul", 1, repeated=True), e("expert_mul", "expert_down", 0, repeated=True), e("expert_down", "weighted_expert", 0, repeated=True), e("topk_values", "weighted_expert", 1, repeated=True), e("weighted_expert", "combine", 0, repeated=True)), + root="combine", + outputs=("combine",), + inputs=(inp("hidden", "router_logits", 0), inp("routing_probs", "routing_probs", 0), inp("topk_indices", "topk_indices", 0), inp("w1_weights", "expert_gate", 1, variadic=True, repeated_group="experts"), inp("w3_weights", "expert_up", 1, variadic=True, repeated_group="experts"), inp("w2_weights", "expert_down", 1, variadic=True, repeated_group="experts")), + output_attrs=(cap("num_experts", default=None, required=False), cap("num_experts_per_tok", "topk", "k"), cap("normalize_routing", default=False), cap("epsilon", default=1e-6), cap("routed_scaling_factor", default=1.0)), + constraints=("topk getitem 0/1 must map values/indices", "expert branch count and weight order must be stable", "router hidden input must be shared with expert branches", "combine/scatter must preserve token order"), + repeated_groups=("experts",), + ), + "dense_mlp_tq_fused": direct("dense_mlp_tq_fused", ("hidden", "gate_weight", "up_weight", "down_weight"), ("product_scale",), "dense_mlp_tq_fused_subgraph", optional_attr_defaults={"product_scale": 1.0}), + "moe_layer_ungated": FusionGraph( + name="moe_layer_ungated_topk_experts", + target="moe_layer_ungated", + cactus_inputs=("hidden", "routing_probs", "topk_indices", "w1_weights", "w2_weights"), + cactus_attrs=("num_experts", "num_experts_per_tok", "normalize_routing", "epsilon", "routed_scaling_factor", "activation"), + nodes={ + "router_logits": n("router_logits", ("aten.mm.default", "aten.matmul.default", "aten.linear.default")), + "routing_probs": n("routing_probs", ("aten._softmax.default", "aten.sigmoid.default")), + "topk": n("topk", "aten.topk.default"), + "topk_values": n("topk_values", ("operator.getitem", "")), + "topk_indices": n("topk_indices", ("operator.getitem", "")), + "expert_up": n("expert_up", ("aten.mm.default", "aten.matmul.default", "aten.linear.default"), repeated=True), + "expert_act": n("expert_act", ("aten.gelu.default", "aten.relu.default", "aten.silu.default"), repeated=True), + "expert_down": n("expert_down", ("aten.mm.default", "aten.matmul.default", "aten.linear.default"), repeated=True), + "weighted_expert": n("weighted_expert", "aten.mul.Tensor", repeated=True), + "combine": n("combine", ("aten.add.Tensor", "aten.index_put.default", "aten.scatter_add.default")), + }, + edges=(e("router_logits", "routing_probs", 0), e("routing_probs", "topk", 0), e("topk", "topk_values", 0, source_output_index=0), e("topk", "topk_indices", 0, source_output_index=1), e("expert_up", "expert_act", 0, repeated=True), e("expert_act", "expert_down", 0, repeated=True), e("expert_down", "weighted_expert", 0, repeated=True), e("topk_values", "weighted_expert", 1, repeated=True), e("weighted_expert", "combine", 0, repeated=True)), + root="combine", + outputs=("combine",), + inputs=(inp("hidden", "router_logits", 0), inp("routing_probs", "routing_probs", 0), inp("topk_indices", "topk_indices", 0), inp("w1_weights", "expert_up", 1, variadic=True, repeated_group="experts"), inp("w2_weights", "expert_down", 1, variadic=True, repeated_group="experts")), + output_attrs=(cap("num_experts", default=None, required=False), cap("num_experts_per_tok", "topk", "k"), cap("normalize_routing", default=False), cap("epsilon", default=1e-6), cap("routed_scaling_factor", default=1.0), cap("activation", "expert_act", "op")), + constraints=("topk tuple outputs must be distinguished", "activation enum must be derived from expert_act op", "expert order must match weight lists"), + repeated_groups=("experts",), + ), + "sample": direct("sample", ("logits",), ("temperature", "top_p", "top_k"), "sampling_subgraph", optional_attr_defaults={"temperature": 0.6, "top_p": 0.95, "top_k": 20}), + "scatter_topk": direct("scatter_topk", ("indices", "values"), ("num_classes",), "scatter_topk_subgraph"), + "persistent": direct("persistent", ("source_node",), (), "persistent_subgraph"), + "rfft": direct("rfft", ("x",), (), "aten.fft_rfft.default"), + "irfft": direct("irfft", ("x",), ("output_length",), "aten.fft_irfft.default"), + "mel_filter_bank": source_op("mel_filter_bank", ("num_frequency_bins", "num_mel_filters", "min_frequency", "max_frequency", "sampling_rate", "norm_type", "scale_type"), "Generated constant/source op for audio preprocessing."), + "spectrogram": direct("spectrogram", ("waveform", "mel_filters"), ("frame_length", "hop_length", "fft_length", "power", "center", "pad_mode", "mel_floor", "log_mel_mode", "dither", "preemphasis", "remove_dc_offset"), "spectrogram_subgraph"), + "image_preprocess": direct("image_preprocess", ("pixel_input",), ("src_width", "src_height", "target_width", "target_height", "patch_size", "channels", "rescale_factor", "mean", "std_dev"), "image_preprocess_subgraph"), +} + + +ALL_FUSION_GRAPHS: tuple[FusionGraph, ...] = tuple(FUSION_GRAPHS.values()) diff --git a/python/cactus/transpiler/ir/fusions.py b/python/cactus/transpiler/ir/fusions.py deleted file mode 100644 index 3559660f1..000000000 --- a/python/cactus/transpiler/ir/fusions.py +++ /dev/null @@ -1,42 +0,0 @@ -import models, nodes - - -RMS_NORM = models.FusionGraph( - target="rms_norm", - root="weighted_output", - edges=( - models.Edge("weighted_output", 0, "normalized_x"), - models.Edge("normalized_x", 1, "inverse_rms"), - models.Edge("inverse_rms", 0, "variance_with_eps"), - models.Edge("variance_with_eps", 0, "variance"), - models.Edge("variance", 0, "squared_x"), - ), - inputs=( - models.InputSpec("normalized_x", 0), - models.InputSpec("weighted_output", 1), - ), - shared_inputs=( - ( - models.InputSpec("normalized_x", 0), - models.InputSpec("squared_x", 0), - ), - ), - output_attrs={ - "eps": models.AttrRef("variance_with_eps", "scalar"), - }, -) - - -ALL_FUSIONS = ( - RMS_NORM, -) - - -ROOT_TARGET_MAP = { - "aten.mul.Tensor": ( - RMS_NORM, - ), -} - -#TODO: Update with str paths and their respective fusion object -OPS_MAP: dict[str, list[models.FusionPattern]] = {} diff --git a/python/cactus/transpiler/ir/fusions_utils.py b/python/cactus/transpiler/ir/fusions_utils.py index df02ecbf0..d7219bc56 100644 --- a/python/cactus/transpiler/ir/fusions_utils.py +++ b/python/cactus/transpiler/ir/fusions_utils.py @@ -1,147 +1,436 @@ -import models, nodes, fusions -"""###################################### FUSION UTILS!!!!!!! ######################################""" -def node_match(node:models.Node, synth_node:models.NodeSpec) -> bool: - if node.underlying_op != synth_node.underlying_op: - return False - for key, expected_value in synth_node.required_attrs.items(): - if key not in node.normalized_attrs: - return False - if node.normalized_attrs[key] != expected_value: - return False - return True -def get_bound_parent(bindings: dict[str, models.Node], node_name: str, parent_index: int) -> models.Node | None: - if node_name not in bindings: - return None - node = bindings[node_name] - if parent_index < 0 or parent_index >= len(node.parents): - return None - return node.parents[parent_index] -def get_input_parent(bindings: dict[str, models.Node], input_spec: models.InputSpec) -> models.Node | None: - return get_bound_parent(bindings, input_spec.node, input_spec.parent_index) -#Match edges -def edges_match(node: models.Node, synth_graph: models.FusionGraph) -> dict[str, models.Node] | None: - bindings:dict[str,models.Node] = {} - if not node_match(node, nodes.NODE_MAP[synth_graph.root]): - return None - bindings[synth_graph.root] = node + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + +# import models, nodes, fusions + + + + + +# """###################################### FUSION UTILS!!!!!!! ######################################""" + +# def attr_match(expected_value, actual_value, node: models.Node) -> bool: +# if hasattr(expected_value, "matches"): +# return expected_value.matches(actual_value, node) + +# return expected_value == actual_value + +# def node_match(node:models.Node, synth_node:models.NodeSpec, bindings: dict[str, models.Node] | None = None) -> bool: +# underlying_op = synth_node.underlying_op + +# if isinstance(underlying_op, tuple): +# if node.underlying_op not in underlying_op: +# return False +# elif node.underlying_op != underlying_op: +# return False + +# allowed_value_kinds = getattr(synth_node, "allowed_value_kinds", ()) +# if allowed_value_kinds and getattr(node, "value_kind", None) not in allowed_value_kinds: +# return False + +# for key, expected_value in synth_node.required_attrs.items(): +# if key not in node.normalized_attrs: +# return False + +# if not attr_match(expected_value, node.normalized_attrs[key], node): +# return False + +# for constraint in getattr(synth_node, "tensor_constraints", ()): +# if not constraint.matches(node, bindings): +# return False + +# return True + +# #Returns the desired parent if it is in the parents list +# def get_bound_parent(bindings: dict[str, models.Node], node_name: str, parent_index: int) -> models.Node | None: +# if node_name not in bindings: +# return None + +# node = bindings[node_name] +# if parent_index < 0 or parent_index >= len(node.parents): +# return None + +# return node.parents[parent_index] + +# #Returns parent base on InputSpec object +# def get_input_parent(bindings: dict[str, models.Node], input_spec: models.InputSpec) -> models.Node | None: +# return get_bound_parent(bindings, input_spec.node, input_spec.parent_index) + +# def get_variadic_input_parents(bindings: dict[str, models.Node], variadic_input) -> list[models.Node] | None: +# if variadic_input.node not in bindings: +# return None + +# node = bindings[variadic_input.node] +# start = variadic_input.start_parent_index +# end = variadic_input.end_parent_index if variadic_input.end_parent_index is not None else len(node.parents) + +# if start < 0 or start > len(node.parents): +# return None + +# if end < start or end > len(node.parents): +# return None + +# parents = node.parents[start:end] + +# if len(parents) < variadic_input.min_count: +# return None + +# if variadic_input.max_count is not None and len(parents) > variadic_input.max_count: +# return None + +# return parents + +# def match_input_constraints(parent: models.Node, input_spec, bindings: dict[str, models.Node]) -> bool: +# allowed_value_kinds = getattr(input_spec, "allowed_value_kinds", ()) +# if allowed_value_kinds and getattr(parent, "value_kind", None) not in allowed_value_kinds: +# return False + +# for constraint in getattr(input_spec, "tensor_constraints", ()): +# if not constraint.matches(parent, bindings): +# return False + +# return True + +# #Match edges and returns binding +# def edges_match(node: models.Node, synth_graph: models.FusionGraph) -> dict[str, models.Node] | None: +# bindings:dict[str,models.Node] = {} + +# if not node_match(node, nodes.NODE_MAP[synth_graph.root]): +# return None + +# bindings[synth_graph.root] = node - for edge in synth_graph.edges: - parent = get_bound_parent(bindings, edge.from_node, edge.parent_index) - if parent is None: - return None +# for edge in synth_graph.edges: +# parent = get_bound_parent(bindings, edge.from_node, edge.parent_index) +# if parent is None: +# if getattr(edge, "required", True): +# return None +# continue + +# if edge.to_node not in nodes.NODE_MAP: +# return None + +# if not node_match(parent, nodes.NODE_MAP[edge.to_node], bindings): +# return None + +# to_output_index = getattr(edge, "to_output_index", None) +# if to_output_index is not None: +# tuple_source = getattr(parent, "tuple_source", None) +# if tuple_source is None or tuple_source.output_index != to_output_index: +# return None + +# if edge.to_node in bindings and bindings[edge.to_node] is not parent: +# return None + +# bindings[edge.to_node] = parent + +# return bindings + +# #Ensures inputs exist within the fusion +# def match_inputs_exist(bindings: dict[str, models.Node], synth_graph: models.FusionGraph) -> bool: +# for input_spec in synth_graph.inputs: +# parent = get_input_parent(bindings, input_spec) +# if parent is None: +# return False + +# if not match_input_constraints(parent, input_spec, bindings): +# return False + +# for variadic_input in getattr(synth_graph, "variadic_inputs", ()): +# parents = get_variadic_input_parents(bindings, variadic_input) +# if parents is None: +# return False + +# for parent in parents: +# if not match_input_constraints(parent, variadic_input, bindings): +# return False + +# return True + +# #Enusres all of a nodes inputs are included in the fusion +# def match_external_inputs_dec(bindings: dict[str, models.Node], synth_graph: models.FusionGraph) -> bool: +# internal_node_ids = {id(node) for node in bindings.values()} +# declared_input_ids = set() + +# for input_spec in synth_graph.inputs: +# parent = get_input_parent(bindings, input_spec) +# if parent is None: +# return False + +# declared_input_ids.add(id(parent)) + +# for variadic_input in getattr(synth_graph, "variadic_inputs", ()): +# parents = get_variadic_input_parents(bindings, variadic_input) +# if parents is None: +# return False + +# for parent in parents: +# declared_input_ids.add(id(parent)) + +# for node in bindings.values(): +# for parent in node.parents: +# if id(parent) not in internal_node_ids and id(parent) not in declared_input_ids: +# return False + +# return True + +# #Ensures named input roles point to the expected kind of value/tensor +# def match_input_roles(bindings: dict[str, models.Node], synth_graph: models.FusionGraph) -> bool: +# role_map = {} + +# for input_spec in synth_graph.inputs: +# role = getattr(input_spec, "role", None) +# if role is None: +# continue + +# parent = get_input_parent(bindings, input_spec) +# if parent is None: +# return False + +# role_map.setdefault(role, []) +# role_map[role].append(parent) + +# for variadic_input in getattr(synth_graph, "variadic_inputs", ()): +# role = getattr(variadic_input, "role", None) +# if role is None: +# continue + +# parents = get_variadic_input_parents(bindings, variadic_input) +# if parents is None: +# return False + +# role_map.setdefault(role, []) +# role_map[role].extend(parents) + +# for role_spec in getattr(synth_graph, "input_roles", ()): +# if role_spec.input_name not in role_map: +# return False + +# for parent in role_map[role_spec.input_name]: +# if not match_input_constraints(parent, role_spec, bindings): +# return False + +# return True + +# #Ensures ops with same inputs are actually pointing to same input node +# def match_shared_inputs(bindings: dict[str, models.Node], synth_graph: models.FusionGraph) -> bool: +# for left_input, right_input in synth_graph.shared_inputs: +# left_parent = get_input_parent(bindings, left_input) +# right_parent = get_input_parent(bindings, right_input) +# if left_parent is None or right_parent is None: +# return False + +# if left_parent is not right_parent: +# return False + +# return True + +# #Ensures that a node within fusion is not called by a differnt node not in the fusion +# def match_no_external_internal_children(bindings: dict[str, models.Node], synth_graph: models.FusionGraph) -> bool: +# if synth_graph.root not in bindings: +# return False - if edge.to_node not in nodes.NODE_MAP: - return None +# root_node = bindings[synth_graph.root] +# internal_node_ids = {id(node) for node in bindings.values()} +# allow_root_external_children = getattr(synth_graph, "allow_root_external_children", True) - if not node_match(parent, nodes.NODE_MAP[edge.to_node]): - return None +# for node in bindings.values(): +# if node is root_node and allow_root_external_children: +# continue - if edge.to_node in bindings and bindings[edge.to_node] is not parent: - return None +# for child in node.children: +# if id(child) not in internal_node_ids: +# return False - bindings[edge.to_node] = parent +# return True - return bindings +# def is_attr_ref(value) -> bool: +# return hasattr(value, "node") and hasattr(value, "attr") +# def match_output_attrs(bindings: dict[str, models.Node], synth_graph: models.FusionGraph) -> bool: +# for attr_value in synth_graph.output_attrs.values(): +# if not is_attr_ref(attr_value): +# continue -def match_inputs_exist(bindings: dict[str, models.Node], synth_graph: models.FusionGraph) -> bool: - for input_spec in synth_graph.inputs: - if get_input_parent(bindings, input_spec) is None: - return False +# if attr_value.node not in bindings: +# if getattr(attr_value, "required", True): +# return False +# continue - return True +# node = bindings[attr_value.node] +# if attr_value.attr not in node.normalized_attrs: +# if getattr(attr_value, "required", True): +# return False +# continue -def match_external_inputs_dec(bindings: dict[str, models.Node], synth_graph: models.FusionGraph) -> bool: - internal_node_ids = {id(node) for node in bindings.values()} - declared_input_ids = set() +# return True - for input_spec in synth_graph.inputs: - parent = get_input_parent(bindings, input_spec) - if parent is None: - return False +# def match_optional_nodes(bindings: dict[str, models.Node], synth_graph: models.FusionGraph) -> bool: +# for optional_node in getattr(synth_graph, "optional_nodes", ()): +# if optional_node.node not in bindings: +# continue - declared_input_ids.add(id(parent)) +# node = bindings[optional_node.node] +# if optional_node.bypass_parent_index < 0 or optional_node.bypass_parent_index >= len(node.parents): +# return False - for node in bindings.values(): - for parent in node.parents: - if id(parent) not in internal_node_ids and id(parent) not in declared_input_ids: - return False +# return True - return True +# def match_tuple_outputs(bindings: dict[str, models.Node], synth_graph: models.FusionGraph) -> bool: +# for tuple_output in getattr(synth_graph, "tuple_outputs", ()): +# if tuple_output.node not in bindings or tuple_output.source_node not in bindings: +# return False -def match_shared_inputs(bindings: dict[str, models.Node], synth_graph: models.FusionGraph) -> bool: - for left_input, right_input in synth_graph.shared_inputs: - left_parent = get_input_parent(bindings, left_input) - right_parent = get_input_parent(bindings, right_input) - if left_parent is None or right_parent is None: - return False +# node = bindings[tuple_output.node] +# source_node = bindings[tuple_output.source_node] +# tuple_source = getattr(node, "tuple_source", None) - if left_parent is not right_parent: - return False +# if tuple_source is None: +# return False - return True +# if tuple_source.source_node_name != source_node.name: +# return False -def match_no_external_internal_children(bindings: dict[str, models.Node], synth_graph: models.FusionGraph) -> bool: - if synth_graph.root not in bindings: - return False +# if tuple_source.output_index != tuple_output.output_index: +# return False - root_node = bindings[synth_graph.root] - internal_node_ids = {id(node) for node in bindings.values()} +# return True - for node in bindings.values(): - if node is root_node: - continue +# def match_repeated_subgraphs(bindings: dict[str, models.Node], synth_graph: models.FusionGraph) -> bool: +# repeated_bindings = getattr(synth_graph, "metadata", {}).get("repeated_bindings", {}) - for child in node.children: - if id(child) not in internal_node_ids: - return False +# for repeated in getattr(synth_graph, "repeated_subgraphs", ()): +# if repeated.anchor_node is not None and repeated.anchor_node not in bindings: +# return False - return True +# matches = repeated_bindings.get(repeated.name, []) -def match_output_attrs(bindings: dict[str, models.Node], synth_graph: models.FusionGraph) -> bool: - for attr_value in synth_graph.output_attrs.values(): - if not isinstance(attr_value, models.AttrRef): - continue +# if len(matches) < repeated.min_count: +# return False - if attr_value.node not in bindings: - return False +# if repeated.max_count is not None and len(matches) > repeated.max_count: +# return False - node = bindings[attr_value.node] - if attr_value.attr not in node.normalized_attrs: - return False +# return True - return True +# #Update this with any additional matching functions if any are added +# MATCHERS = ( +# match_inputs_exist, +# match_external_inputs_dec, +# match_input_roles, +# match_shared_inputs, +# match_no_external_internal_children, +# match_output_attrs, +# match_optional_nodes, +# match_tuple_outputs, +# match_repeated_subgraphs, +# ) -#Update this with any additional matching functions if any are added -MATCHERS = ( - match_inputs_exist, - match_external_inputs_dec, - match_shared_inputs, - match_no_external_internal_children, - match_output_attrs, -) +# def fusion_match(node: models.Node) -> models.FusionResult | None: -def fusion_match(node: models.Node) -> models.FusionGraph | None: +# for fusion in fusions.ROOT_TARGET_MAP.get(node.underlying_op, ()): +# bindings = edges_match(node, fusion) +# if bindings != None and all(matcher(bindings, fusion) for matcher in MATCHERS): +# return models.FusionResult.from_fusion(fusion, bindings) - for fusion in fusions.ROOT_TARGET_MAP.get(node.underlying_op, ()): - bindings = edges_match(node, fusion) - if bindings != None and all(matcher(bindings, fusion) for matcher in MATCHERS): - #Will change this to return FusionResult object which will have a list of all nodes being fused - return fusion +# return None diff --git a/python/cactus/transpiler/ir/models.py b/python/cactus/transpiler/ir/models.py index 2dca41c07..f56548ae0 100644 --- a/python/cactus/transpiler/ir/models.py +++ b/python/cactus/transpiler/ir/models.py @@ -1,128 +1,834 @@ -from dataclasses import dataclass, field -from typing import Any, Optional -import converter.models as CVModels -import constants -@dataclass(slots=True) -class Node: - layer: Optional[CVModels.LayerRecord] - underlying_op: str - normalized_attrs: dict[str, Any] - parents: list["Node"] - children: list["Node"] -@dataclass(slots=True) -class Graph: - name_map:dict[str, Node] - source_nodes: list[Node] - consumed_ids: set[str] - @classmethod - def from_layer_map(cls, map:CVModels.LayerMap) -> "Graph": - return generate_graph(map) -@dataclass(slots=True) -class NodeSpec: - underlying_op:str - required_attrs: dict[str, Any] = field(default_factory=dict) -@dataclass(slots=True) -class Edge: - from_node:str - parent_index:int - to_node:str -@dataclass(slots=True) -class InputSpec: - node:str - parent_index:int -@dataclass(slots=True) -class AttrRef: - node:str - attr:str -@dataclass(slots=True) -class FusionGraph: - target:str - root:str - edges:tuple[Edge, ...] - inputs:tuple[InputSpec, ...] = () - shared_inputs:tuple[tuple[InputSpec, InputSpec], ...] = () - output_attrs:dict[str, Any] = field(default_factory=dict) -class FusionResult: - node:Node - edges_consumed:list[str] #Can be changed to -"""###################################### MODEL UTILS!!!!!!! ######################################""" -def extract_node_names(value: Any) -> list[str]: - - if isinstance(value, dict): - refs = [] - if "node" in value: - refs.append(value["node"]) - for item in value.values(): - refs.extend(extract_node_names(item)) - return refs - if isinstance(value, list): - refs = [] - for item in value: - refs.extend(extract_node_names(item)) - return refs - return [] -def extract_attrs(layer_: CVModels.LayerRecord) -> dict[str, Any]: - - def contains_node_ref(value: Any) -> bool: - if isinstance(value, dict): - if "node" in value: - return True - return any(contains_node_ref(item) for item in value.values()) - if isinstance(value, list): - return any(contains_node_ref(item) for item in value) - return False - attrs = {} - if isinstance(layer_.kwargs, dict): - attrs.update(layer_.kwargs) - attr_names = constants.LAYER_ATTRS_MAP.get(layer_.target, []) - positional_attrs = [item for item in layer_.args if not contains_node_ref(item)] - for i, attr_name in enumerate(attr_names): - if i < len(positional_attrs) and attr_name not in attrs: - attrs[attr_name] = positional_attrs[i] - return attrs -def generate_graph(map: CVModels.LayerMap): - temp_map:dict[str, Node] = {} - temp_source = [] - for layer_ in map.nodes: - temp = Node(layer=layer_, normalized_attrs=extract_attrs(layer_), underlying_op= layer_.target, parents=[], children=[]) - temp_map[layer_.name] = temp - temp_source.append(temp) if layer_.node_type == "output" else None - for layer_ in map.nodes: - for name in layer_.users: - temp_map[layer_.name].children.append(temp_map[name]) - for name in extract_node_names(layer_.args): - temp_map[layer_.name].parents.append(temp_map[name]) - return Graph(name_map=temp_map, source_nodes=temp_source, consumed_ids=set()) + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + +# from __future__ import annotations + +# from dataclasses import dataclass, field +# from typing import Any, Callable, Optional +# import converter.models as CVModels +# import constants + + +# class ValueKind: +# UNKNOWN = "unknown" +# PARAMETER = "parameter" +# BUFFER = "buffer" +# USER_INPUT = "user_input" +# LIFTED_CONSTANT = "lifted_constant" +# ACTIVATION = "activation" +# OUTPUT = "output" + + +# @dataclass(slots=True) +# class TensorMeta: +# shape: list[int] = field(default_factory=list) +# dtype: str | None = None +# device: str | None = None +# requires_grad: bool | None = None +# stride: list[int] | None = None + +# @property +# def rank(self) -> int: +# return len(self.shape) + +# def dim(self, index: int) -> int | None: +# if index < 0: +# index += len(self.shape) + +# if index < 0 or index >= len(self.shape): +# return None + +# return self.shape[index] + + +# @dataclass(slots=True) +# class TupleOutputRef: +# source_node_name: str +# output_index: int + + +# @dataclass(slots=True) +# class Node: +# name: str +# underlying_op: str +# normalized_attrs: dict[str, Any] +# parents: list["Node"] +# children: list["Node"] +# tensor_meta: TensorMeta | None = None +# tuple_output_metas: list[TensorMeta | None] = field(default_factory=list) +# value_kind: str = ValueKind.UNKNOWN +# tuple_source: TupleOutputRef | None = None +# raw_layer: Optional[CVModels.LayerRecord] = None + + +# @dataclass(frozen=True, slots=True) +# class AttrPredicate: +# description: str +# predicate: Callable[[Any, Node], bool] = field(repr=False, compare=False) + +# def matches(self, value: Any, node: Node) -> bool: +# return self.predicate(value, node) + + +# def attr_any() -> AttrPredicate: +# return AttrPredicate("any value", lambda _value, _node: True) + + +# def attr_in(values: tuple[Any, ...]) -> AttrPredicate: +# return AttrPredicate(f"in {values}", lambda value, _node: value in values) + + +# def attr_rank_dim(dim_values: tuple[int, ...]) -> AttrPredicate: +# def matches(value: Any, node: Node) -> bool: +# if value in dim_values: +# return True + +# if node.tensor_meta is None: +# return False + +# normalized_values = tuple(dim + node.tensor_meta.rank if dim < 0 else dim for dim in dim_values) +# normalized_value = value + node.tensor_meta.rank if isinstance(value, int) and value < 0 else value +# return normalized_value in normalized_values + +# return AttrPredicate(f"rank-aware dim in {dim_values}", matches) + + +# @dataclass(frozen=True, slots=True) +# class DimRef: +# node: str +# dim: int + + +# @dataclass(frozen=True, slots=True) +# class TensorConstraint: +# rank: int | None = None +# min_rank: int | None = None +# max_rank: int | None = None +# dtype: str | None = None +# dim_equals: tuple[tuple[int, int], ...] = () +# same_dim_as: tuple[tuple[int, DimRef], ...] = () + +# def matches(self, node: Node, bindings: dict[str, Node] | None = None) -> bool: +# if node.tensor_meta is None: +# return False + +# meta = node.tensor_meta + +# if self.rank is not None and meta.rank != self.rank: +# return False + +# if self.min_rank is not None and meta.rank < self.min_rank: +# return False + +# if self.max_rank is not None and meta.rank > self.max_rank: +# return False + +# if self.dtype is not None and meta.dtype != self.dtype: +# return False + +# for dim_index, expected in self.dim_equals: +# if meta.dim(dim_index) != expected: +# return False + +# if bindings is None and self.same_dim_as: +# return False + +# for dim_index, ref in self.same_dim_as: +# other = bindings.get(ref.node) if bindings is not None else None +# if other is None or other.tensor_meta is None: +# return False + +# if meta.dim(dim_index) != other.tensor_meta.dim(ref.dim): +# return False + +# return True + + +# @dataclass(slots=True) +# class NodeSpec: +# underlying_op: str | tuple[str, ...] +# required_attrs: dict[str, Any] = field(default_factory=dict) +# allowed_value_kinds: tuple[str, ...] = () +# tensor_constraints: tuple[TensorConstraint, ...] = () +# optional: bool = False + + +# @dataclass(slots=True) +# class Edge: +# from_node: str +# parent_index: int +# to_node: str +# required: bool = True +# to_output_index: int | None = None + + +# @dataclass(slots=True) +# class InputSpec: +# node: str +# parent_index: int +# role: str | None = None +# allowed_value_kinds: tuple[str, ...] = () +# tensor_constraints: tuple[TensorConstraint, ...] = () + + +# @dataclass(slots=True) +# class VariadicInputSpec: +# node: str +# start_parent_index: int = 0 +# end_parent_index: int | None = None +# min_count: int = 0 +# max_count: int | None = None +# role: str | None = None +# allowed_value_kinds: tuple[str, ...] = () +# tensor_constraints: tuple[TensorConstraint, ...] = () + + +# @dataclass(slots=True) +# class InputRoleSpec: +# input_name: str +# role: str +# allowed_value_kinds: tuple[str, ...] = () +# tensor_constraints: tuple[TensorConstraint, ...] = () + + +# @dataclass(slots=True) +# class OptionalNodeSpec: +# node: str +# bypass_parent_index: int = 0 + + +# @dataclass(slots=True) +# class TupleOutputSpec: +# node: str +# source_node: str +# output_index: int + + +# @dataclass(slots=True) +# class RepeatedSubgraphSpec: +# name: str +# pattern: "FusionGraph" +# min_count: int = 1 +# max_count: int | None = None +# anchor_node: str | None = None + + +# @dataclass(slots=True) +# class AttrRef: +# node: str +# attr: str +# default: Any = None +# required: bool = True + + +# @dataclass(slots=True) +# class FusionGraph: +# target: str +# root: str +# edges: tuple[Edge, ...] +# inputs: tuple[InputSpec, ...] = () +# shared_inputs: tuple[tuple[InputSpec, InputSpec], ...] = () +# output_attrs: dict[str, Any] = field(default_factory=dict) +# optional_nodes: tuple[OptionalNodeSpec, ...] = () +# variadic_inputs: tuple[VariadicInputSpec, ...] = () +# input_roles: tuple[InputRoleSpec, ...] = () +# tuple_outputs: tuple[TupleOutputSpec, ...] = () +# repeated_subgraphs: tuple[RepeatedSubgraphSpec, ...] = () +# metadata: dict[str, Any] = field(default_factory=dict) +# allow_root_external_children: bool = True + + +# @dataclass(slots=True) +# class FusionResult: +# node: Node +# bindings: dict[str, Node] +# internal_nodes: list[Node] +# external_input: list[Node] +# external_consumers: list[Node] +# edges_consumed: list[str] +# fusion: FusionGraph | None = None +# variadic_external_inputs: dict[str, list[Node]] = field(default_factory=dict) +# input_role_nodes: dict[str, list[Node]] = field(default_factory=dict) + +# @classmethod +# def from_fusion(cls, fusion: FusionGraph, bindings: dict[str, Node]): +# return build_fusion_result(fusion, bindings) + + +# @dataclass(slots=True) +# class Graph: +# name_map: dict[str, Node] +# source_nodes: list[Node] +# consumed_ids: set[str] +# fusedNodes: list[FusionResult] = field(default_factory=list) + +# @classmethod +# def from_layer_map(cls, map: CVModels.LayerMap) -> "Graph": +# return generate_graph(map) + +# def queue_fusion(self, fusion_result: FusionResult) -> None: +# self.fusedNodes.append(fusion_result) + +# def update_graph_with_fusions(self) -> bool: +# if not len(self.fusedNodes): +# return False + +# update_graph(self) +# return True + + +# """###################################### MODEL UTILS!!!!!!! ######################################""" + +# # TODO: Understand +# def extract_node_names(value: Any) -> list[str]: +# if isinstance(value, dict): +# refs = [] +# if "node" in value: +# refs.append(value["node"]) + +# for item in value.values(): +# refs.extend(extract_node_names(item)) + +# return refs + +# if isinstance(value, list): +# refs = [] +# for item in value: +# refs.extend(extract_node_names(item)) +# return refs + +# return [] + +# # TODO: Understand +# def extract_attrs(layer_: CVModels.LayerRecord) -> dict[str, Any]: +# def contains_node_ref(value: Any) -> bool: +# if isinstance(value, dict): +# if "node" in value: +# return True +# return any(contains_node_ref(item) for item in value.values()) + +# if isinstance(value, list): +# return any(contains_node_ref(item) for item in value) + +# return False + +# attrs = {} + +# if isinstance(layer_.kwargs, dict): +# attrs.update(layer_.kwargs) + +# attr_names = constants.LAYER_ATTRS_MAP.get(layer_.target, []) + +# if None in attr_names: +# for i, attr_name in enumerate(attr_names): +# if attr_name is None or i >= len(layer_.args): +# continue + +# if contains_node_ref(layer_.args[i]) or attr_name in attrs: +# continue + +# attrs[attr_name] = layer_.args[i] +# else: +# positional_attrs = [item for item in layer_.args if not contains_node_ref(item)] + +# for i, attr_name in enumerate(attr_names): +# if i < len(positional_attrs) and attr_name not in attrs: +# attrs[attr_name] = positional_attrs[i] + +# return attrs + +# # TODO: Understand +# def attr_value_matches(expected_value: Any, actual_value: Any, node: Node) -> bool: +# if isinstance(expected_value, AttrPredicate): +# return expected_value.matches(actual_value, node) + +# return actual_value == expected_value + +# # TODO: Understand +# def node_spec_matches(node: Node, spec: NodeSpec, bindings: dict[str, Node] | None = None) -> bool: +# allowed_ops = spec.underlying_op if isinstance(spec.underlying_op, tuple) else (spec.underlying_op,) +# if node.underlying_op not in allowed_ops: +# return False + +# if spec.allowed_value_kinds and node.value_kind not in spec.allowed_value_kinds: +# return False + +# for key, expected_value in spec.required_attrs.items(): +# if key not in node.normalized_attrs: +# return False + +# if not attr_value_matches(expected_value, node.normalized_attrs[key], node): +# return False + +# for constraint in spec.tensor_constraints: +# if not constraint.matches(node, bindings): +# return False + +# return True + +# # TODO: Understand +# def _plain(value: Any) -> Any: +# if hasattr(value, "model_dump"): +# return value.model_dump() + +# if hasattr(value, "dict"): +# return value.dict() + +# return value + +# # TODO: Understand +# def _parse_one_tensor_meta(value: Any) -> TensorMeta | None: +# value = _plain(value) + +# if not isinstance(value, dict): +# return None + +# shape = value.get("shape") +# if shape is None: +# return None + +# return TensorMeta( +# shape=[int(item) for item in shape], +# dtype=value.get("dtype"), +# device=value.get("device"), +# requires_grad=value.get("requires_grad"), +# stride=value.get("stride"), +# ) + +# # TODO: Optimize this function further. +# def extract_tensor_metas(layer_: CVModels.LayerRecord) -> tuple[TensorMeta | None, list[TensorMeta | None]]: +# meta = _plain(layer_.tensor_output_meta) + +# single_meta = _parse_one_tensor_meta(meta) +# if single_meta is not None: +# return single_meta, [] + +# if isinstance(meta, list): +# tuple_metas = [_parse_one_tensor_meta(item) for item in meta] +# if tuple_metas: +# return None, tuple_metas + +# return None, [] + +# # TODO: Optimize this function further. +# def infer_value_kind(layer_: CVModels.LayerRecord) -> str: +# name = layer_.name +# target = layer_.target + +# if layer_.node_type == "output": +# return ValueKind.OUTPUT + +# if layer_.node_type == "placeholder": +# if name.startswith("p_") or "_weight" in name or name.endswith("weight"): +# return ValueKind.PARAMETER + +# if name.startswith("b_") or "buffer" in name: +# return ValueKind.BUFFER + +# return ValueKind.USER_INPUT + +# if layer_.node_type == "get_attr": +# return ValueKind.PARAMETER + +# if target in { +# "aten.scalar_tensor.default", +# "aten.full.default", +# "aten.arange.start_step", +# }: +# return ValueKind.LIFTED_CONSTANT + +# return ValueKind.ACTIVATION + +# # TODO: Optimize this function further. +# def extract_tuple_source(layer_: CVModels.LayerRecord) -> TupleOutputRef | None: +# target = layer_.target +# if target not in {"operator.getitem", ""}: +# return None + +# node_refs = extract_node_names(layer_.args) +# if not node_refs: +# return None + +# output_index = None +# if isinstance(layer_.args, list) and len(layer_.args) >= 2 and isinstance(layer_.args[1], int): +# output_index = layer_.args[1] + +# if output_index is None: +# return None + +# return TupleOutputRef(source_node_name=node_refs[0], output_index=output_index) + + +# def append_unique(nodes: list[Node], node: Node) -> None: +# if all(existing is not node for existing in nodes): +# nodes.append(node) + + +# def resolve_attr_ref(attr_name: str, attr_value: Any, bindings: dict[str, Node]) -> Any: +# if not isinstance(attr_value, AttrRef): +# return attr_value + +# if attr_value.node not in bindings: +# if attr_value.required: +# raise ValueError(f"Output attr {attr_name} references unbound node {attr_value.node}") +# return attr_value.default + +# bound_node = bindings[attr_value.node] +# if attr_value.attr not in bound_node.normalized_attrs: +# if attr_value.required: +# raise ValueError(f"Output attr {attr_name} references missing attr {attr_value.attr}") +# return attr_value.default + +# return bound_node.normalized_attrs[attr_value.attr] + + +# # TODO: Optimize this function further. +# def build_fusion_result(fusion: FusionGraph, bindings: dict[str, Node]) -> FusionResult: +# root_node = bindings[fusion.root] +# internal_nodes = [] + +# for node in bindings.values(): +# append_unique(internal_nodes, node) + +# internal_node_ids = {id(node) for node in internal_nodes} +# output_attrs = {} + +# for attr_name, attr_value in fusion.output_attrs.items(): +# output_attrs[attr_name] = resolve_attr_ref(attr_name, attr_value, bindings) + +# external_input = [] +# input_role_nodes = {} + +# for input_spec in fusion.inputs: +# if input_spec.node not in bindings: +# raise ValueError(f"Fusion input references unbound node {input_spec.node}") + +# node = bindings[input_spec.node] +# if input_spec.parent_index < 0 or input_spec.parent_index >= len(node.parents): +# raise ValueError(f"Fusion input has invalid parent index {input_spec.parent_index}") + +# parent = node.parents[input_spec.parent_index] +# append_unique(external_input, parent) + +# if input_spec.role is not None: +# input_role_nodes.setdefault(input_spec.role, []) +# append_unique(input_role_nodes[input_spec.role], parent) + +# variadic_external_inputs = {} + +# for variadic_input in fusion.variadic_inputs: +# if variadic_input.node not in bindings: +# raise ValueError(f"Variadic fusion input references unbound node {variadic_input.node}") + +# node = bindings[variadic_input.node] +# end = variadic_input.end_parent_index if variadic_input.end_parent_index is not None else len(node.parents) +# parents = node.parents[variadic_input.start_parent_index:end] + +# if len(parents) < variadic_input.min_count: +# raise ValueError(f"Variadic fusion input {variadic_input.node} has too few parents") + +# if variadic_input.max_count is not None and len(parents) > variadic_input.max_count: +# raise ValueError(f"Variadic fusion input {variadic_input.node} has too many parents") + +# role = variadic_input.role or variadic_input.node +# variadic_external_inputs[role] = [] + +# for parent in parents: +# append_unique(external_input, parent) +# append_unique(variadic_external_inputs[role], parent) + +# if variadic_input.role is not None: +# input_role_nodes.setdefault(variadic_input.role, []) +# append_unique(input_role_nodes[variadic_input.role], parent) + +# external_consumers = [] + +# for node in internal_nodes: +# for child in node.children: +# if id(child) in internal_node_ids: +# continue + +# if node is not root_node or not fusion.allow_root_external_children: +# raise ValueError("Cannot fuse node with external consumers outside the fusion boundary") + +# append_unique(external_consumers, child) + +# edges_consumed = [] + +# for child in internal_nodes: +# for parent in child.parents: +# if id(parent) in internal_node_ids: +# edges_consumed.append(f"{parent.name}->{child.name}") + +# fused_node = Node( +# name=f"{fusion.target}_{root_node.name}", +# underlying_op=fusion.target, +# normalized_attrs=output_attrs, +# parents=external_input, +# children=external_consumers, +# tensor_meta=root_node.tensor_meta, +# tuple_output_metas=root_node.tuple_output_metas.copy(), +# value_kind=ValueKind.ACTIVATION, +# tuple_source=None, +# raw_layer=None, +# ) + +# return FusionResult( +# node=fused_node, +# bindings=bindings, +# internal_nodes=internal_nodes, +# external_input=external_input, +# external_consumers=external_consumers, +# edges_consumed=edges_consumed, +# fusion=fusion, +# variadic_external_inputs=variadic_external_inputs, +# input_role_nodes=input_role_nodes, +# ) + +# # TODO: Optimize this function further. +# def generate_graph(map: CVModels.LayerMap) -> Graph: +# temp_map: dict[str, Node] = {} +# temp_source = [] + +# for layer_ in map.nodes: +# tensor_meta, tuple_output_metas = extract_tensor_metas(layer_) +# temp = Node( +# name=layer_.name, +# normalized_attrs=extract_attrs(layer_), +# underlying_op=layer_.target, +# parents=[], +# children=[], +# tensor_meta=tensor_meta, +# tuple_output_metas=tuple_output_metas, +# value_kind=infer_value_kind(layer_), +# tuple_source=extract_tuple_source(layer_), +# raw_layer=layer_, +# ) +# temp_map[layer_.name] = temp +# temp_source.append(temp) if layer_.node_type == "output" else None + +# for layer_ in map.nodes: +# for name in layer_.users: +# temp_map[layer_.name].children.append(temp_map[name]) + +# for name in extract_node_names(layer_.args): +# temp_map[layer_.name].parents.append(temp_map[name]) + +# for node in temp_map.values(): +# if node.tuple_source is None: +# continue + +# source_node = temp_map.get(node.tuple_source.source_node_name) +# output_index = node.tuple_source.output_index +# if source_node is None or output_index < 0 or output_index >= len(source_node.tuple_output_metas): +# continue + +# node.tensor_meta = source_node.tuple_output_metas[output_index] + +# return Graph(name_map=temp_map, source_nodes=temp_source, consumed_ids=set(), fusedNodes=[]) + + +# # TODO: Optimize this function further. +# def update_graph(graph: Graph) -> None: +# replacement_by_id: dict[int, Node] = {} +# fused_nodes = [] + +# for fusion_result in graph.fusedNodes: +# append_unique(fused_nodes, fusion_result.node) + +# for internal_node in fusion_result.internal_nodes: +# internal_id = id(internal_node) +# existing_replacement = replacement_by_id.get(internal_id) + +# if existing_replacement is not None and existing_replacement is not fusion_result.node: +# raise ValueError(f"Node {internal_node.name} is consumed by multiple fusions") + +# replacement_by_id[internal_id] = fusion_result.node + +# def resolve_node(node: Node) -> Node: +# seen_ids = set() + +# while id(node) in replacement_by_id and id(node) not in seen_ids: +# seen_ids.add(id(node)) +# node = replacement_by_id[id(node)] + +# return node + +# final_nodes = [] + +# for node in graph.name_map.values(): +# if id(node) not in replacement_by_id: +# append_unique(final_nodes, node) + +# for node in fused_nodes: +# append_unique(final_nodes, node) + +# new_name_map = {} + +# for node in final_nodes: +# if node.name in new_name_map and new_name_map[node.name] is not node: +# raise ValueError(f"Duplicate node name after fusion: {node.name}") + +# new_name_map[node.name] = node + +# final_node_ids = {id(node) for node in final_nodes} +# fusion_result_by_node_id = {id(fusion_result.node): fusion_result for fusion_result in graph.fusedNodes} + +# for node in final_nodes: +# fusion_result = fusion_result_by_node_id.get(id(node)) +# old_parents = fusion_result.external_input if fusion_result is not None else node.parents +# new_parents = [] + +# for parent in old_parents: +# resolved_parent = resolve_node(parent) + +# if resolved_parent is node: +# continue + +# if id(resolved_parent) not in final_node_ids: +# continue + +# new_parents.append(resolved_parent) + +# node.parents = new_parents +# node.children = [] + +# if node.tuple_source is not None: +# source_node = graph.name_map.get(node.tuple_source.source_node_name) +# if source_node is not None: +# resolved_source = resolve_node(source_node) +# if resolved_source is not source_node: +# node.tuple_source = TupleOutputRef( +# source_node_name=resolved_source.name, +# output_index=node.tuple_source.output_index, +# ) + +# for node in final_nodes: +# for parent in node.parents: +# append_unique(parent.children, node) + +# new_source_nodes = [] + +# for source_node in graph.source_nodes: +# resolved_source = resolve_node(source_node) + +# if id(resolved_source) in final_node_ids: +# append_unique(new_source_nodes, resolved_source) + +# graph.name_map = new_name_map +# graph.source_nodes = new_source_nodes +# graph.consumed_ids.clear() +# graph.fusedNodes.clear() \ No newline at end of file diff --git a/python/cactus/transpiler/ir/nodes.py b/python/cactus/transpiler/ir/nodes.py deleted file mode 100644 index 39561b473..000000000 --- a/python/cactus/transpiler/ir/nodes.py +++ /dev/null @@ -1,3 +0,0 @@ -import models - -NODE_MAP = {} \ No newline at end of file diff --git a/python/cactus/transpiler/ir/simplify_ir.py b/python/cactus/transpiler/ir/simplify_ir.py index e2eaba665..b77f1f04a 100644 --- a/python/cactus/transpiler/ir/simplify_ir.py +++ b/python/cactus/transpiler/ir/simplify_ir.py @@ -4,232 +4,12 @@ on the DAG to generate an in-order execution list. """ -import models, fusions_utils, fusions +import models import converter.models as CVModels -def try_match_from_node(node: models.Node, graph: models.Graph) -> models.FusionResult | None: - if node.layer.name in graph.consumed_ids: - return None - - best_match = None - best_match_len = 0 - checked_candidates = set() - - for path in fusions_utils.dfs_path_gen([node], graph): - for curr_len in range(len(path), 0, -1): - candidate_path = path[:curr_len] - candidate_ids = tuple(candidate_node.layer.name for candidate_node in candidate_path) - - if candidate_ids in checked_candidates: - continue - checked_candidates.add(candidate_ids) - - ops_key = "->".join(candidate_node.layer.target for candidate_node in candidate_path) - if ops_key not in fusions_utils.OPS_MAP: - continue - - fusion = fusions_utils.fusion_match(ops_key, candidate_path) - if fusion is None: - continue - - if curr_len > best_match_len: - best_match = models.FusionResult.from_match(fusion, candidate_path) - best_match_len = curr_len - - if best_match_len == len(path): - return best_match - - return best_match - - -def rev_top_sort(graph:models.Graph) -> models.SimplifiedGraph: - worklist = list(graph.source_nodes) - visited = set() - fused_results = [] - - while worklist: - node = worklist.pop() - node_id = node.layer.name - - if node_id in visited: - continue - visited.add(node_id) - - if node_id in graph.consumed_ids: - continue - - match = try_match_from_node(node, graph) - - if match is not None: - fused_results.append(match) - graph.consumed_ids.update(match.matched_node_ids) - worklist.extend(input_node for input_node in match.external_input_nodes if input_node.layer.name not in graph.consumed_ids) - continue - - worklist.extend(parent for parent in node.parents if parent.layer.name not in graph.consumed_ids) - - simplified_graph = models.SimplifiedGraph.from_graph(graph) - simplified_graph.fused_results = fused_results - - for match in fused_results: - simplified_graph.name_map[match.node.layer.name] = match.node - for node_id in match.matched_node_ids: - simplified_graph.replacement_map[node_id] = match.node.layer.name - - for node_id, node in graph.name_map.items(): - if node_id not in graph.consumed_ids: - simplified_graph.name_map[node_id] = node - - simplified_graph.source_nodes = [ - simplified_graph.name_map[simplified_graph.replacement_map.get(node.layer.name, node.layer.name)] - for node in graph.source_nodes - if simplified_graph.replacement_map.get(node.layer.name, node.layer.name) in simplified_graph.name_map - ] - - return simplified_graph - - -# - - def simplify(messy_ir:CVModels.LayerMap) -> CVModels.LayerMap: - graph: models.Graph = models.Graph.from_layer_map(messy_ir) - simplified_graph = rev_top_sort(graph) - return layer_map_from_simplified_graph(simplified_graph, messy_ir) - - - - -# def resolve_name(name: str, replacement_map: dict[str, str]) -> str: -# seen = set() - -# while name in replacement_map and name not in seen: -# seen.add(name) -# name = replacement_map[name] - -# return name - - -# def rewrite_node_refs(value, replacement_map: dict[str, str]): -# if isinstance(value, dict): -# rewritten = {key: rewrite_node_refs(item, replacement_map) for key, item in value.items()} -# if "node" in rewritten: -# rewritten["node"] = resolve_name(rewritten["node"], replacement_map) -# return rewritten - -# if isinstance(value, list): -# return [rewrite_node_refs(item, replacement_map) for item in value] - -# return value - - -# def extract_node_refs(value) -> list[str]: -# if isinstance(value, dict): -# refs = [] -# if "node" in value: -# refs.append(value["node"]) - -# for item in value.values(): -# refs.extend(extract_node_refs(item)) - -# return refs - -# if isinstance(value, list): -# refs = [] -# for item in value: -# refs.extend(extract_node_refs(item)) -# return refs - -# return [] - - -# def topo_sort(graph: models.SimplifiedGraph, args_map: dict[str, object], kwargs_map: dict[str, object]) -> list[str]: -# deps = {} -# users = {node_id: set() for node_id in graph.name_map} - -# for node_id in graph.name_map: -# node_refs = extract_node_refs(args_map[node_id]) + extract_node_refs(kwargs_map[node_id]) -# deps[node_id] = { -# ref -# for ref in node_refs -# if ref in graph.name_map and ref != node_id -# } - -# for dep in deps[node_id]: -# users[dep].add(node_id) - -# ready = sorted( -# [node_id for node_id, node_deps in deps.items() if not node_deps], -# key=lambda node_id: graph.name_map[node_id].layer.index, -# ) -# ordered = [] - -# while ready: -# node_id = ready.pop(0) -# ordered.append(node_id) - -# for user_id in sorted(users[node_id], key=lambda item: graph.name_map[item].layer.index): -# deps[user_id].remove(node_id) -# if not deps[user_id]: -# ready.append(user_id) - -# ready.sort(key=lambda item: graph.name_map[item].layer.index) - -# if len(ordered) != len(graph.name_map): -# raise ValueError("Simplified graph contains a cycle or unresolved dependency") - -# return ordered - - -# def layer_map_from_simplified_graph(graph: models.SimplifiedGraph, original: CVModels.LayerMap) -> CVModels.LayerMap: -# fused_node_ids = {result.node.layer.name for result in graph.fused_results} -# args_map = {} -# kwargs_map = {} - -# for node_id, node in graph.name_map.items(): -# args_map[node_id] = rewrite_node_refs(node.layer.args, graph.replacement_map) - -# if node_id in fused_node_ids: -# kwargs_map[node_id] = rewrite_node_refs(node.normalized_attrs, graph.replacement_map) -# else: -# kwargs_map[node_id] = rewrite_node_refs(node.layer.kwargs, graph.replacement_map) - -# ordered_node_ids = topo_sort(graph, args_map, kwargs_map) -# users_map = {node_id: [] for node_id in graph.name_map} - -# for node_id in ordered_node_ids: -# node_refs = extract_node_refs(args_map[node_id]) + extract_node_refs(kwargs_map[node_id]) -# for ref in node_refs: -# if ref in users_map and node_id not in users_map[ref]: -# users_map[ref].append(node_id) - -# nodes = [] -# for index, node_id in enumerate(ordered_node_ids): -# node = graph.name_map[node_id] -# nodes.append( -# CVModels.LayerRecord( -# index=index, -# name=node.layer.name, -# node_type=node.layer.node_type, -# target=node.layer.target, -# args=args_map[node_id], -# kwargs=kwargs_map[node_id], -# users=users_map[node_id], -# tensor_output_meta=node.layer.tensor_output_meta, -# module_stack=node.layer.module_stack, -# ) -# ) - - -# return CVModels.LayerMap( -# model_name=original.model_name, -# task=original.task, -# graph_signature=original.graph_signature, -# range_constants=original.range_constants, -# nodes=nodes, -# ) - + pass diff --git a/python/cactus/transpiler/main.py b/python/cactus/transpiler/main.py index 6e2408627..e12744367 100644 --- a/python/cactus/transpiler/main.py +++ b/python/cactus/transpiler/main.py @@ -1,6 +1,7 @@ import converter.convert as cv import converter.models as CVModels +#TODO: Update converter to also downlaod model config.txt if __name__ == "__main__": - model_map:CVModels = cv.model_ops_to_json("google/gemma-4-E2B", "/Users/sandhup/Documents/personal/cactus/python/cactus/transpiler/converter/jsons/") + model_map:CVModels = cv.model_ops_to_json("LiquidAI/LFM2.5-8B-A1B", "/Users/sandhup/Documents/personal/cactus/python/cactus/transpiler/converter/jsons/") \ No newline at end of file From edd39a596630acc59935ab9be487d34460674822 Mon Sep 17 00:00:00 2001 From: Parkirat Sandhu Date: Thu, 9 Jul 2026 21:26:35 -0400 Subject: [PATCH 17/84] Changes to transpiler converter --- python/cactus/transpiler/converter/models.py | 20 +- python/cactus/transpiler/converter/temp.py | 395 ------------------- 2 files changed, 14 insertions(+), 401 deletions(-) delete mode 100644 python/cactus/transpiler/converter/temp.py diff --git a/python/cactus/transpiler/converter/models.py b/python/cactus/transpiler/converter/models.py index e7f7f2dca..57ae14691 100644 --- a/python/cactus/transpiler/converter/models.py +++ b/python/cactus/transpiler/converter/models.py @@ -1,19 +1,31 @@ from pydantic import BaseModel from typing import Any, Optional +from dataclasses import dataclass import torch +model_presets: dict[str, str] = { + "PREFILL_NO_CACHE": "prefill_no_cache", + "PREFILL_WITH_CACHE": "prefill_with_cache", + "DECODING_WITH_CACHE": "decoding_with_cache", +} + +@dataclass(slots=True) +class ExportModel: + name: str + + input_modalities: list[str] + output_modalities: list[str] class InputtedModel(BaseModel): name: str input_modalities: list[str] -# output_modalities: list[str] -> Possible future need for this, but not yet + output_modalities: list[str] class LayerRecord(BaseModel): index: int name: str node_type: str target: str - #model_schema: str | None TODO: Need to add this field back and include model_schema during conversion args: Any kwargs: Any users: list[str] @@ -107,10 +119,6 @@ def jsonable(x: Any) -> Any: def aten_name(target: Any) -> str: - """ - Convert default naming to simpler name. - For non-ATen targets, falls back to str(target). - """ schema = getattr(target, "_schema", None) if schema is not None: diff --git a/python/cactus/transpiler/converter/temp.py b/python/cactus/transpiler/converter/temp.py deleted file mode 100644 index 2d997cb96..000000000 --- a/python/cactus/transpiler/converter/temp.py +++ /dev/null @@ -1,395 +0,0 @@ -# import argparse -# import json -# from collections import Counter -# from typing import Any, Dict, Optional - -# import torch -# from transformers import ( -# AutoModel, -# AutoModelForCausalLM, -# AutoModelForSequenceClassification, -# AutoTokenizer, -# ) - - -# def jsonable(x: Any) -> Any: -# """ -# Convert graph/export objects into JSON-safe values. -# """ -# if isinstance(x, torch.fx.Node): -# return {"node": x.name} - -# if isinstance(x, torch.Tensor): -# return { -# "shape": list(x.shape), -# "dtype": str(x.dtype), -# "device": str(x.device), -# "requires_grad": bool(x.requires_grad), -# } - -# if isinstance(x, torch.Size): -# return list(x) - -# if isinstance(x, torch.dtype): -# return str(x) - -# if isinstance(x, torch.device): -# return str(x) - -# if isinstance(x, slice): -# return { -# "slice": { -# "start": jsonable(x.start), -# "stop": jsonable(x.stop), -# "step": jsonable(x.step), -# } -# } - -# if isinstance(x, range): -# return list(x) - -# if isinstance(x, (list, tuple)): -# return [jsonable(v) for v in x] - -# if isinstance(x, dict): -# return {str(k): jsonable(v) for k, v in x.items()} - -# if isinstance(x, (str, int, float, bool)) or x is None: -# return x - -# return repr(x) - - -# def aten_name(target: Any) -> str: -# """ -# Convert torch.ops.aten.linear.default into: -# aten.linear.default - -# For non-ATen targets, falls back to str(target). -# """ -# schema = getattr(target, "_schema", None) - -# if schema is not None: -# # Example: -# # schema.name = "aten::linear" -# # schema.overload_name = "" or "Tensor" or "int" -# if "::" in schema.name: -# namespace, op = schema.name.split("::", 1) -# overload = schema.overload_name if schema.overload_name else "default" -# return f"{namespace}.{op}.{overload}" - -# if hasattr(target, "name"): -# try: -# return target.name() -# except Exception: -# pass - -# return str(target) - - -# def schema_string(target: Any) -> Optional[str]: -# schema = getattr(target, "_schema", None) -# return str(schema) if schema is not None else None - - -# def extract_tensor_meta(node: torch.fx.Node) -> Optional[Any]: -# """ -# torch.export usually stores fake tensor output metadata in node.meta["val"]. -# """ -# if "val" not in node.meta: -# return None - -# return jsonable(node.meta["val"]) - - -# def extract_module_stack(node: torch.fx.Node) -> Optional[Any]: -# """ -# Sometimes torch.export nodes contain original module context here. -# This can help later when mapping ops back to model.layers.X.self_attn, mlp, etc. -# """ -# stack = node.meta.get("nn_module_stack", None) - -# if stack is None: -# return None - -# out = [] - -# for key, value in stack.items(): -# if isinstance(value, tuple) and len(value) >= 2: -# module_path = value[0] -# module_type = value[1] -# out.append( -# { -# "key": str(key), -# "module_path": str(module_path), -# "module_type": getattr(module_type, "__name__", str(module_type)), -# } -# ) -# else: -# out.append( -# { -# "key": str(key), -# "value": repr(value), -# } -# ) - -# return out - - -# def make_graph_json( -# exported: torch.export.ExportedProgram, -# model_name: str, -# task: str, -# attn_implementation: str, -# ran_decompositions: bool, -# ) -> Dict[str, Any]: -# graph_module = exported.graph_module -# nodes_json = [] - -# for i, node in enumerate(graph_module.graph.nodes): -# target = aten_name(node.target) - -# node_json = { -# "index": i, -# "name": node.name, -# "node_type": node.op, -# "target": target, -# "schema": schema_string(node.target), -# "args": jsonable(node.args), -# "kwargs": jsonable(node.kwargs), -# "users": [user.name for user in node.users], -# } - -# tensor_meta = extract_tensor_meta(node) -# if tensor_meta is not None: -# node_json["tensor_meta"] = tensor_meta - -# module_stack = extract_module_stack(node) -# if module_stack is not None: -# node_json["module_stack"] = module_stack - -# nodes_json.append(node_json) - -# op_counts = Counter( -# node["target"] -# for node in nodes_json -# if node["node_type"] == "call_function" -# ) - -# return { -# "model_name": model_name, -# "task": task, -# "attn_implementation": attn_implementation, -# "ran_decompositions": ran_decompositions, -# "num_nodes": len(nodes_json), -# "num_call_function_nodes": sum(op_counts.values()), -# "op_counts": dict(op_counts.most_common()), -# "graph_signature": repr(exported.graph_signature), -# "range_constraints": repr(exported.range_constraints), -# "nodes": nodes_json, -# } - - -# def load_model(model_name: str, task: str, attn_implementation: str): -# """ -# Load HF model with eager attention when supported. -# """ -# model_cls_map = { -# "base": AutoModel, -# "causal-lm": AutoModelForCausalLM, -# "sequence-classification": AutoModelForSequenceClassification, -# } - -# if task not in model_cls_map: -# raise ValueError(f"Unsupported task: {task}") - -# model_cls = model_cls_map[task] - -# try: -# model = model_cls.from_pretrained( -# model_name, -# attn_implementation=attn_implementation, -# ) -# except TypeError: -# print( -# f"Warning: this model class did not accept " -# f'attn_implementation="{attn_implementation}". ' -# f"Retrying without it." -# ) -# model = model_cls.from_pretrained(model_name) - -# model.eval() - -# # Hugging Face outputs are often ModelOutput dataclasses. -# # Tuples are easier for torch.export to serialize. -# if hasattr(model, "config"): -# model.config.return_dict = False - -# # Important for GPT/LLaMA/Mistral/etc. -# # Disables KV-cache outputs and cache-specific forward paths. -# if hasattr(model.config, "use_cache"): -# model.config.use_cache = False - -# return model - - -# def export_hf_model_to_json( -# model_name: str, -# output_path: str, -# task: str = "causal-lm", -# sample_text: str = "Hello world", -# max_length: int = 32, -# attn_implementation: str = "eager", -# decompose: bool = True, -# strict: bool = False, -# ): -# tokenizer = AutoTokenizer.from_pretrained(model_name) - -# model = load_model( -# model_name=model_name, -# task=task, -# attn_implementation=attn_implementation, -# ) - -# inputs = tokenizer( -# sample_text, -# return_tensors="pt", -# truncation=True, -# max_length=max_length, -# ) - -# # Some decoder-only tokenizers do not define pad_token. -# # Not needed for a single example, but useful if you later batch examples. -# if tokenizer.pad_token is None and tokenizer.eos_token is not None: -# tokenizer.pad_token = tokenizer.eos_token - -# with torch.no_grad(): -# exported = torch.export.export( -# model, -# args=(), -# kwargs=dict(inputs), -# strict=strict, -# ) - -# ran_decompositions = False - -# if decompose: -# try: -# exported = exported.run_decompositions() -# ran_decompositions = True -# except Exception as e: -# print(f"Warning: run_decompositions() failed: {repr(e)}") -# print("Continuing with the original exported graph.") - -# graph_json = make_graph_json( -# exported=exported, -# model_name=model_name, -# task=task, -# attn_implementation=attn_implementation, -# ran_decompositions=ran_decompositions, -# ) - -# with open(output_path, "w") as f: -# json.dump(graph_json, f, indent=2) - -# print(f"Saved graph JSON to: {output_path}") -# print(f"Total nodes: {graph_json['num_nodes']}") -# print(f"call_function nodes: {graph_json['num_call_function_nodes']}") -# print() -# print("Top ops:") -# for op, count in list(graph_json["op_counts"].items())[:30]: -# print(f"{count:5d} {op}") - -# attention_ops = [ -# op for op in graph_json["op_counts"] -# if "attention" in op.lower() or "scaled_dot_product" in op.lower() -# ] - -# if attention_ops: -# print() -# print("Attention-like ops still present:") -# for op in attention_ops: -# print(f" {op}") -# print() -# print( -# "If you still see aten.scaled_dot_product_attention.default, " -# "attention was not fully broken into matmul/softmax/matmul." -# ) - - -# # def parse_args(): -# # parser = argparse.ArgumentParser( -# # description="Export a Hugging Face PyTorch model to lower-level ATen JSON." -# # ) - -# # parser.add_argument( -# # "--model", -# # type=str, -# # required=True, -# # help="HF model name, e.g. bert-base-uncased or gpt2.", -# # ) - -# # parser.add_argument( -# # "--output", -# # type=str, -# # default="model_aten_graph.json", -# # help="Output JSON path.", -# # ) - -# # parser.add_argument( -# # "--task", -# # type=str, -# # default="causal-lm", -# # choices=["base", "causal-lm", "sequence-classification"], -# # help="Which AutoModel class to use.", -# # ) - -# # parser.add_argument( -# # "--text", -# # type=str, -# # default="Hello world", -# # help="Sample text used to trace/export the model.", -# # ) - -# # parser.add_argument( -# # "--max-length", -# # type=int, -# # default=32, -# # help="Max token length for the sample input.", -# # ) - -# # parser.add_argument( -# # "--attn-implementation", -# # type=str, -# # default="eager", -# # help='Attention backend. For graph inspection, use "eager".', -# # ) - -# # parser.add_argument( -# # "--no-decompose", -# # action="store_true", -# # help="Disable exported.run_decompositions().", -# # ) - -# # parser.add_argument( -# # "--strict", -# # action="store_true", -# # help="Use strict=True in torch.export.export.", -# # ) - -# # return parser.parse_args() - - -# if __name__ == "__main__": -# args = parse_args() - -# export_hf_model_to_json( -# model_name=args.model, -# output_path=args.output, -# task=args.task, -# sample_text=args.text, -# max_length=args.max_length, -# attn_implementation=args.attn_implementation, -# decompose=not args.no_decompose, -# strict=args.strict, -# ) \ No newline at end of file From a4c7af667ea3daa884508b06c7b2ad0d469fb15d Mon Sep 17 00:00:00 2001 From: Parkirat Date: Fri, 10 Jul 2026 11:44:36 -0400 Subject: [PATCH 18/84] Reset --- .../cactus/transpiler/converter/constants.py | 2 - python/cactus/transpiler/converter/convert.py | 33 - .../converter/model_preprocessing.py | 13 - .../transpiler/converter/model_utils.py | 5 - python/cactus/transpiler/converter/models.py | 182 ---- python/cactus/transpiler/ir/constants.py | 74 -- .../transpiler/ir/fusion_graph_catalog.py | 530 ----------- python/cactus/transpiler/ir/fusions_utils.py | 436 --------- python/cactus/transpiler/ir/models.py | 834 ------------------ python/cactus/transpiler/ir/simplify_ir.py | 15 - python/cactus/transpiler/main.py | 7 - 11 files changed, 2131 deletions(-) delete mode 100644 python/cactus/transpiler/converter/constants.py delete mode 100644 python/cactus/transpiler/converter/convert.py delete mode 100644 python/cactus/transpiler/converter/model_preprocessing.py delete mode 100644 python/cactus/transpiler/converter/model_utils.py delete mode 100644 python/cactus/transpiler/converter/models.py delete mode 100644 python/cactus/transpiler/ir/constants.py delete mode 100644 python/cactus/transpiler/ir/fusion_graph_catalog.py delete mode 100644 python/cactus/transpiler/ir/fusions_utils.py delete mode 100644 python/cactus/transpiler/ir/models.py delete mode 100644 python/cactus/transpiler/ir/simplify_ir.py diff --git a/python/cactus/transpiler/converter/constants.py b/python/cactus/transpiler/converter/constants.py deleted file mode 100644 index 9060d96d5..000000000 --- a/python/cactus/transpiler/converter/constants.py +++ /dev/null @@ -1,2 +0,0 @@ -DEFAULT_PROMPT = "What is the capital of France?" -MODEL_TASK = "Causal_LLM" \ No newline at end of file diff --git a/python/cactus/transpiler/converter/convert.py b/python/cactus/transpiler/converter/convert.py deleted file mode 100644 index 08e4bd77c..000000000 --- a/python/cactus/transpiler/converter/convert.py +++ /dev/null @@ -1,33 +0,0 @@ -import torch -from .models import LayerRecord, LayerMap -from .model_preprocessing import load_model, load_input -from . import constants - - - -def model_ops_to_json(id:str, output_path:str) -> LayerMap: - model = load_model(id) - input = load_input(id) - - model.config.return_dict = False if hasattr(model.config, "return_dict") else None - model.config.use_cache = False if hasattr(model.config, "use_cache") else None - - with torch.no_grad(): - exported_model = torch.export.export(model, args=input, kwargs={"use_cache":False}, strict=False).run_decompositions() - - #Iterates across every node in the graph and stores it as a layer object - nodes = [] - for i, record in enumerate(exported_model.graph_module.graph.nodes): - nodes.append(LayerRecord.from_node(num = i, x = record)) - - json_ = LayerMap.from_data(x=exported_model, name=id, model_task = constants.MODEL_TASK, nodes_list=nodes).model_dump_json(indent=4) - - with open(output_path+"output1.json", "w", encoding="utf-8") as file: - file.write(json_) - - - -if __name__ == "__main__": - model_ops_to_json("LiquidAI/LFM2.5-8B-A1B", "/Users/sandhup/Documents/personal/cactus/python/cactus/transpiler/converter/jsons/") - print("Done") - diff --git a/python/cactus/transpiler/converter/model_preprocessing.py b/python/cactus/transpiler/converter/model_preprocessing.py deleted file mode 100644 index 4d66129f9..000000000 --- a/python/cactus/transpiler/converter/model_preprocessing.py +++ /dev/null @@ -1,13 +0,0 @@ -from . import constants -from transformers import AutoTokenizer, AutoModelForCausalLM - -def load_model(id:str) -> object: - model = AutoModelForCausalLM.from_pretrained(id, attn_implementation = "eager") - model.eval() - return model - -#Currently just handles single input type (strings), but will change to make input a list to allow mutli-modal models to run multiple different inputs -def load_input(id:str) -> tuple: - tokenizer = AutoTokenizer.from_pretrained(id) - input = (tokenizer(constants.DEFAULT_PROMPT, return_tensors="pt")["input_ids"],) - return input diff --git a/python/cactus/transpiler/converter/model_utils.py b/python/cactus/transpiler/converter/model_utils.py deleted file mode 100644 index 2c5824019..000000000 --- a/python/cactus/transpiler/converter/model_utils.py +++ /dev/null @@ -1,5 +0,0 @@ -from typing import Any, Optional -import torch -# from models import TensorInstance, Slice - - diff --git a/python/cactus/transpiler/converter/models.py b/python/cactus/transpiler/converter/models.py deleted file mode 100644 index 57ae14691..000000000 --- a/python/cactus/transpiler/converter/models.py +++ /dev/null @@ -1,182 +0,0 @@ -from pydantic import BaseModel -from typing import Any, Optional -from dataclasses import dataclass -import torch - -model_presets: dict[str, str] = { - "PREFILL_NO_CACHE": "prefill_no_cache", - "PREFILL_WITH_CACHE": "prefill_with_cache", - "DECODING_WITH_CACHE": "decoding_with_cache", -} - -@dataclass(slots=True) -class ExportModel: - name: str - - input_modalities: list[str] - output_modalities: list[str] - -class InputtedModel(BaseModel): - name: str - input_modalities: list[str] - output_modalities: list[str] - -class LayerRecord(BaseModel): - index: int - name: str - node_type: str - target: str - args: Any - kwargs: Any - users: list[str] - tensor_output_meta: Any | None - module_stack: Any | None - - @classmethod - def from_node(cls, num:int, x: torch.Node) -> "LayerRecord": - return cls( - index = num, - name = str(x.name), - node_type = str(x.op), - target = str(x.target), - args = jsonable(x.args), - kwargs = jsonable(x.kwargs), - users = [user.name for user in x.users], - tensor_output_meta = extract_tensor_meta(x), - module_stack = extract_module_stack(x) - ) - -class TensorInstance(BaseModel): - shape: list - dtype: str - device: str - requires_grad: bool - - @classmethod - def from_tensor(cls, x: torch.Tensor) -> "TensorInstance": - return cls(shape=list(x.shape), dtype=str(x.dtype), device=str(x.device), requires_grad = bool(x.requires_grad)) - -class Slice(BaseModel): - start: Any - stop: Any - step: Any - - @classmethod - def from_slice(cls, x: slice) -> "Slice": - return cls(start = jsonable(x.start), stop = jsonable(x.stop), step = jsonable(x.step)) - - -class LayerMap(BaseModel): - model_name: str - task: str - graph_signature: str - range_constants: str - nodes: list[LayerRecord] - - @classmethod - def from_data(cls, x:torch.export.ExportedProgram, name:str, model_task:str, nodes_list:list[LayerRecord]) -> "LayerMap": - return cls(model_name = name, task = model_task, graph_signature = repr(x.graph_signature), range_constants = repr(x.range_constraints), nodes = nodes_list) - - - -"""###################################### MODEL UTILS!!!!!!! ######################################""" - -def jsonable(x: Any) -> Any: - """ - Convert graph/export objects into JSON-safe values. - """ - if isinstance(x, torch.fx.Node): - return {"node": x.name} - - if isinstance(x, torch.Tensor): - return TensorInstance.from_tensor(x) - - if isinstance(x, torch.Size): - return list(x) - - if isinstance(x, torch.dtype): - return str(x) - - if isinstance(x, torch.device): - return str(x) - - if isinstance(x, slice): - return Slice.from_slice(x) - - if isinstance(x, range): - return list(x) - - if isinstance(x, (list, tuple)): - return [jsonable(v) for v in x] - - if isinstance(x, dict): - return {str(k): jsonable(v) for k, v in x.items()} - - if isinstance(x, (str, int, float, bool)) or x is None: - return x - - return repr(x) - - -def aten_name(target: Any) -> str: - schema = getattr(target, "_schema", None) - - if schema is not None: - if "::" in schema.name: - namespace, op = schema.name.split("::", 1) - overload = schema.overload_name if schema.overload_name else "default" - return f"{namespace}.{op}.{overload}" - - if hasattr(target, "name"): - try: - return target.name() - except Exception: - pass - - return str(target) - - -def extract_tensor_meta(node: torch.fx.Node) -> Optional[Any]: - """ - torch.export usually stores fake tensor output metadata in node.meta["val"]. - """ - if "val" not in node.meta: - return None - - return jsonable(node.meta["val"]) - - -#This function needs to be further optimized and cleaned up -#CLEANUP: Create object for output dictionary (potentially one outter object as well); Condense down for loop -def extract_module_stack(node: torch.fx.Node) -> Optional[Any]: - """ - Sometimes torch.export nodes contain original module context here. - This can help later when mapping ops back to model.layers.X.self_attn, mlp, etc. - """ - stack = node.meta.get("nn_module_stack", None) - - if stack is None: - return None - - out = [] - - for key, value in stack.items(): - if isinstance(value, tuple) and len(value) >= 2: - module_path = value[0] - module_type = value[1] - out.append( - { - "key": str(key), - "module_path": str(module_path), - "module_type": getattr(module_type, "__name__", str(module_type)), - } - ) - else: - out.append( - { - "key": str(key), - "value": repr(value), - } - ) - - return out \ No newline at end of file diff --git a/python/cactus/transpiler/ir/constants.py b/python/cactus/transpiler/ir/constants.py deleted file mode 100644 index 4ecf94e6f..000000000 --- a/python/cactus/transpiler/ir/constants.py +++ /dev/null @@ -1,74 +0,0 @@ -LAYER_ATTRS_MAP = { - "aten._assert_tensor_metadata.default": ["size", "stride", "dtype"], - "aten._softmax.default": ["dim", "half_to_float"], - "aten._to_copy.default": [None, "dtype", "layout", "device", "pin_memory", "non_blocking", "memory_format"], - "aten.add.Scalar": ["scalar"], - "aten.add.Tensor": ["scalar"], - "aten.addmm.default": [None, None, None, "beta", "alpha"], - "aten.arange.start_step": ["start", "end", "step"], - "aten.avg_pool1d.default": [None, "kernel_size", "stride", "padding", "ceil_mode", "count_include_pad"], - "aten.avg_pool2d.default": [None, "kernel_size", "stride", "padding", "ceil_mode", "count_include_pad", "divisor_override"], - "aten.batch_norm.default": [None, None, None, None, None, "training", "momentum", "eps", "cudnn_enabled"], - "aten.bitwise_and.Tensor": [], - "aten.bitwise_or.Tensor": [], - "aten.bmm.default": [], - "aten.cat.default": ["dim"], - "aten.clone.default": [], - "aten.clamp.default": [None, "min", "max"], - "aten.clamp_max.default": [None, "max"], - "aten.clamp_min.default": [None, "min"], - "aten.convolution.default": [None, None, None, "stride", "padding", "dilation", "transposed", "output_padding", "groups"], - "aten.conv1d.default": [None, None, None, "stride", "padding", "dilation", "groups"], - "aten.conv2d.default": [None, None, None, "stride", "padding", "dilation", "groups"], - "aten.cos.default": [], - "aten.cumsum.default": ["dim"], - "aten.div.Scalar": ["scalar"], - "aten.div.Tensor": ["scalar"], - "aten.embedding.default": ["padding_idx"], - "aten.eq.Scalar": ["scalar"], - "aten.eq.Tensor": [], - "aten.expand.default": ["shape"], - "aten.full.default": ["shape", "fill_value"], - "aten.gelu.default": [], - "aten.glu.default": ["dim"], - "aten.group_norm.default": [None, "num_groups", None, None, "eps", "cudnn_enabled"], - "aten.gt.Tensor": [], - "aten.index.Tensor": [], - "aten.index_put.default": [], - "aten.layer_norm.default": [None, "normalized_shape", None, None, "eps", "cudnn_enabled"], - "aten.le.Tensor": [], - "aten.linear.default": [], - "aten.max.dim": ["dim", "keepdim"], - "aten.max_pool1d.default": [None, "kernel_size", "stride", "padding", "dilation", "ceil_mode"], - "aten.max_pool2d.default": [None, "kernel_size", "stride", "padding", "dilation", "ceil_mode"], - "aten.mean.dim": ["dim", "keepdim"], - "aten.mm.default": [], - "aten.mul.Scalar": ["scalar"], - "aten.mul.Tensor": ["scalar"], - "aten.native_batch_norm.default": [None, None, None, None, None, "training", "momentum", "eps"], - "aten.native_layer_norm.default": [None, "normalized_shape", None, None, "eps"], - "aten.ne.Scalar": ["scalar"], - "aten.ne.Tensor": [], - "aten.neg.default": [], - "aten.permute.default": ["dims"], - "aten.pow.Tensor_Scalar": ["exponent"], - "aten.rms_norm.default": [None, "normalized_shape", None, "eps"], - "aten.scalar_tensor.default": ["value"], - "aten.scaled_dot_product_attention.default": [None, None, None, "attn_mask", "dropout_p", "is_causal", "scale", "enable_gqa"], - "aten.select.int": ["dim", "index"], - "aten.sin.default": [], - "aten.slice.Tensor": ["dim", "start", "end", "step"], - "aten.sub.Scalar": ["scalar"], - "aten.sub.Tensor": ["scalar"], - "aten.sum.dim_IntList": ["dim", "keepdim"], - "aten.tanh.default": [], - "aten.topk.default": [None, "k", "dim", "largest", "sorted"], - "aten.transpose.int": ["dim0", "dim1"], - "aten.unsqueeze.default": ["dim"], - "aten.var.dim": ["dim", "unbiased", "keepdim"], - "aten.view.default": ["shape"], - "aten.where.self": [], -} - -#Dictates maximum DFS depth during fusion search -DFS_DEPTH = 10 diff --git a/python/cactus/transpiler/ir/fusion_graph_catalog.py b/python/cactus/transpiler/ir/fusion_graph_catalog.py deleted file mode 100644 index 4bc13a0f5..000000000 --- a/python/cactus/transpiler/ir/fusion_graph_catalog.py +++ /dev/null @@ -1,530 +0,0 @@ -from __future__ import annotations - -from dataclasses import dataclass, field -from typing import Any - - -@dataclass(frozen=True, slots=True) -class FusionNode: - name: str - op: str | tuple[str, ...] - attrs: dict[str, Any] = field(default_factory=dict) - optional: bool = False - transparent: bool = False - repeated: bool = False - notes: str = "" - - -@dataclass(frozen=True, slots=True) -class FusionEdge: - source: str - dest: str - dest_input_index: int - source_output_index: int | None = None - optional: bool = False - repeated: bool = False - - -@dataclass(frozen=True, slots=True) -class FusionInput: - role: str - node: str - parent_index: int | None = None - variadic: bool = False - optional: bool = False - repeated_group: str | None = None - - -@dataclass(frozen=True, slots=True) -class AttrCapture: - name: str - source_node: str | None = None - source_attr: str | None = None - default: Any = None - required: bool = True - notes: str = "" - - -@dataclass(frozen=True, slots=True) -class FusionGraph: - name: str - target: str - cactus_inputs: tuple[str, ...] - cactus_attrs: tuple[str, ...] - nodes: dict[str, FusionNode] = field(default_factory=dict) - edges: tuple[FusionEdge, ...] = () - root: str | None = None - outputs: tuple[str, ...] = () - inputs: tuple[FusionInput, ...] = () - output_attrs: tuple[AttrCapture, ...] = () - constraints: tuple[str, ...] = () - variants: tuple[str, ...] = () - repeated_groups: tuple[str, ...] = () - notes: str = "" - - -def n( - name: str, - op: str | tuple[str, ...], - *, - attrs: dict[str, Any] | None = None, - optional: bool = False, - transparent: bool = False, - repeated: bool = False, - notes: str = "", -) -> FusionNode: - return FusionNode( - name=name, - op=op, - attrs=attrs or {}, - optional=optional, - transparent=transparent, - repeated=repeated, - notes=notes, - ) - - -def e( - source: str, - dest: str, - dest_input_index: int, - *, - source_output_index: int | None = None, - optional: bool = False, - repeated: bool = False, -) -> FusionEdge: - return FusionEdge( - source=source, - dest=dest, - dest_input_index=dest_input_index, - source_output_index=source_output_index, - optional=optional, - repeated=repeated, - ) - - -def inp( - role: str, - node: str, - parent_index: int | None = None, - *, - variadic: bool = False, - optional: bool = False, - repeated_group: str | None = None, -) -> FusionInput: - return FusionInput( - role=role, - node=node, - parent_index=parent_index, - variadic=variadic, - optional=optional, - repeated_group=repeated_group, - ) - - -def cap( - name: str, - source_node: str | None = None, - source_attr: str | None = None, - *, - default: Any = None, - required: bool = True, - notes: str = "", -) -> AttrCapture: - return AttrCapture( - name=name, - source_node=source_node, - source_attr=source_attr, - default=default, - required=required, - notes=notes, - ) - - -def direct( - target: str, - cactus_inputs: tuple[str, ...], - cactus_attrs: tuple[str, ...], - raw_ops: str | tuple[str, ...], - *, - attr_map: dict[str, str] | None = None, - optional_attr_defaults: dict[str, Any] | None = None, - notes: str = "", -) -> FusionGraph: - root = f"{target}_root" - output_attrs = [] - - for cactus_attr in cactus_attrs: - source_attr = (attr_map or {}).get(cactus_attr, cactus_attr) - output_attrs.append( - cap( - cactus_attr, - root, - source_attr, - default=(optional_attr_defaults or {}).get(cactus_attr), - required=cactus_attr not in (optional_attr_defaults or {}), - ) - ) - - return FusionGraph( - name=f"{target}_direct", - target=target, - cactus_inputs=cactus_inputs, - cactus_attrs=cactus_attrs, - nodes={root: n(root, raw_ops)}, - root=root, - outputs=(root,), - inputs=tuple(inp(role, root, i) for i, role in enumerate(cactus_inputs)), - output_attrs=tuple(output_attrs), - notes=notes, - ) - - -def source_op(target: str, cactus_attrs: tuple[str, ...], notes: str) -> FusionGraph: - root = f"{target}_source" - return FusionGraph( - name=f"{target}_source", - target=target, - cactus_inputs=(), - cactus_attrs=cactus_attrs, - nodes={root: n(root, "cactus.source_op", notes=notes)}, - root=root, - outputs=(root,), - output_attrs=tuple(cap(attr, default=None, notes="Caller/runtime supplied attr") for attr in cactus_attrs), - notes=notes, - ) - - -FUSION_GRAPHS: dict[str, FusionGraph] = { - "add": direct("add", ("a", "b"), (), "aten.add.Tensor"), - "add_clipped": direct( - "add_clipped", - ("a", "b"), - (), - "aten.add.Tensor", - notes="Needs post-add clipping semantics from a trace before this can be safely fused.", - ), - "subtract": direct("subtract", ("a", "b"), (), "aten.sub.Tensor"), - "multiply": direct("multiply", ("a", "b"), (), "aten.mul.Tensor"), - "divide": direct("divide", ("a", "b"), (), "aten.div.Tensor"), - "not_equal": direct("not_equal", ("a", "b"), (), "aten.ne.Tensor"), - "abs": direct("abs", ("x",), (), "aten.abs.default"), - "pow": direct("pow", ("x",), ("exponent",), "aten.pow.Tensor_Scalar"), - "precision_cast": direct( - "precision_cast", - ("x",), - ("dtype",), - "aten._to_copy.default", - attr_map={"dtype": "dtype"}, - ), - "quantize_activations": direct( - "quantize_activations", - ("x",), - (), - ("aten.quantize_per_tensor.default", "aten._to_copy.default"), - notes="Needs quantization-specific attrs/scales once a quantized export is available.", - ), - "scalar_add": direct("scalar_add", ("x",), ("value",), "aten.add.Scalar", attr_map={"value": "scalar"}), - "scalar_subtract": direct("scalar_subtract", ("x",), ("value",), "aten.sub.Scalar", attr_map={"value": "scalar"}), - "scalar_multiply": direct("scalar_multiply", ("x",), ("value",), ("aten.mul.Scalar", "aten.neg.default"), attr_map={"value": "scalar"}), - "scalar_divide": direct("scalar_divide", ("x",), ("value",), "aten.div.Scalar", attr_map={"value": "scalar"}), - "scalar_floor_divide": direct("scalar_floor_divide", ("x",), ("value",), "aten.floor_divide.default", attr_map={"value": "scalar"}), - "scalar_not_equal": direct("scalar_not_equal", ("x",), ("value",), "aten.ne.Scalar", attr_map={"value": "scalar"}), - "scalar_exp": direct("scalar_exp", ("x",), (), "aten.exp.default"), - "scalar_sqrt": direct("scalar_sqrt", ("x",), (), "aten.sqrt.default"), - "scalar_cos": direct("scalar_cos", ("x",), (), "aten.cos.default"), - "scalar_sin": direct("scalar_sin", ("x",), (), "aten.sin.default"), - "scalar_log": direct("scalar_log", ("x",), (), "aten.log.default"), - "clamp": direct("clamp", ("x",), ("lo", "hi"), "aten.clamp.default", attr_map={"lo": "min", "hi": "max"}), - "masked_select_prefix": direct( - "masked_select_prefix", - ("x", "mask"), - (), - ("aten.masked_select.default", "aten.index.Tensor"), - notes="Needs prefix semantics check, not just generic masked_select.", - ), - "masked_scatter": direct("masked_scatter", ("x", "mask", "source"), (), "aten.masked_scatter.default"), - "view": direct("view", ("x",), ("shape",), "aten.view.default"), - "reshape": direct("reshape", ("x",), ("shape",), "aten.reshape.default"), - "expand": direct("expand", ("x",), ("shape",), "aten.expand.default"), - "flatten": direct( - "flatten", - ("x",), - ("start_dim", "end_dim"), - "aten.flatten.using_ints", - optional_attr_defaults={"start_dim": 0, "end_dim": -1}, - ), - "slice": direct("slice", ("x",), ("axis", "start", "end", "step"), "aten.slice.Tensor", attr_map={"axis": "dim"}), - "index": direct("index", ("x",), ("index_value", "axis"), "aten.select.int", attr_map={"axis": "dim", "index_value": "index"}), - "transpose": direct( - "transpose", - ("x",), - ("backend",), - "aten.transpose.int", - optional_attr_defaults={"backend": 0}, - notes="Needs dim0/dim1 capture if Cactus transpose grows beyond backend-only attrs.", - ), - "permute": direct( - "permute", - ("x",), - ("permutation", "backend"), - "aten.permute.default", - attr_map={"permutation": "dims"}, - optional_attr_defaults={"backend": 0}, - ), - "matmul": direct( - "matmul", - ("a", "b"), - ("pretransposed_rhs", "backend", "output_dtype"), - ("aten.mm.default", "aten.matmul.default", "aten.bmm.default"), - optional_attr_defaults={"pretransposed_rhs": False, "backend": 0, "output_dtype": None}, - ), - "gather": direct("gather", ("tensor", "indices"), (), "aten.gather.default"), - "embedding_from_tensor": direct("embedding_from_tensor", ("embedding_tensor", "indices"), (), "aten.embedding.default"), - "embedding_from_file": source_op("embedding_from_file", ("filename",), "Runtime/source op. No ATen subgraph unless file loading is represented in IR."), - "mmap_embeddings": source_op("mmap_embeddings", ("filename",), "Runtime/source op for memory-mapped embedding weights."), - "mmap_weights": source_op("mmap_weights", ("filename",), "Runtime/source op for memory-mapped weights."), - "bilinear_interpolation": direct( - "bilinear_interpolation", - ("pos_embeds",), - ("dst_height", "dst_width"), - ("aten.upsample_bilinear2d.vec", "aten._upsample_bilinear2d_aa.vec"), - ), - "concat": direct("concat", ("a", "b"), ("axis",), "aten.cat.default", attr_map={"axis": "dim"}), - "cat": FusionGraph( - name="cat_variadic", - target="cat", - cactus_inputs=("tensors",), - cactus_attrs=("axis",), - nodes={"cat": n("cat", "aten.cat.default")}, - root="cat", - outputs=("cat",), - inputs=(inp("tensors", "cat", variadic=True),), - output_attrs=(cap("axis", "cat", "dim"),), - ), - "groupnorm": direct("groupnorm", ("x", "weight", "bias"), ("num_groups", "eps"), "aten.group_norm.default"), - "layernorm": direct("layernorm", ("x", "weight", "bias"), ("eps",), ("aten.layer_norm.default", "aten.native_layer_norm.default")), - "batchnorm": direct("batchnorm", ("x", "weight", "bias", "running_mean", "running_var"), ("axis", "eps"), ("aten.batch_norm.default", "aten.native_batch_norm.default"), optional_attr_defaults={"axis": 1}), - "rms_norm": FusionGraph( - name="rms_norm_decomposed_or_direct", - target="rms_norm", - cactus_inputs=("x", "weight"), - cactus_attrs=("eps",), - nodes={ - "squared_x": n("squared_x", "aten.pow.Tensor_Scalar", attrs={"exponent": 2}), - "variance": n("variance", "aten.mean.dim"), - "variance_with_eps": n("variance_with_eps", "aten.add.Tensor"), - "inverse_rms": n("inverse_rms", "aten.pow.Tensor_Scalar", attrs={"exponent": -0.5}), - "normalized_x": n("normalized_x", "aten.mul.Tensor"), - "weighted_output": n("weighted_output", "aten.mul.Tensor"), - }, - edges=( - e("squared_x", "variance", 0), - e("variance", "variance_with_eps", 0), - e("variance_with_eps", "inverse_rms", 0), - e("inverse_rms", "normalized_x", 1), - e("normalized_x", "weighted_output", 0), - ), - root="weighted_output", - outputs=("weighted_output",), - inputs=(inp("x", "normalized_x", 0), inp("weight", "weighted_output", 1)), - output_attrs=(cap("eps", "variance_with_eps", "scalar"),), - constraints=("normalized_x input 0 and squared_x input 0 must be the same node", "mean axis must be hidden dimension"), - variants=("aten.rms_norm.default", "pow_mean_add_pow_mul_mul"), - ), - "topk": direct("topk", ("x",), ("k",), "aten.topk.default"), - "rope": FusionGraph( - name="rope_decomposed", - target="rope", - cactus_inputs=("x",), - cactus_attrs=("theta", "position_offset", "backend"), - nodes={ - "x": n("x", "external"), - "sin": n("sin", ("aten.sin.default", "constant_sin_cache")), - "cos": n("cos", ("aten.cos.default", "constant_cos_cache")), - "rotate_half": n("rotate_half", "rope_rotate_half_subgraph"), - "scaled_x": n("scaled_x", "aten.mul.Tensor"), - "scaled_rot": n("scaled_rot", "aten.mul.Tensor"), - "out": n("out", "aten.add.Tensor"), - }, - edges=(e("x", "scaled_x", 0), e("cos", "scaled_x", 1), e("x", "rotate_half", 0), e("rotate_half", "scaled_rot", 0), e("sin", "scaled_rot", 1), e("scaled_x", "out", 0), e("scaled_rot", "out", 1)), - root="out", - outputs=("out",), - inputs=(inp("x", "scaled_x", 0),), - output_attrs=(cap("theta", default=None, required=False), cap("position_offset", default=0, required=False), cap("backend", default=0, required=False)), - constraints=("pairwise rotary dims must match Cactus layout", "sin/cos cache or generated constants must be identified"), - ), - "rope_gptj": FusionGraph( - name="rope_gptj_decomposed", - target="rope_gptj", - cactus_inputs=("x",), - cactus_attrs=("theta", "position_offset", "rot_dim", "backend"), - nodes={"rope": n("rope", "rope_gptj_subgraph")}, - root="rope", - outputs=("rope",), - inputs=(inp("x", "rope", 0),), - output_attrs=(cap("theta", default=None, required=False), cap("position_offset", default=0), cap("rot_dim", default=0), cap("backend", default=0)), - constraints=("GPT-J interleaved rotary layout check",), - ), - "sum": direct("sum", ("x",), ("axis",), "aten.sum.dim_IntList", attr_map={"axis": "dim"}), - "mean": direct("mean", ("x",), ("axis", "keepdim"), "aten.mean.dim", attr_map={"axis": "dim"}), - "variance": direct("variance", ("x",), ("axis", "keepdim"), "aten.var.dim", attr_map={"axis": "dim"}), - "min": direct("min", ("x",), ("axis",), "aten.min.dim", attr_map={"axis": "dim"}), - "max": direct("max", ("x",), ("axis",), "aten.max.dim", attr_map={"axis": "dim"}), - "cumsum": direct("cumsum", ("x",), ("axis",), "aten.cumsum.default", attr_map={"axis": "dim"}), - "softmax": direct("softmax", ("x",), ("axis",), "aten._softmax.default", attr_map={"axis": "dim"}), - "attention": FusionGraph( - name="attention_direct_or_decomposed", - target="attention", - cactus_inputs=("query", "key", "value", "mask"), - cactus_attrs=("scale", "is_causal", "position_offset", "window_size", "backend", "use_mask", "additive_mask"), - nodes={ - "query": n("query", "external"), - "key": n("key", "external"), - "value": n("value", "external"), - "k_transpose": n("k_transpose", ("aten.transpose.int", "aten.permute.default"), optional=True, transparent=True), - "scores": n("scores", ("aten.matmul.default", "aten.bmm.default")), - "scale": n("scale", ("aten.mul.Scalar", "aten.div.Scalar"), optional=True), - "mask_apply": n("mask_apply", ("aten.add.Tensor", "aten.where.self"), optional=True), - "softmax": n("softmax", "aten._softmax.default"), - "context": n("context", ("aten.matmul.default", "aten.bmm.default")), - "direct_sdpa": n("direct_sdpa", "aten.scaled_dot_product_attention.default", optional=True), - }, - edges=(e("key", "k_transpose", 0, optional=True), e("query", "scores", 0), e("k_transpose", "scores", 1, optional=True), e("scores", "scale", 0, optional=True), e("scale", "mask_apply", 0, optional=True), e("mask_apply", "softmax", 0, optional=True), e("scores", "softmax", 0, optional=True), e("softmax", "context", 0), e("value", "context", 1)), - root="context", - outputs=("context",), - inputs=(inp("query", "scores", 0), inp("key", "scores", 1), inp("value", "context", 1), inp("mask", "mask_apply", 1, optional=True)), - output_attrs=(cap("scale", "scale", "scalar", default=None, required=False), cap("is_causal", "direct_sdpa", "is_causal", default=False, required=False), cap("position_offset", default=0), cap("window_size", default=0), cap("backend", default=0), cap("use_mask", default=False), cap("additive_mask", default=False)), - constraints=("q/k head_dim match", "softmax axis is score last dimension", "mask semantics are boolean/additive/causal", "RoPE, if fused separately, must already be applied to q and k"), - variants=("aten.scaled_dot_product_attention.default", "matmul_scale_mask_softmax_matmul", "bmm_scale_mask_softmax_bmm"), - ), - "kv_cache_state": source_op("kv_cache_state", ("max_seq_len", "num_kv_heads", "head_dim", "window_size", "sink_size", "num_slots"), "State allocation op; not discovered from ordinary ATen math subgraph."), - "kv_cache_append": direct("kv_cache_append", ("new_kv", "cache_state"), ("window_size", "sink_size"), "kv_cache_append_subgraph", optional_attr_defaults={"window_size": 0, "sink_size": 0}), - "attention_cached": FusionGraph( - name="attention_cached", - target="attention_cached", - cactus_inputs=("query", "key_new", "value_new", "k_cache_state", "v_cache_state"), - cactus_attrs=("scale", "position_offset", "window_size", "v_head_dim"), - nodes={"attention_cached": n("attention_cached", "cached_attention_subgraph")}, - root="attention_cached", - outputs=("attention_cached",), - inputs=tuple(inp(role, "attention_cached", i) for i, role in enumerate(("query", "key_new", "value_new", "k_cache_state", "v_cache_state"))), - output_attrs=(cap("scale", default=None, required=False), cap("position_offset", default=0), cap("window_size", default=0), cap("v_head_dim", default=None, required=False)), - constraints=("must identify cache state nodes", "must verify append/read ordering"), - ), - "conv_cache_state": source_op("conv_cache_state", ("window_size", "hidden_dim"), "State allocation op for streaming conv."), - "conv_cache_append": direct("conv_cache_append", ("new_data", "cache_state"), (), "conv_cache_append_subgraph"), - "conv_cache_initialize": direct("conv_cache_initialize", ("rows", "cache_state"), (), "conv_cache_initialize_subgraph"), - "recurrent_cache_state": source_op("recurrent_cache_state", ("shape", "dtype"), "State allocation op for recurrent kernels."), - "recurrent_cache_write": direct("recurrent_cache_write", ("new_value", "cache_input"), (), "recurrent_cache_write_subgraph"), - "rel_pos_bias": direct("rel_pos_bias", ("query", "relative_key"), ("scale",), "relative_position_bias_subgraph"), - "attention_int8_hybrid": direct( - "attention_int8_hybrid", - ("query", "key_new", "value_new"), - ("scale", "position_offset", "cache_len", "num_kv_heads", "head_dim", "window_size"), - "attention_int8_hybrid_subgraph", - notes="Requires quantized/cache-aware attention trace.", - ), - "relu": direct("relu", ("x",), (), "aten.relu.default"), - "silu": FusionGraph( - name="silu_direct_or_sigmoid_mul", - target="silu", - cactus_inputs=("x",), - cactus_attrs=(), - nodes={"sigmoid": n("sigmoid", "aten.sigmoid.default", optional=True), "mul": n("mul", "aten.mul.Tensor")}, - edges=(e("sigmoid", "mul", 0, optional=True),), - root="mul", - outputs=("mul",), - inputs=(inp("x", "mul", 1),), - constraints=("mul other input and sigmoid input must be same x",), - variants=("aten.silu.default", "sigmoid_mul"), - ), - "gelu": direct("gelu", ("x",), (), "aten.gelu.default"), - "gelu_erf": direct("gelu_erf", ("x",), (), "gelu_erf_subgraph", notes="Detect erf-based GELU decomposition."), - "sigmoid": direct("sigmoid", ("x",), (), "aten.sigmoid.default"), - "tanh": direct("tanh", ("x",), (), "aten.tanh.default"), - "glu": direct("glu", ("x",), ("axis",), "aten.glu.default", attr_map={"axis": "dim"}), - "conv1d_causal": direct("conv1d_causal", ("x", "weight"), ("kernel_size", "dilation"), "aten.convolution.default", notes="Requires padding/causal layout check."), - "conv1d_k3": direct("conv1d_k3", ("x", "weight"), ("stride",), "aten.convolution.default", optional_attr_defaults={"stride": 1}, notes="Requires weight kernel_size == 3."), - "conv1d_k7s3": direct("conv1d_k7s3", ("x", "weight", "bias"), (), "aten.convolution.default", notes="Requires kernel_size == 7 and stride == 3."), - "conv1d": direct("conv1d", ("x", "weight", "bias"), ("stride",), ("aten.conv1d.default", "aten.convolution.default"), optional_attr_defaults={"stride": 1}), - "conv1d_same_depthwise_k9": direct("conv1d_same_depthwise_k9", ("x", "weight", "bias"), (), "aten.convolution.default", notes="Requires groups == channels, kernel_size == 9, same padding."), - "conv1d_pointwise": direct("conv1d_pointwise", ("x", "weight", "bias"), (), "aten.convolution.default", notes="Requires kernel_size == 1."), - "conv2d_k3s2p1": direct("conv2d_k3s2p1", ("x", "weight", "bias"), (), "aten.convolution.default", notes="Requires kernel_size == 3, stride == 2, padding == 1."), - "conv2d_depthwise_k3s2p1": direct("conv2d_depthwise_k3s2p1", ("x", "weight", "bias"), (), "aten.convolution.default", notes="Requires depthwise groups and k3s2p1."), - "conv2d_pointwise_1x1": direct("conv2d_pointwise_1x1", ("x", "weight", "bias"), (), "aten.convolution.default", notes="Requires kernel_size == 1x1."), - "conv2d_k3s1p1": direct("conv2d_k3s1p1", ("x", "weight", "bias"), (), "aten.convolution.default", notes="Requires kernel_size == 3, stride == 1, padding == 1."), - "conv2d": direct("conv2d", ("x", "weight", "bias"), ("stride", "padding", "dilation", "groups"), ("aten.conv2d.default", "aten.convolution.default"), optional_attr_defaults={"stride": 1, "padding": 0, "dilation": 1, "groups": 1}), - "lstm_cell": direct("lstm_cell", ("input", "h_prev", "c_prev", "weight_ih", "weight_hh", "bias_ih", "bias_hh"), (), "lstm_cell_subgraph", notes="Requires gate split/order checks: input, forget, cell, output."), - "gated_deltanet_decode": direct("gated_deltanet_decode", ("query", "key", "value", "gate_log", "beta", "initial_state"), ("scale",), "gated_deltanet_decode_subgraph"), - "gated_deltanet_prefill": direct("gated_deltanet_prefill", ("query", "key", "value", "gate_log", "beta", "initial_state"), ("chunk_size", "scale"), "gated_deltanet_prefill_subgraph"), - "stft": direct("stft", ("x", "weight"), ("stride", "num_fft_bins"), ("aten.stft.default", "stft_subgraph")), - "altup_predict": direct("altup_predict", ("coefs", "streams"), (), "altup_predict_subgraph"), - "altup_correct": direct("altup_correct", ("coefs", "innovation", "predictions"), (), "altup_correct_subgraph"), - "gaussian_topk": direct("gaussian_topk", ("x",), ("ppf",), "gaussian_topk_subgraph"), - "moe_layer_gated": FusionGraph( - name="moe_layer_gated_topk_experts", - target="moe_layer_gated", - cactus_inputs=("hidden", "routing_probs", "topk_indices", "w1_weights", "w3_weights", "w2_weights"), - cactus_attrs=("num_experts", "num_experts_per_tok", "normalize_routing", "epsilon", "routed_scaling_factor"), - nodes={ - "router_logits": n("router_logits", ("aten.mm.default", "aten.matmul.default", "aten.linear.default")), - "routing_probs": n("routing_probs", ("aten._softmax.default", "aten.sigmoid.default")), - "topk": n("topk", "aten.topk.default"), - "topk_values": n("topk_values", ("operator.getitem", "")), - "topk_indices": n("topk_indices", ("operator.getitem", "")), - "renorm": n("renorm", ("aten.sum.dim_IntList", "aten.div.Tensor"), optional=True), - "expert_gate": n("expert_gate", ("aten.mm.default", "aten.matmul.default", "aten.linear.default"), repeated=True), - "expert_up": n("expert_up", ("aten.mm.default", "aten.matmul.default", "aten.linear.default"), repeated=True), - "expert_act": n("expert_act", ("aten.silu.default", "aten.gelu.default", "aten.relu.default"), repeated=True), - "expert_mul": n("expert_mul", "aten.mul.Tensor", repeated=True), - "expert_down": n("expert_down", ("aten.mm.default", "aten.matmul.default", "aten.linear.default"), repeated=True), - "weighted_expert": n("weighted_expert", "aten.mul.Tensor", repeated=True), - "combine": n("combine", ("aten.add.Tensor", "aten.index_put.default", "aten.scatter_add.default")), - }, - edges=(e("router_logits", "routing_probs", 0), e("routing_probs", "topk", 0), e("topk", "topk_values", 0, source_output_index=0), e("topk", "topk_indices", 0, source_output_index=1), e("expert_gate", "expert_act", 0, repeated=True), e("expert_act", "expert_mul", 0, repeated=True), e("expert_up", "expert_mul", 1, repeated=True), e("expert_mul", "expert_down", 0, repeated=True), e("expert_down", "weighted_expert", 0, repeated=True), e("topk_values", "weighted_expert", 1, repeated=True), e("weighted_expert", "combine", 0, repeated=True)), - root="combine", - outputs=("combine",), - inputs=(inp("hidden", "router_logits", 0), inp("routing_probs", "routing_probs", 0), inp("topk_indices", "topk_indices", 0), inp("w1_weights", "expert_gate", 1, variadic=True, repeated_group="experts"), inp("w3_weights", "expert_up", 1, variadic=True, repeated_group="experts"), inp("w2_weights", "expert_down", 1, variadic=True, repeated_group="experts")), - output_attrs=(cap("num_experts", default=None, required=False), cap("num_experts_per_tok", "topk", "k"), cap("normalize_routing", default=False), cap("epsilon", default=1e-6), cap("routed_scaling_factor", default=1.0)), - constraints=("topk getitem 0/1 must map values/indices", "expert branch count and weight order must be stable", "router hidden input must be shared with expert branches", "combine/scatter must preserve token order"), - repeated_groups=("experts",), - ), - "dense_mlp_tq_fused": direct("dense_mlp_tq_fused", ("hidden", "gate_weight", "up_weight", "down_weight"), ("product_scale",), "dense_mlp_tq_fused_subgraph", optional_attr_defaults={"product_scale": 1.0}), - "moe_layer_ungated": FusionGraph( - name="moe_layer_ungated_topk_experts", - target="moe_layer_ungated", - cactus_inputs=("hidden", "routing_probs", "topk_indices", "w1_weights", "w2_weights"), - cactus_attrs=("num_experts", "num_experts_per_tok", "normalize_routing", "epsilon", "routed_scaling_factor", "activation"), - nodes={ - "router_logits": n("router_logits", ("aten.mm.default", "aten.matmul.default", "aten.linear.default")), - "routing_probs": n("routing_probs", ("aten._softmax.default", "aten.sigmoid.default")), - "topk": n("topk", "aten.topk.default"), - "topk_values": n("topk_values", ("operator.getitem", "")), - "topk_indices": n("topk_indices", ("operator.getitem", "")), - "expert_up": n("expert_up", ("aten.mm.default", "aten.matmul.default", "aten.linear.default"), repeated=True), - "expert_act": n("expert_act", ("aten.gelu.default", "aten.relu.default", "aten.silu.default"), repeated=True), - "expert_down": n("expert_down", ("aten.mm.default", "aten.matmul.default", "aten.linear.default"), repeated=True), - "weighted_expert": n("weighted_expert", "aten.mul.Tensor", repeated=True), - "combine": n("combine", ("aten.add.Tensor", "aten.index_put.default", "aten.scatter_add.default")), - }, - edges=(e("router_logits", "routing_probs", 0), e("routing_probs", "topk", 0), e("topk", "topk_values", 0, source_output_index=0), e("topk", "topk_indices", 0, source_output_index=1), e("expert_up", "expert_act", 0, repeated=True), e("expert_act", "expert_down", 0, repeated=True), e("expert_down", "weighted_expert", 0, repeated=True), e("topk_values", "weighted_expert", 1, repeated=True), e("weighted_expert", "combine", 0, repeated=True)), - root="combine", - outputs=("combine",), - inputs=(inp("hidden", "router_logits", 0), inp("routing_probs", "routing_probs", 0), inp("topk_indices", "topk_indices", 0), inp("w1_weights", "expert_up", 1, variadic=True, repeated_group="experts"), inp("w2_weights", "expert_down", 1, variadic=True, repeated_group="experts")), - output_attrs=(cap("num_experts", default=None, required=False), cap("num_experts_per_tok", "topk", "k"), cap("normalize_routing", default=False), cap("epsilon", default=1e-6), cap("routed_scaling_factor", default=1.0), cap("activation", "expert_act", "op")), - constraints=("topk tuple outputs must be distinguished", "activation enum must be derived from expert_act op", "expert order must match weight lists"), - repeated_groups=("experts",), - ), - "sample": direct("sample", ("logits",), ("temperature", "top_p", "top_k"), "sampling_subgraph", optional_attr_defaults={"temperature": 0.6, "top_p": 0.95, "top_k": 20}), - "scatter_topk": direct("scatter_topk", ("indices", "values"), ("num_classes",), "scatter_topk_subgraph"), - "persistent": direct("persistent", ("source_node",), (), "persistent_subgraph"), - "rfft": direct("rfft", ("x",), (), "aten.fft_rfft.default"), - "irfft": direct("irfft", ("x",), ("output_length",), "aten.fft_irfft.default"), - "mel_filter_bank": source_op("mel_filter_bank", ("num_frequency_bins", "num_mel_filters", "min_frequency", "max_frequency", "sampling_rate", "norm_type", "scale_type"), "Generated constant/source op for audio preprocessing."), - "spectrogram": direct("spectrogram", ("waveform", "mel_filters"), ("frame_length", "hop_length", "fft_length", "power", "center", "pad_mode", "mel_floor", "log_mel_mode", "dither", "preemphasis", "remove_dc_offset"), "spectrogram_subgraph"), - "image_preprocess": direct("image_preprocess", ("pixel_input",), ("src_width", "src_height", "target_width", "target_height", "patch_size", "channels", "rescale_factor", "mean", "std_dev"), "image_preprocess_subgraph"), -} - - -ALL_FUSION_GRAPHS: tuple[FusionGraph, ...] = tuple(FUSION_GRAPHS.values()) diff --git a/python/cactus/transpiler/ir/fusions_utils.py b/python/cactus/transpiler/ir/fusions_utils.py deleted file mode 100644 index d7219bc56..000000000 --- a/python/cactus/transpiler/ir/fusions_utils.py +++ /dev/null @@ -1,436 +0,0 @@ - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -# import models, nodes, fusions - - - - - -# """###################################### FUSION UTILS!!!!!!! ######################################""" - -# def attr_match(expected_value, actual_value, node: models.Node) -> bool: -# if hasattr(expected_value, "matches"): -# return expected_value.matches(actual_value, node) - -# return expected_value == actual_value - -# def node_match(node:models.Node, synth_node:models.NodeSpec, bindings: dict[str, models.Node] | None = None) -> bool: -# underlying_op = synth_node.underlying_op - -# if isinstance(underlying_op, tuple): -# if node.underlying_op not in underlying_op: -# return False -# elif node.underlying_op != underlying_op: -# return False - -# allowed_value_kinds = getattr(synth_node, "allowed_value_kinds", ()) -# if allowed_value_kinds and getattr(node, "value_kind", None) not in allowed_value_kinds: -# return False - -# for key, expected_value in synth_node.required_attrs.items(): -# if key not in node.normalized_attrs: -# return False - -# if not attr_match(expected_value, node.normalized_attrs[key], node): -# return False - -# for constraint in getattr(synth_node, "tensor_constraints", ()): -# if not constraint.matches(node, bindings): -# return False - -# return True - -# #Returns the desired parent if it is in the parents list -# def get_bound_parent(bindings: dict[str, models.Node], node_name: str, parent_index: int) -> models.Node | None: -# if node_name not in bindings: -# return None - -# node = bindings[node_name] -# if parent_index < 0 or parent_index >= len(node.parents): -# return None - -# return node.parents[parent_index] - -# #Returns parent base on InputSpec object -# def get_input_parent(bindings: dict[str, models.Node], input_spec: models.InputSpec) -> models.Node | None: -# return get_bound_parent(bindings, input_spec.node, input_spec.parent_index) - -# def get_variadic_input_parents(bindings: dict[str, models.Node], variadic_input) -> list[models.Node] | None: -# if variadic_input.node not in bindings: -# return None - -# node = bindings[variadic_input.node] -# start = variadic_input.start_parent_index -# end = variadic_input.end_parent_index if variadic_input.end_parent_index is not None else len(node.parents) - -# if start < 0 or start > len(node.parents): -# return None - -# if end < start or end > len(node.parents): -# return None - -# parents = node.parents[start:end] - -# if len(parents) < variadic_input.min_count: -# return None - -# if variadic_input.max_count is not None and len(parents) > variadic_input.max_count: -# return None - -# return parents - -# def match_input_constraints(parent: models.Node, input_spec, bindings: dict[str, models.Node]) -> bool: -# allowed_value_kinds = getattr(input_spec, "allowed_value_kinds", ()) -# if allowed_value_kinds and getattr(parent, "value_kind", None) not in allowed_value_kinds: -# return False - -# for constraint in getattr(input_spec, "tensor_constraints", ()): -# if not constraint.matches(parent, bindings): -# return False - -# return True - -# #Match edges and returns binding -# def edges_match(node: models.Node, synth_graph: models.FusionGraph) -> dict[str, models.Node] | None: -# bindings:dict[str,models.Node] = {} - -# if not node_match(node, nodes.NODE_MAP[synth_graph.root]): -# return None - -# bindings[synth_graph.root] = node - -# for edge in synth_graph.edges: -# parent = get_bound_parent(bindings, edge.from_node, edge.parent_index) -# if parent is None: -# if getattr(edge, "required", True): -# return None -# continue - -# if edge.to_node not in nodes.NODE_MAP: -# return None - -# if not node_match(parent, nodes.NODE_MAP[edge.to_node], bindings): -# return None - -# to_output_index = getattr(edge, "to_output_index", None) -# if to_output_index is not None: -# tuple_source = getattr(parent, "tuple_source", None) -# if tuple_source is None or tuple_source.output_index != to_output_index: -# return None - -# if edge.to_node in bindings and bindings[edge.to_node] is not parent: -# return None - -# bindings[edge.to_node] = parent - -# return bindings - -# #Ensures inputs exist within the fusion -# def match_inputs_exist(bindings: dict[str, models.Node], synth_graph: models.FusionGraph) -> bool: -# for input_spec in synth_graph.inputs: -# parent = get_input_parent(bindings, input_spec) -# if parent is None: -# return False - -# if not match_input_constraints(parent, input_spec, bindings): -# return False - -# for variadic_input in getattr(synth_graph, "variadic_inputs", ()): -# parents = get_variadic_input_parents(bindings, variadic_input) -# if parents is None: -# return False - -# for parent in parents: -# if not match_input_constraints(parent, variadic_input, bindings): -# return False - -# return True - -# #Enusres all of a nodes inputs are included in the fusion -# def match_external_inputs_dec(bindings: dict[str, models.Node], synth_graph: models.FusionGraph) -> bool: -# internal_node_ids = {id(node) for node in bindings.values()} -# declared_input_ids = set() - -# for input_spec in synth_graph.inputs: -# parent = get_input_parent(bindings, input_spec) -# if parent is None: -# return False - -# declared_input_ids.add(id(parent)) - -# for variadic_input in getattr(synth_graph, "variadic_inputs", ()): -# parents = get_variadic_input_parents(bindings, variadic_input) -# if parents is None: -# return False - -# for parent in parents: -# declared_input_ids.add(id(parent)) - -# for node in bindings.values(): -# for parent in node.parents: -# if id(parent) not in internal_node_ids and id(parent) not in declared_input_ids: -# return False - -# return True - -# #Ensures named input roles point to the expected kind of value/tensor -# def match_input_roles(bindings: dict[str, models.Node], synth_graph: models.FusionGraph) -> bool: -# role_map = {} - -# for input_spec in synth_graph.inputs: -# role = getattr(input_spec, "role", None) -# if role is None: -# continue - -# parent = get_input_parent(bindings, input_spec) -# if parent is None: -# return False - -# role_map.setdefault(role, []) -# role_map[role].append(parent) - -# for variadic_input in getattr(synth_graph, "variadic_inputs", ()): -# role = getattr(variadic_input, "role", None) -# if role is None: -# continue - -# parents = get_variadic_input_parents(bindings, variadic_input) -# if parents is None: -# return False - -# role_map.setdefault(role, []) -# role_map[role].extend(parents) - -# for role_spec in getattr(synth_graph, "input_roles", ()): -# if role_spec.input_name not in role_map: -# return False - -# for parent in role_map[role_spec.input_name]: -# if not match_input_constraints(parent, role_spec, bindings): -# return False - -# return True - -# #Ensures ops with same inputs are actually pointing to same input node -# def match_shared_inputs(bindings: dict[str, models.Node], synth_graph: models.FusionGraph) -> bool: -# for left_input, right_input in synth_graph.shared_inputs: -# left_parent = get_input_parent(bindings, left_input) -# right_parent = get_input_parent(bindings, right_input) -# if left_parent is None or right_parent is None: -# return False - -# if left_parent is not right_parent: -# return False - -# return True - -# #Ensures that a node within fusion is not called by a differnt node not in the fusion -# def match_no_external_internal_children(bindings: dict[str, models.Node], synth_graph: models.FusionGraph) -> bool: -# if synth_graph.root not in bindings: -# return False - -# root_node = bindings[synth_graph.root] -# internal_node_ids = {id(node) for node in bindings.values()} -# allow_root_external_children = getattr(synth_graph, "allow_root_external_children", True) - -# for node in bindings.values(): -# if node is root_node and allow_root_external_children: -# continue - -# for child in node.children: -# if id(child) not in internal_node_ids: -# return False - -# return True - -# def is_attr_ref(value) -> bool: -# return hasattr(value, "node") and hasattr(value, "attr") - -# def match_output_attrs(bindings: dict[str, models.Node], synth_graph: models.FusionGraph) -> bool: -# for attr_value in synth_graph.output_attrs.values(): -# if not is_attr_ref(attr_value): -# continue - -# if attr_value.node not in bindings: -# if getattr(attr_value, "required", True): -# return False -# continue - -# node = bindings[attr_value.node] -# if attr_value.attr not in node.normalized_attrs: -# if getattr(attr_value, "required", True): -# return False -# continue - -# return True - -# def match_optional_nodes(bindings: dict[str, models.Node], synth_graph: models.FusionGraph) -> bool: -# for optional_node in getattr(synth_graph, "optional_nodes", ()): -# if optional_node.node not in bindings: -# continue - -# node = bindings[optional_node.node] -# if optional_node.bypass_parent_index < 0 or optional_node.bypass_parent_index >= len(node.parents): -# return False - -# return True - -# def match_tuple_outputs(bindings: dict[str, models.Node], synth_graph: models.FusionGraph) -> bool: -# for tuple_output in getattr(synth_graph, "tuple_outputs", ()): -# if tuple_output.node not in bindings or tuple_output.source_node not in bindings: -# return False - -# node = bindings[tuple_output.node] -# source_node = bindings[tuple_output.source_node] -# tuple_source = getattr(node, "tuple_source", None) - -# if tuple_source is None: -# return False - -# if tuple_source.source_node_name != source_node.name: -# return False - -# if tuple_source.output_index != tuple_output.output_index: -# return False - -# return True - -# def match_repeated_subgraphs(bindings: dict[str, models.Node], synth_graph: models.FusionGraph) -> bool: -# repeated_bindings = getattr(synth_graph, "metadata", {}).get("repeated_bindings", {}) - -# for repeated in getattr(synth_graph, "repeated_subgraphs", ()): -# if repeated.anchor_node is not None and repeated.anchor_node not in bindings: -# return False - -# matches = repeated_bindings.get(repeated.name, []) - -# if len(matches) < repeated.min_count: -# return False - -# if repeated.max_count is not None and len(matches) > repeated.max_count: -# return False - -# return True - - -# #Update this with any additional matching functions if any are added -# MATCHERS = ( -# match_inputs_exist, -# match_external_inputs_dec, -# match_input_roles, -# match_shared_inputs, -# match_no_external_internal_children, -# match_output_attrs, -# match_optional_nodes, -# match_tuple_outputs, -# match_repeated_subgraphs, -# ) - -# def fusion_match(node: models.Node) -> models.FusionResult | None: - -# for fusion in fusions.ROOT_TARGET_MAP.get(node.underlying_op, ()): -# bindings = edges_match(node, fusion) -# if bindings != None and all(matcher(bindings, fusion) for matcher in MATCHERS): -# return models.FusionResult.from_fusion(fusion, bindings) - -# return None diff --git a/python/cactus/transpiler/ir/models.py b/python/cactus/transpiler/ir/models.py deleted file mode 100644 index f56548ae0..000000000 --- a/python/cactus/transpiler/ir/models.py +++ /dev/null @@ -1,834 +0,0 @@ - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -# from __future__ import annotations - -# from dataclasses import dataclass, field -# from typing import Any, Callable, Optional -# import converter.models as CVModels -# import constants - - -# class ValueKind: -# UNKNOWN = "unknown" -# PARAMETER = "parameter" -# BUFFER = "buffer" -# USER_INPUT = "user_input" -# LIFTED_CONSTANT = "lifted_constant" -# ACTIVATION = "activation" -# OUTPUT = "output" - - -# @dataclass(slots=True) -# class TensorMeta: -# shape: list[int] = field(default_factory=list) -# dtype: str | None = None -# device: str | None = None -# requires_grad: bool | None = None -# stride: list[int] | None = None - -# @property -# def rank(self) -> int: -# return len(self.shape) - -# def dim(self, index: int) -> int | None: -# if index < 0: -# index += len(self.shape) - -# if index < 0 or index >= len(self.shape): -# return None - -# return self.shape[index] - - -# @dataclass(slots=True) -# class TupleOutputRef: -# source_node_name: str -# output_index: int - - -# @dataclass(slots=True) -# class Node: -# name: str -# underlying_op: str -# normalized_attrs: dict[str, Any] -# parents: list["Node"] -# children: list["Node"] -# tensor_meta: TensorMeta | None = None -# tuple_output_metas: list[TensorMeta | None] = field(default_factory=list) -# value_kind: str = ValueKind.UNKNOWN -# tuple_source: TupleOutputRef | None = None -# raw_layer: Optional[CVModels.LayerRecord] = None - - -# @dataclass(frozen=True, slots=True) -# class AttrPredicate: -# description: str -# predicate: Callable[[Any, Node], bool] = field(repr=False, compare=False) - -# def matches(self, value: Any, node: Node) -> bool: -# return self.predicate(value, node) - - -# def attr_any() -> AttrPredicate: -# return AttrPredicate("any value", lambda _value, _node: True) - - -# def attr_in(values: tuple[Any, ...]) -> AttrPredicate: -# return AttrPredicate(f"in {values}", lambda value, _node: value in values) - - -# def attr_rank_dim(dim_values: tuple[int, ...]) -> AttrPredicate: -# def matches(value: Any, node: Node) -> bool: -# if value in dim_values: -# return True - -# if node.tensor_meta is None: -# return False - -# normalized_values = tuple(dim + node.tensor_meta.rank if dim < 0 else dim for dim in dim_values) -# normalized_value = value + node.tensor_meta.rank if isinstance(value, int) and value < 0 else value -# return normalized_value in normalized_values - -# return AttrPredicate(f"rank-aware dim in {dim_values}", matches) - - -# @dataclass(frozen=True, slots=True) -# class DimRef: -# node: str -# dim: int - - -# @dataclass(frozen=True, slots=True) -# class TensorConstraint: -# rank: int | None = None -# min_rank: int | None = None -# max_rank: int | None = None -# dtype: str | None = None -# dim_equals: tuple[tuple[int, int], ...] = () -# same_dim_as: tuple[tuple[int, DimRef], ...] = () - -# def matches(self, node: Node, bindings: dict[str, Node] | None = None) -> bool: -# if node.tensor_meta is None: -# return False - -# meta = node.tensor_meta - -# if self.rank is not None and meta.rank != self.rank: -# return False - -# if self.min_rank is not None and meta.rank < self.min_rank: -# return False - -# if self.max_rank is not None and meta.rank > self.max_rank: -# return False - -# if self.dtype is not None and meta.dtype != self.dtype: -# return False - -# for dim_index, expected in self.dim_equals: -# if meta.dim(dim_index) != expected: -# return False - -# if bindings is None and self.same_dim_as: -# return False - -# for dim_index, ref in self.same_dim_as: -# other = bindings.get(ref.node) if bindings is not None else None -# if other is None or other.tensor_meta is None: -# return False - -# if meta.dim(dim_index) != other.tensor_meta.dim(ref.dim): -# return False - -# return True - - -# @dataclass(slots=True) -# class NodeSpec: -# underlying_op: str | tuple[str, ...] -# required_attrs: dict[str, Any] = field(default_factory=dict) -# allowed_value_kinds: tuple[str, ...] = () -# tensor_constraints: tuple[TensorConstraint, ...] = () -# optional: bool = False - - -# @dataclass(slots=True) -# class Edge: -# from_node: str -# parent_index: int -# to_node: str -# required: bool = True -# to_output_index: int | None = None - - -# @dataclass(slots=True) -# class InputSpec: -# node: str -# parent_index: int -# role: str | None = None -# allowed_value_kinds: tuple[str, ...] = () -# tensor_constraints: tuple[TensorConstraint, ...] = () - - -# @dataclass(slots=True) -# class VariadicInputSpec: -# node: str -# start_parent_index: int = 0 -# end_parent_index: int | None = None -# min_count: int = 0 -# max_count: int | None = None -# role: str | None = None -# allowed_value_kinds: tuple[str, ...] = () -# tensor_constraints: tuple[TensorConstraint, ...] = () - - -# @dataclass(slots=True) -# class InputRoleSpec: -# input_name: str -# role: str -# allowed_value_kinds: tuple[str, ...] = () -# tensor_constraints: tuple[TensorConstraint, ...] = () - - -# @dataclass(slots=True) -# class OptionalNodeSpec: -# node: str -# bypass_parent_index: int = 0 - - -# @dataclass(slots=True) -# class TupleOutputSpec: -# node: str -# source_node: str -# output_index: int - - -# @dataclass(slots=True) -# class RepeatedSubgraphSpec: -# name: str -# pattern: "FusionGraph" -# min_count: int = 1 -# max_count: int | None = None -# anchor_node: str | None = None - - -# @dataclass(slots=True) -# class AttrRef: -# node: str -# attr: str -# default: Any = None -# required: bool = True - - -# @dataclass(slots=True) -# class FusionGraph: -# target: str -# root: str -# edges: tuple[Edge, ...] -# inputs: tuple[InputSpec, ...] = () -# shared_inputs: tuple[tuple[InputSpec, InputSpec], ...] = () -# output_attrs: dict[str, Any] = field(default_factory=dict) -# optional_nodes: tuple[OptionalNodeSpec, ...] = () -# variadic_inputs: tuple[VariadicInputSpec, ...] = () -# input_roles: tuple[InputRoleSpec, ...] = () -# tuple_outputs: tuple[TupleOutputSpec, ...] = () -# repeated_subgraphs: tuple[RepeatedSubgraphSpec, ...] = () -# metadata: dict[str, Any] = field(default_factory=dict) -# allow_root_external_children: bool = True - - -# @dataclass(slots=True) -# class FusionResult: -# node: Node -# bindings: dict[str, Node] -# internal_nodes: list[Node] -# external_input: list[Node] -# external_consumers: list[Node] -# edges_consumed: list[str] -# fusion: FusionGraph | None = None -# variadic_external_inputs: dict[str, list[Node]] = field(default_factory=dict) -# input_role_nodes: dict[str, list[Node]] = field(default_factory=dict) - -# @classmethod -# def from_fusion(cls, fusion: FusionGraph, bindings: dict[str, Node]): -# return build_fusion_result(fusion, bindings) - - -# @dataclass(slots=True) -# class Graph: -# name_map: dict[str, Node] -# source_nodes: list[Node] -# consumed_ids: set[str] -# fusedNodes: list[FusionResult] = field(default_factory=list) - -# @classmethod -# def from_layer_map(cls, map: CVModels.LayerMap) -> "Graph": -# return generate_graph(map) - -# def queue_fusion(self, fusion_result: FusionResult) -> None: -# self.fusedNodes.append(fusion_result) - -# def update_graph_with_fusions(self) -> bool: -# if not len(self.fusedNodes): -# return False - -# update_graph(self) -# return True - - -# """###################################### MODEL UTILS!!!!!!! ######################################""" - -# # TODO: Understand -# def extract_node_names(value: Any) -> list[str]: -# if isinstance(value, dict): -# refs = [] -# if "node" in value: -# refs.append(value["node"]) - -# for item in value.values(): -# refs.extend(extract_node_names(item)) - -# return refs - -# if isinstance(value, list): -# refs = [] -# for item in value: -# refs.extend(extract_node_names(item)) -# return refs - -# return [] - -# # TODO: Understand -# def extract_attrs(layer_: CVModels.LayerRecord) -> dict[str, Any]: -# def contains_node_ref(value: Any) -> bool: -# if isinstance(value, dict): -# if "node" in value: -# return True -# return any(contains_node_ref(item) for item in value.values()) - -# if isinstance(value, list): -# return any(contains_node_ref(item) for item in value) - -# return False - -# attrs = {} - -# if isinstance(layer_.kwargs, dict): -# attrs.update(layer_.kwargs) - -# attr_names = constants.LAYER_ATTRS_MAP.get(layer_.target, []) - -# if None in attr_names: -# for i, attr_name in enumerate(attr_names): -# if attr_name is None or i >= len(layer_.args): -# continue - -# if contains_node_ref(layer_.args[i]) or attr_name in attrs: -# continue - -# attrs[attr_name] = layer_.args[i] -# else: -# positional_attrs = [item for item in layer_.args if not contains_node_ref(item)] - -# for i, attr_name in enumerate(attr_names): -# if i < len(positional_attrs) and attr_name not in attrs: -# attrs[attr_name] = positional_attrs[i] - -# return attrs - -# # TODO: Understand -# def attr_value_matches(expected_value: Any, actual_value: Any, node: Node) -> bool: -# if isinstance(expected_value, AttrPredicate): -# return expected_value.matches(actual_value, node) - -# return actual_value == expected_value - -# # TODO: Understand -# def node_spec_matches(node: Node, spec: NodeSpec, bindings: dict[str, Node] | None = None) -> bool: -# allowed_ops = spec.underlying_op if isinstance(spec.underlying_op, tuple) else (spec.underlying_op,) -# if node.underlying_op not in allowed_ops: -# return False - -# if spec.allowed_value_kinds and node.value_kind not in spec.allowed_value_kinds: -# return False - -# for key, expected_value in spec.required_attrs.items(): -# if key not in node.normalized_attrs: -# return False - -# if not attr_value_matches(expected_value, node.normalized_attrs[key], node): -# return False - -# for constraint in spec.tensor_constraints: -# if not constraint.matches(node, bindings): -# return False - -# return True - -# # TODO: Understand -# def _plain(value: Any) -> Any: -# if hasattr(value, "model_dump"): -# return value.model_dump() - -# if hasattr(value, "dict"): -# return value.dict() - -# return value - -# # TODO: Understand -# def _parse_one_tensor_meta(value: Any) -> TensorMeta | None: -# value = _plain(value) - -# if not isinstance(value, dict): -# return None - -# shape = value.get("shape") -# if shape is None: -# return None - -# return TensorMeta( -# shape=[int(item) for item in shape], -# dtype=value.get("dtype"), -# device=value.get("device"), -# requires_grad=value.get("requires_grad"), -# stride=value.get("stride"), -# ) - -# # TODO: Optimize this function further. -# def extract_tensor_metas(layer_: CVModels.LayerRecord) -> tuple[TensorMeta | None, list[TensorMeta | None]]: -# meta = _plain(layer_.tensor_output_meta) - -# single_meta = _parse_one_tensor_meta(meta) -# if single_meta is not None: -# return single_meta, [] - -# if isinstance(meta, list): -# tuple_metas = [_parse_one_tensor_meta(item) for item in meta] -# if tuple_metas: -# return None, tuple_metas - -# return None, [] - -# # TODO: Optimize this function further. -# def infer_value_kind(layer_: CVModels.LayerRecord) -> str: -# name = layer_.name -# target = layer_.target - -# if layer_.node_type == "output": -# return ValueKind.OUTPUT - -# if layer_.node_type == "placeholder": -# if name.startswith("p_") or "_weight" in name or name.endswith("weight"): -# return ValueKind.PARAMETER - -# if name.startswith("b_") or "buffer" in name: -# return ValueKind.BUFFER - -# return ValueKind.USER_INPUT - -# if layer_.node_type == "get_attr": -# return ValueKind.PARAMETER - -# if target in { -# "aten.scalar_tensor.default", -# "aten.full.default", -# "aten.arange.start_step", -# }: -# return ValueKind.LIFTED_CONSTANT - -# return ValueKind.ACTIVATION - -# # TODO: Optimize this function further. -# def extract_tuple_source(layer_: CVModels.LayerRecord) -> TupleOutputRef | None: -# target = layer_.target -# if target not in {"operator.getitem", ""}: -# return None - -# node_refs = extract_node_names(layer_.args) -# if not node_refs: -# return None - -# output_index = None -# if isinstance(layer_.args, list) and len(layer_.args) >= 2 and isinstance(layer_.args[1], int): -# output_index = layer_.args[1] - -# if output_index is None: -# return None - -# return TupleOutputRef(source_node_name=node_refs[0], output_index=output_index) - - -# def append_unique(nodes: list[Node], node: Node) -> None: -# if all(existing is not node for existing in nodes): -# nodes.append(node) - - -# def resolve_attr_ref(attr_name: str, attr_value: Any, bindings: dict[str, Node]) -> Any: -# if not isinstance(attr_value, AttrRef): -# return attr_value - -# if attr_value.node not in bindings: -# if attr_value.required: -# raise ValueError(f"Output attr {attr_name} references unbound node {attr_value.node}") -# return attr_value.default - -# bound_node = bindings[attr_value.node] -# if attr_value.attr not in bound_node.normalized_attrs: -# if attr_value.required: -# raise ValueError(f"Output attr {attr_name} references missing attr {attr_value.attr}") -# return attr_value.default - -# return bound_node.normalized_attrs[attr_value.attr] - - -# # TODO: Optimize this function further. -# def build_fusion_result(fusion: FusionGraph, bindings: dict[str, Node]) -> FusionResult: -# root_node = bindings[fusion.root] -# internal_nodes = [] - -# for node in bindings.values(): -# append_unique(internal_nodes, node) - -# internal_node_ids = {id(node) for node in internal_nodes} -# output_attrs = {} - -# for attr_name, attr_value in fusion.output_attrs.items(): -# output_attrs[attr_name] = resolve_attr_ref(attr_name, attr_value, bindings) - -# external_input = [] -# input_role_nodes = {} - -# for input_spec in fusion.inputs: -# if input_spec.node not in bindings: -# raise ValueError(f"Fusion input references unbound node {input_spec.node}") - -# node = bindings[input_spec.node] -# if input_spec.parent_index < 0 or input_spec.parent_index >= len(node.parents): -# raise ValueError(f"Fusion input has invalid parent index {input_spec.parent_index}") - -# parent = node.parents[input_spec.parent_index] -# append_unique(external_input, parent) - -# if input_spec.role is not None: -# input_role_nodes.setdefault(input_spec.role, []) -# append_unique(input_role_nodes[input_spec.role], parent) - -# variadic_external_inputs = {} - -# for variadic_input in fusion.variadic_inputs: -# if variadic_input.node not in bindings: -# raise ValueError(f"Variadic fusion input references unbound node {variadic_input.node}") - -# node = bindings[variadic_input.node] -# end = variadic_input.end_parent_index if variadic_input.end_parent_index is not None else len(node.parents) -# parents = node.parents[variadic_input.start_parent_index:end] - -# if len(parents) < variadic_input.min_count: -# raise ValueError(f"Variadic fusion input {variadic_input.node} has too few parents") - -# if variadic_input.max_count is not None and len(parents) > variadic_input.max_count: -# raise ValueError(f"Variadic fusion input {variadic_input.node} has too many parents") - -# role = variadic_input.role or variadic_input.node -# variadic_external_inputs[role] = [] - -# for parent in parents: -# append_unique(external_input, parent) -# append_unique(variadic_external_inputs[role], parent) - -# if variadic_input.role is not None: -# input_role_nodes.setdefault(variadic_input.role, []) -# append_unique(input_role_nodes[variadic_input.role], parent) - -# external_consumers = [] - -# for node in internal_nodes: -# for child in node.children: -# if id(child) in internal_node_ids: -# continue - -# if node is not root_node or not fusion.allow_root_external_children: -# raise ValueError("Cannot fuse node with external consumers outside the fusion boundary") - -# append_unique(external_consumers, child) - -# edges_consumed = [] - -# for child in internal_nodes: -# for parent in child.parents: -# if id(parent) in internal_node_ids: -# edges_consumed.append(f"{parent.name}->{child.name}") - -# fused_node = Node( -# name=f"{fusion.target}_{root_node.name}", -# underlying_op=fusion.target, -# normalized_attrs=output_attrs, -# parents=external_input, -# children=external_consumers, -# tensor_meta=root_node.tensor_meta, -# tuple_output_metas=root_node.tuple_output_metas.copy(), -# value_kind=ValueKind.ACTIVATION, -# tuple_source=None, -# raw_layer=None, -# ) - -# return FusionResult( -# node=fused_node, -# bindings=bindings, -# internal_nodes=internal_nodes, -# external_input=external_input, -# external_consumers=external_consumers, -# edges_consumed=edges_consumed, -# fusion=fusion, -# variadic_external_inputs=variadic_external_inputs, -# input_role_nodes=input_role_nodes, -# ) - -# # TODO: Optimize this function further. -# def generate_graph(map: CVModels.LayerMap) -> Graph: -# temp_map: dict[str, Node] = {} -# temp_source = [] - -# for layer_ in map.nodes: -# tensor_meta, tuple_output_metas = extract_tensor_metas(layer_) -# temp = Node( -# name=layer_.name, -# normalized_attrs=extract_attrs(layer_), -# underlying_op=layer_.target, -# parents=[], -# children=[], -# tensor_meta=tensor_meta, -# tuple_output_metas=tuple_output_metas, -# value_kind=infer_value_kind(layer_), -# tuple_source=extract_tuple_source(layer_), -# raw_layer=layer_, -# ) -# temp_map[layer_.name] = temp -# temp_source.append(temp) if layer_.node_type == "output" else None - -# for layer_ in map.nodes: -# for name in layer_.users: -# temp_map[layer_.name].children.append(temp_map[name]) - -# for name in extract_node_names(layer_.args): -# temp_map[layer_.name].parents.append(temp_map[name]) - -# for node in temp_map.values(): -# if node.tuple_source is None: -# continue - -# source_node = temp_map.get(node.tuple_source.source_node_name) -# output_index = node.tuple_source.output_index -# if source_node is None or output_index < 0 or output_index >= len(source_node.tuple_output_metas): -# continue - -# node.tensor_meta = source_node.tuple_output_metas[output_index] - -# return Graph(name_map=temp_map, source_nodes=temp_source, consumed_ids=set(), fusedNodes=[]) - - -# # TODO: Optimize this function further. -# def update_graph(graph: Graph) -> None: -# replacement_by_id: dict[int, Node] = {} -# fused_nodes = [] - -# for fusion_result in graph.fusedNodes: -# append_unique(fused_nodes, fusion_result.node) - -# for internal_node in fusion_result.internal_nodes: -# internal_id = id(internal_node) -# existing_replacement = replacement_by_id.get(internal_id) - -# if existing_replacement is not None and existing_replacement is not fusion_result.node: -# raise ValueError(f"Node {internal_node.name} is consumed by multiple fusions") - -# replacement_by_id[internal_id] = fusion_result.node - -# def resolve_node(node: Node) -> Node: -# seen_ids = set() - -# while id(node) in replacement_by_id and id(node) not in seen_ids: -# seen_ids.add(id(node)) -# node = replacement_by_id[id(node)] - -# return node - -# final_nodes = [] - -# for node in graph.name_map.values(): -# if id(node) not in replacement_by_id: -# append_unique(final_nodes, node) - -# for node in fused_nodes: -# append_unique(final_nodes, node) - -# new_name_map = {} - -# for node in final_nodes: -# if node.name in new_name_map and new_name_map[node.name] is not node: -# raise ValueError(f"Duplicate node name after fusion: {node.name}") - -# new_name_map[node.name] = node - -# final_node_ids = {id(node) for node in final_nodes} -# fusion_result_by_node_id = {id(fusion_result.node): fusion_result for fusion_result in graph.fusedNodes} - -# for node in final_nodes: -# fusion_result = fusion_result_by_node_id.get(id(node)) -# old_parents = fusion_result.external_input if fusion_result is not None else node.parents -# new_parents = [] - -# for parent in old_parents: -# resolved_parent = resolve_node(parent) - -# if resolved_parent is node: -# continue - -# if id(resolved_parent) not in final_node_ids: -# continue - -# new_parents.append(resolved_parent) - -# node.parents = new_parents -# node.children = [] - -# if node.tuple_source is not None: -# source_node = graph.name_map.get(node.tuple_source.source_node_name) -# if source_node is not None: -# resolved_source = resolve_node(source_node) -# if resolved_source is not source_node: -# node.tuple_source = TupleOutputRef( -# source_node_name=resolved_source.name, -# output_index=node.tuple_source.output_index, -# ) - -# for node in final_nodes: -# for parent in node.parents: -# append_unique(parent.children, node) - -# new_source_nodes = [] - -# for source_node in graph.source_nodes: -# resolved_source = resolve_node(source_node) - -# if id(resolved_source) in final_node_ids: -# append_unique(new_source_nodes, resolved_source) - -# graph.name_map = new_name_map -# graph.source_nodes = new_source_nodes -# graph.consumed_ids.clear() -# graph.fusedNodes.clear() \ No newline at end of file diff --git a/python/cactus/transpiler/ir/simplify_ir.py b/python/cactus/transpiler/ir/simplify_ir.py deleted file mode 100644 index b77f1f04a..000000000 --- a/python/cactus/transpiler/ir/simplify_ir.py +++ /dev/null @@ -1,15 +0,0 @@ -""" -The code in this directory is responsible for taking in the raw components lists from the converter folder and transform -it into a valid computation DAG (consisting of both operation nodes and value nodes). Topological sort will then be performed -on the DAG to generate an in-order execution list. -""" - -import models -import converter.models as CVModels - - -def simplify(messy_ir:CVModels.LayerMap) -> CVModels.LayerMap: - pass - - - diff --git a/python/cactus/transpiler/main.py b/python/cactus/transpiler/main.py index e12744367..e69de29bb 100644 --- a/python/cactus/transpiler/main.py +++ b/python/cactus/transpiler/main.py @@ -1,7 +0,0 @@ -import converter.convert as cv -import converter.models as CVModels - -#TODO: Update converter to also downlaod model config.txt -if __name__ == "__main__": - model_map:CVModels = cv.model_ops_to_json("LiquidAI/LFM2.5-8B-A1B", "/Users/sandhup/Documents/personal/cactus/python/cactus/transpiler/converter/jsons/") - \ No newline at end of file From 802b326acffb0e8c6c506dde6718fafe8eb85ba4 Mon Sep 17 00:00:00 2001 From: Parkirat Date: Sun, 12 Jul 2026 11:25:18 -0400 Subject: [PATCH 19/84] File restructure --- python/cactus/transpiler/ModelProfiles/models.py | 0 1 file changed, 0 insertions(+), 0 deletions(-) create mode 100644 python/cactus/transpiler/ModelProfiles/models.py diff --git a/python/cactus/transpiler/ModelProfiles/models.py b/python/cactus/transpiler/ModelProfiles/models.py new file mode 100644 index 000000000..e69de29bb From 386ff84c3084ffebda87c14bb123d2937ad05c51 Mon Sep 17 00:00:00 2001 From: Parkirat Date: Sun, 12 Jul 2026 11:29:22 -0400 Subject: [PATCH 20/84] Folder restructuring --- python/cactus/transpiler/Converter/models.py | 0 python/cactus/transpiler/Fusions/models.py | 0 python/cactus/transpiler/Generator/models.py | 0 python/cactus/transpiler/IR/models.py | 0 python/cactus/transpiler/main.py | 7 +++++++ 5 files changed, 7 insertions(+) create mode 100644 python/cactus/transpiler/Converter/models.py create mode 100644 python/cactus/transpiler/Fusions/models.py create mode 100644 python/cactus/transpiler/Generator/models.py create mode 100644 python/cactus/transpiler/IR/models.py diff --git a/python/cactus/transpiler/Converter/models.py b/python/cactus/transpiler/Converter/models.py new file mode 100644 index 000000000..e69de29bb diff --git a/python/cactus/transpiler/Fusions/models.py b/python/cactus/transpiler/Fusions/models.py new file mode 100644 index 000000000..e69de29bb diff --git a/python/cactus/transpiler/Generator/models.py b/python/cactus/transpiler/Generator/models.py new file mode 100644 index 000000000..e69de29bb diff --git a/python/cactus/transpiler/IR/models.py b/python/cactus/transpiler/IR/models.py new file mode 100644 index 000000000..e69de29bb diff --git a/python/cactus/transpiler/main.py b/python/cactus/transpiler/main.py index e69de29bb..b36194922 100644 --- a/python/cactus/transpiler/main.py +++ b/python/cactus/transpiler/main.py @@ -0,0 +1,7 @@ +from dataclasses import dataclass + +@dataclass(slots=True) +class ModelProfile: + model_family:str + component_graphs:tuple[str] + export_modes:ep \ No newline at end of file From a74d3fab45d23d6558e21538a804a5960260ebf8 Mon Sep 17 00:00:00 2001 From: Parkirat Date: Mon, 13 Jul 2026 16:32:16 -0400 Subject: [PATCH 21/84] Model profiles --- .../transpiler/ModelProfiles/components.py | 48 ++++++++ .../cactus/transpiler/ModelProfiles/models.py | 41 +++++++ .../transpiler/ModelProfiles/profiles.py | 50 ++++++++ .../cactus/transpiler/ModelProfiles/routes.py | 116 ++++++++++++++++++ 4 files changed, 255 insertions(+) create mode 100644 python/cactus/transpiler/ModelProfiles/components.py create mode 100644 python/cactus/transpiler/ModelProfiles/profiles.py create mode 100644 python/cactus/transpiler/ModelProfiles/routes.py diff --git a/python/cactus/transpiler/ModelProfiles/components.py b/python/cactus/transpiler/ModelProfiles/components.py new file mode 100644 index 000000000..08280c300 --- /dev/null +++ b/python/cactus/transpiler/ModelProfiles/components.py @@ -0,0 +1,48 @@ +from __future__ import annotations + +from .models import Components + + +TEXT_EMBED = Components( + name="text_embed", + patterns=("embed_tokens", "embedding"), +) + +VISION_ENCODER = Components( + name="vision_encoder", + patterns=("vision", "image"), +) + +AUDIO_ENCODER = Components( + name="audio_encoder", + patterns=("audio",), +) + +TOKEN_MERGE = Components( + name="token_merge", + patterns=("multi_modal", "multimodal", "projector", "merge"), +) + +TEXT_DECODER = Components( + name="text_decoder", + patterns=("language_model", "text", "decoder", "model.layers"), +) + +LM_HEAD = Components( + name="lm_head", + patterns=("lm_head", "logits"), +) + + +TEXT_COMPONENTS = { + TEXT_EMBED.name: TEXT_EMBED, + TEXT_DECODER.name: TEXT_DECODER, + LM_HEAD.name: LM_HEAD, +} + +MULTIMODAL_COMPONENTS = { + **TEXT_COMPONENTS, + VISION_ENCODER.name: VISION_ENCODER, + AUDIO_ENCODER.name: AUDIO_ENCODER, + TOKEN_MERGE.name: TOKEN_MERGE, +} diff --git a/python/cactus/transpiler/ModelProfiles/models.py b/python/cactus/transpiler/ModelProfiles/models.py index e69de29bb..82e1980f3 100644 --- a/python/cactus/transpiler/ModelProfiles/models.py +++ b/python/cactus/transpiler/ModelProfiles/models.py @@ -0,0 +1,41 @@ +from dataclasses import dataclass +from typing import Optional + + +@dataclass(slots=True) +class Components: + name:str + patterns:tuple[str, ...] + +@dataclass(slots=True) +class Combinations: + input:tuple[str, ...] + output:str + +@dataclass(slots=True) +class Route: + route:tuple[Combinations, ...] + +@dataclass(slots=True) +class InferencePattern: + name:str + route:Route + +@dataclass(slots=True) +class Cache: + types:tuple[str, ...] + +@dataclass(slots=True) +class Files: + required: tuple[str, ...] + optional: tuple[str, ...] + +@dataclass(slots=True) +class ModelProfile: + model_profiles:str + components: dict[str, Components] + inference_type: dict[str, InferencePattern] + cache_type: Cache + files: Files + fusion_fields: tuple[str, ...] #(gemma, generic) + features: tuple[str, ...] #String of what major ops (attention, RMSNorm, ) diff --git a/python/cactus/transpiler/ModelProfiles/profiles.py b/python/cactus/transpiler/ModelProfiles/profiles.py new file mode 100644 index 000000000..bd736ca9f --- /dev/null +++ b/python/cactus/transpiler/ModelProfiles/profiles.py @@ -0,0 +1,50 @@ +from __future__ import annotations + +from .components import MULTIMODAL_COMPONENTS +from .models import Cache, Files, ModelProfile +from .routes import GEMMA4_INFERENCE_PATTERNS + + +GEMMA4_E2B_PROFILE = ModelProfile( + model_profiles="gemma4_e2b", + components=MULTIMODAL_COMPONENTS, + inference_type=GEMMA4_INFERENCE_PATTERNS, + cache_type=Cache( + types=("attention_kv", "sliding_window_attention_kv"), + ), + files=Files( + required=("config.json",), + optional=( + "generation_config.json", + "processor_config.json", + "tokenizer_config.json", + ), + ), + fusion_fields=( + "generic", + "embedding", + "vision", + "audio", + "gemma4_token_merge", + "gemma4_rmsnorm", + "gemma4_rope", + "gemma4_attention", + "gemma4_mlp", + "linear", + ), + features=( + "text", + "vision", + "audio", + "attention", + "sliding_window_attention", + "rope", + "rms_norm", + "conv", + "kv_cache", + "multimodal_token_merge", + ), +) + + +ALL_PROFILES = (GEMMA4_E2B_PROFILE,) diff --git a/python/cactus/transpiler/ModelProfiles/routes.py b/python/cactus/transpiler/ModelProfiles/routes.py new file mode 100644 index 000000000..e556b83c8 --- /dev/null +++ b/python/cactus/transpiler/ModelProfiles/routes.py @@ -0,0 +1,116 @@ +from __future__ import annotations + +from .models import Combinations, InferencePattern, Route + + +TEXT_PREFILL_ROUTE = Route( + route=( + Combinations( + input=("text_embed",), + output="text_decoder", + ), + Combinations( + input=("text_decoder",), + output="lm_head", + ), + ) +) + +TEXT_VISION_PREFILL_ROUTE = Route( + route=( + Combinations( + input=("text_embed", "vision_encoder"), + output="token_merge", + ), + Combinations( + input=("token_merge",), + output="text_decoder", + ), + Combinations( + input=("text_decoder",), + output="lm_head", + ), + ) +) + +TEXT_AUDIO_PREFILL_ROUTE = Route( + route=( + Combinations( + input=("text_embed", "audio_encoder"), + output="token_merge", + ), + Combinations( + input=("token_merge",), + output="text_decoder", + ), + Combinations( + input=("text_decoder",), + output="lm_head", + ), + ) +) + +TEXT_VISION_AUDIO_PREFILL_ROUTE = Route( + route=( + Combinations( + input=("text_embed", "vision_encoder", "audio_encoder"), + output="token_merge", + ), + Combinations( + input=("token_merge",), + output="text_decoder", + ), + Combinations( + input=("text_decoder",), + output="lm_head", + ), + ) +) + +TEXT_DECODE_ROUTE = Route( + route=( + Combinations( + input=("text_embed",), + output="text_decoder", + ), + Combinations( + input=("text_decoder",), + output="lm_head", + ), + ) +) + + +TEXT_PREFILL = InferencePattern( + name="text_prefill", + route=TEXT_PREFILL_ROUTE, +) + +TEXT_VISION_PREFILL = InferencePattern( + name="text_vision_prefill", + route=TEXT_VISION_PREFILL_ROUTE, +) + +TEXT_AUDIO_PREFILL = InferencePattern( + name="text_audio_prefill", + route=TEXT_AUDIO_PREFILL_ROUTE, +) + +TEXT_VISION_AUDIO_PREFILL = InferencePattern( + name="text_vision_audio_prefill", + route=TEXT_VISION_AUDIO_PREFILL_ROUTE, +) + +TEXT_DECODE = InferencePattern( + name="text_decode", + route=TEXT_DECODE_ROUTE, +) + + +GEMMA4_INFERENCE_PATTERNS = { + TEXT_PREFILL.name: TEXT_PREFILL, + TEXT_VISION_PREFILL.name: TEXT_VISION_PREFILL, + TEXT_AUDIO_PREFILL.name: TEXT_AUDIO_PREFILL, + TEXT_VISION_AUDIO_PREFILL.name: TEXT_VISION_AUDIO_PREFILL, + TEXT_DECODE.name: TEXT_DECODE, +} From 4d0e6a6b78774c84c5fe6c9f8d45b9df3252e542 Mon Sep 17 00:00:00 2001 From: Parkirat Date: Mon, 13 Jul 2026 16:41:17 -0400 Subject: [PATCH 22/84] Added additional profiles --- .../transpiler/ModelProfiles/components.py | 29 +++++ .../transpiler/ModelProfiles/profiles.py | 117 +++++++++++++++++- .../cactus/transpiler/ModelProfiles/routes.py | 69 +++++++++++ 3 files changed, 212 insertions(+), 3 deletions(-) diff --git a/python/cactus/transpiler/ModelProfiles/components.py b/python/cactus/transpiler/ModelProfiles/components.py index 08280c300..3dc551eae 100644 --- a/python/cactus/transpiler/ModelProfiles/components.py +++ b/python/cactus/transpiler/ModelProfiles/components.py @@ -18,6 +18,16 @@ patterns=("audio",), ) +ASR_DECODER = Components( + name="asr_decoder", + patterns=("decoder", "tdt_decoder", "rnnt_decoder", "prediction"), +) + +ASR_HEAD = Components( + name="asr_head", + patterns=("joint", "tdt_head", "rnnt_head", "ctc", "log_probs"), +) + TOKEN_MERGE = Components( name="token_merge", patterns=("multi_modal", "multimodal", "projector", "merge"), @@ -40,6 +50,25 @@ LM_HEAD.name: LM_HEAD, } +SPEECH_SEQ2SEQ_COMPONENTS = { + TEXT_EMBED.name: TEXT_EMBED, + AUDIO_ENCODER.name: AUDIO_ENCODER, + TEXT_DECODER.name: TEXT_DECODER, + LM_HEAD.name: LM_HEAD, +} + +SPEECH_TRANSCRIBER_COMPONENTS = { + AUDIO_ENCODER.name: AUDIO_ENCODER, + ASR_DECODER.name: ASR_DECODER, + ASR_HEAD.name: ASR_HEAD, +} + +VISION_LANGUAGE_COMPONENTS = { + **TEXT_COMPONENTS, + VISION_ENCODER.name: VISION_ENCODER, + TOKEN_MERGE.name: TOKEN_MERGE, +} + MULTIMODAL_COMPONENTS = { **TEXT_COMPONENTS, VISION_ENCODER.name: VISION_ENCODER, diff --git a/python/cactus/transpiler/ModelProfiles/profiles.py b/python/cactus/transpiler/ModelProfiles/profiles.py index bd736ca9f..b924c6690 100644 --- a/python/cactus/transpiler/ModelProfiles/profiles.py +++ b/python/cactus/transpiler/ModelProfiles/profiles.py @@ -1,8 +1,13 @@ from __future__ import annotations -from .components import MULTIMODAL_COMPONENTS +from .components import MULTIMODAL_COMPONENTS, SPEECH_SEQ2SEQ_COMPONENTS, SPEECH_TRANSCRIBER_COMPONENTS, VISION_LANGUAGE_COMPONENTS from .models import Cache, Files, ModelProfile -from .routes import GEMMA4_INFERENCE_PATTERNS +from .routes import ( + GEMMA4_INFERENCE_PATTERNS, + LFM_VLM_INFERENCE_PATTERNS, + PARAKEET_INFERENCE_PATTERNS, + WHISPER_INFERENCE_PATTERNS, +) GEMMA4_E2B_PROFILE = ModelProfile( @@ -46,5 +51,111 @@ ), ) +WHISPER_PROFILE = ModelProfile( + model_profiles="whisper", + components=SPEECH_SEQ2SEQ_COMPONENTS, + inference_type=WHISPER_INFERENCE_PATTERNS, + cache_type=Cache( + types=("cross_attention_kv", "decoder_attention_kv"), + ), + files=Files( + required=("config.json",), + optional=( + "generation_config.json", + "preprocessor_config.json", + "tokenizer_config.json", + ), + ), + fusion_fields=( + "generic", + "audio", + "whisper_audio_encoder", + "whisper_attention", + "whisper_mlp", + "linear", + ), + features=( + "audio", + "speech_to_text", + "encoder_decoder", + "attention", + "cross_attention", + "conv", + "kv_cache", + ), +) -ALL_PROFILES = (GEMMA4_E2B_PROFILE,) +PARAKEET_PROFILE = ModelProfile( + model_profiles="parakeet", + components=SPEECH_TRANSCRIBER_COMPONENTS, + inference_type=PARAKEET_INFERENCE_PATTERNS, + cache_type=Cache( + types=(), + ), + files=Files( + required=("config.json",), + optional=( + "preprocessor_config.json", + "tokenizer_config.json", + ), + ), + fusion_fields=( + "generic", + "audio", + "parakeet_fastconformer", + "parakeet_tdt", + "conv", + "linear", + ), + features=( + "audio", + "speech_to_text", + "conformer", + "tdt", + "conv", + "attention", + ), +) + +LFM_VLM_PROFILE = ModelProfile( + model_profiles="lfm_vlm", + components=VISION_LANGUAGE_COMPONENTS, + inference_type=LFM_VLM_INFERENCE_PATTERNS, + cache_type=Cache( + types=("attention_kv",), + ), + files=Files( + required=("config.json",), + optional=( + "generation_config.json", + "processor_config.json", + "tokenizer_config.json", + ), + ), + fusion_fields=( + "generic", + "embedding", + "vision", + "lfm_vlm_token_merge", + "lfm_attention", + "lfm_mlp", + "linear", + ), + features=( + "text", + "vision", + "vision_language", + "attention", + "rope", + "rms_norm", + "kv_cache", + "multimodal_token_merge", + ), +) + +ALL_PROFILES = ( + GEMMA4_E2B_PROFILE, + WHISPER_PROFILE, + PARAKEET_PROFILE, + LFM_VLM_PROFILE, +) diff --git a/python/cactus/transpiler/ModelProfiles/routes.py b/python/cactus/transpiler/ModelProfiles/routes.py index e556b83c8..071573ada 100644 --- a/python/cactus/transpiler/ModelProfiles/routes.py +++ b/python/cactus/transpiler/ModelProfiles/routes.py @@ -80,6 +80,45 @@ ) ) +SPEECH_SEQ2SEQ_PREFILL_ROUTE = Route( + route=( + Combinations( + input=("audio_encoder", "text_embed"), + output="text_decoder", + ), + Combinations( + input=("text_decoder",), + output="lm_head", + ), + ) +) + +SPEECH_SEQ2SEQ_DECODE_ROUTE = Route( + route=( + Combinations( + input=("audio_encoder", "text_embed"), + output="text_decoder", + ), + Combinations( + input=("text_decoder",), + output="lm_head", + ), + ) +) + +SPEECH_TRANSCRIBE_ROUTE = Route( + route=( + Combinations( + input=("audio_encoder",), + output="asr_decoder", + ), + Combinations( + input=("audio_encoder", "asr_decoder"), + output="asr_head", + ), + ) +) + TEXT_PREFILL = InferencePattern( name="text_prefill", @@ -106,6 +145,21 @@ route=TEXT_DECODE_ROUTE, ) +SPEECH_SEQ2SEQ_PREFILL = InferencePattern( + name="speech_seq2seq_prefill", + route=SPEECH_SEQ2SEQ_PREFILL_ROUTE, +) + +SPEECH_SEQ2SEQ_DECODE = InferencePattern( + name="speech_seq2seq_decode", + route=SPEECH_SEQ2SEQ_DECODE_ROUTE, +) + +SPEECH_TRANSCRIBE = InferencePattern( + name="speech_transcribe", + route=SPEECH_TRANSCRIBE_ROUTE, +) + GEMMA4_INFERENCE_PATTERNS = { TEXT_PREFILL.name: TEXT_PREFILL, @@ -114,3 +168,18 @@ TEXT_VISION_AUDIO_PREFILL.name: TEXT_VISION_AUDIO_PREFILL, TEXT_DECODE.name: TEXT_DECODE, } + +WHISPER_INFERENCE_PATTERNS = { + SPEECH_SEQ2SEQ_PREFILL.name: SPEECH_SEQ2SEQ_PREFILL, + SPEECH_SEQ2SEQ_DECODE.name: SPEECH_SEQ2SEQ_DECODE, +} + +PARAKEET_INFERENCE_PATTERNS = { + SPEECH_TRANSCRIBE.name: SPEECH_TRANSCRIBE, +} + +LFM_VLM_INFERENCE_PATTERNS = { + TEXT_PREFILL.name: TEXT_PREFILL, + TEXT_VISION_PREFILL.name: TEXT_VISION_PREFILL, + TEXT_DECODE.name: TEXT_DECODE, +} From d54bda56f87d33c10a08338c49804413335f7b2c Mon Sep 17 00:00:00 2001 From: Parkirat Date: Tue, 14 Jul 2026 09:55:48 -0400 Subject: [PATCH 23/84] Added qwen2.5-0.5b profile --- python/cactus/transpiler/Converter/models.py | 1 + .../transpiler/ModelProfiles/profiles.py | 47 ++++++++++++++++++- .../cactus/transpiler/ModelProfiles/routes.py | 5 ++ 3 files changed, 52 insertions(+), 1 deletion(-) diff --git a/python/cactus/transpiler/Converter/models.py b/python/cactus/transpiler/Converter/models.py index e69de29bb..141c5db23 100644 --- a/python/cactus/transpiler/Converter/models.py +++ b/python/cactus/transpiler/Converter/models.py @@ -0,0 +1 @@ +from dataclasses import dataclass \ No newline at end of file diff --git a/python/cactus/transpiler/ModelProfiles/profiles.py b/python/cactus/transpiler/ModelProfiles/profiles.py index b924c6690..e1a9b2ce0 100644 --- a/python/cactus/transpiler/ModelProfiles/profiles.py +++ b/python/cactus/transpiler/ModelProfiles/profiles.py @@ -1,11 +1,18 @@ from __future__ import annotations -from .components import MULTIMODAL_COMPONENTS, SPEECH_SEQ2SEQ_COMPONENTS, SPEECH_TRANSCRIBER_COMPONENTS, VISION_LANGUAGE_COMPONENTS +from .components import ( + MULTIMODAL_COMPONENTS, + SPEECH_SEQ2SEQ_COMPONENTS, + SPEECH_TRANSCRIBER_COMPONENTS, + TEXT_COMPONENTS, + VISION_LANGUAGE_COMPONENTS, +) from .models import Cache, Files, ModelProfile from .routes import ( GEMMA4_INFERENCE_PATTERNS, LFM_VLM_INFERENCE_PATTERNS, PARAKEET_INFERENCE_PATTERNS, + QWEN2_5_INFERENCE_PATTERNS, WHISPER_INFERENCE_PATTERNS, ) @@ -153,9 +160,47 @@ ), ) +QWEN2_5_0_5B_PROFILE = ModelProfile( + model_profiles="qwen2_5_0_5b", + components=TEXT_COMPONENTS, + inference_type=QWEN2_5_INFERENCE_PATTERNS, + cache_type=Cache( + types=("attention_kv",), + ), + files=Files( + required=("config.json",), + optional=( + "generation_config.json", + "tokenizer.json", + "tokenizer_config.json", + "vocab.json", + "merges.txt", + ), + ), + fusion_fields=( + "generic", + "embedding", + "qwen2_5_rmsnorm", + "qwen2_5_rope", + "qwen2_5_attention", + "qwen2_5_mlp", + "linear", + ), + features=( + "text", + "causal_lm", + "attention", + "rope", + "rms_norm", + "mlp", + "kv_cache", + ), +) + ALL_PROFILES = ( GEMMA4_E2B_PROFILE, WHISPER_PROFILE, PARAKEET_PROFILE, LFM_VLM_PROFILE, + QWEN2_5_0_5B_PROFILE, ) diff --git a/python/cactus/transpiler/ModelProfiles/routes.py b/python/cactus/transpiler/ModelProfiles/routes.py index 071573ada..fb59970e5 100644 --- a/python/cactus/transpiler/ModelProfiles/routes.py +++ b/python/cactus/transpiler/ModelProfiles/routes.py @@ -183,3 +183,8 @@ TEXT_VISION_PREFILL.name: TEXT_VISION_PREFILL, TEXT_DECODE.name: TEXT_DECODE, } + +QWEN2_5_INFERENCE_PATTERNS = { + TEXT_PREFILL.name: TEXT_PREFILL, + TEXT_DECODE.name: TEXT_DECODE, +} From 7fd1cfc3c1d0e280df08fcae69d2ee893856fb6b Mon Sep 17 00:00:00 2001 From: Parkirat Date: Tue, 14 Jul 2026 15:53:27 -0400 Subject: [PATCH 24/84] Added converter logic --- python/cactus/transpiler/Converter/models.py | 170 +++++++++++++++++- .../cactus/transpiler/ModelProfiles/models.py | 4 +- 2 files changed, 171 insertions(+), 3 deletions(-) diff --git a/python/cactus/transpiler/Converter/models.py b/python/cactus/transpiler/Converter/models.py index 141c5db23..55f0ebade 100644 --- a/python/cactus/transpiler/Converter/models.py +++ b/python/cactus/transpiler/Converter/models.py @@ -1 +1,169 @@ -from dataclasses import dataclass \ No newline at end of file +from __future__ import annotations + +import json +from dataclasses import dataclass, asdict +from typing import Any +from pydantic import BaseModel +import torch + +@dataclass(slots=True) +class Input: + args: tuple + kwargs: dict + modalities: tuple[str, ...] + inference_mode: str + +@dataclass(slots=True) +class Model: + name: str + input: Input + + def export(self, input: Input) -> LayerMap: + export_(model=self, input=input) + + +@dataclass(slots=True) +class TensorInstance: + shape: list[int] + dtype: str + + @classmethod + def from_tensor(cls, x: torch.Tensor) -> "TensorInstance": + return cls(shape=list(x.shape), dtype=str(x.dtype)) + +@dataclass(slots=True) +class Slice: + start: Any + stop: Any + step: Any + + @classmethod + def from_slice(cls, x: slice) -> "Slice": + return cls(start=jsonable(x.start), stop=jsonable(x.stop), step=jsonable(x.step)) + + +@dataclass(slots=True) +class LayerRecord(BaseModel): + index: int + name: str + node_type: str + target: str + args: Any + kwargs: Any + users: list[str] + tensor_output_meta: Any | None + module_stack: Any | None + + @classmethod + def from_node(cls, num: int, x: torch.fx.Node) -> "LayerRecord": + return cls( + index=num, + name=str(x.name), + node_type=str(x.op), + target=str(x.target), + args=jsonable(x.args), + kwargs=jsonable(x.kwargs), + users=[user.name for user in x.users], + tensor_output_meta=extract_tensor_meta(x), + module_stack=extract_module_stack(x), + ) + + + +@dataclass(slots=True) +class LayerMap(BaseModel): + model_name: str + task: str + graph_signature: str + range_constants: str + nodes: list[LayerRecord] + + @classmethod + def from_data(cls, x: torch.export.ExportedProgram, name: str, model_task: str, nodes_list: list[LayerRecord]) -> "LayerMap": + return cls(model_name=name, task=model_task, graph_signature=repr(x.graph_signature), range_constants=repr(x.range_constraints), nodes=nodes_list) + + def model_dump(self) -> dict[str, Any]: + return asdict(self) + + def model_dump_json(self, indent: int | None = None) -> str: + return json.dumps(self.model_dump(), indent=indent) + +"""#####################################Model Utils#####################################""" + +def jsonable(x: Any) -> Any: + if isinstance(x, torch.fx.Node): + return {"node": x.name} + + if isinstance(x, torch.Tensor): + return asdict(TensorInstance.from_tensor(x)) + + if isinstance(x, torch.Size): + return list(x) + + if isinstance(x, torch.dtype): + return str(x) + + if isinstance(x, torch.device): + return str(x) + + if isinstance(x, slice): + return asdict(Slice.from_slice(x)) + + if isinstance(x, range): + return list(x) + + if isinstance(x, (list, tuple)): + return [jsonable(v) for v in x] + + if isinstance(x, dict): + return {str(k): jsonable(v) for k, v in x.items()} + + if isinstance(x, (str, int, float, bool)) or x is None: + return x + + return repr(x) + + +def aten_name(target: Any) -> str: + schema = getattr(target, "_schema", None) + + if schema is not None and "::" in schema.name: + namespace, op = schema.name.split("::", 1) + overload = schema.overload_name if schema.overload_name else "default" + return f"{namespace}.{op}.{overload}" + + if hasattr(target, "name"): + return target.name + + return str(target) + + +def extract_tensor_meta(node: torch.fx.Node) -> Any | None: + if "val" not in node.meta: + return None + + return jsonable(node.meta["val"]) + + +def extract_module_stack(node: torch.fx.Node) -> Any | None: + stack = node.meta.get("nn_module_stack", None) + if stack is None: + return None + + out = [] + + for key, value in stack.items(): + if isinstance(value, tuple) and len(value) >= 2: + module_path = value[0] + module_type = value[1] + out.append({"key": str(key), "module_path": str(module_path), "module_type": getattr(module_type, "__name__", str(module_type))}) + else: + out.append({"key": str(key), "value": repr(value)}) + + return out + + +def export_(model: Model, input: Input) -> LayerMap: + exported = torch.export.export(model.model, args=input.args, kwargs=input.kwargs, strict=False,).run_decompositions() + records = [LayerRecord.from_node(i, node) for i, node in enumerate(exported.graph.nodes)] + return LayerMap.from_exported_model(model=model, input=input, exported_model=exported, nodes=records) diff --git a/python/cactus/transpiler/ModelProfiles/models.py b/python/cactus/transpiler/ModelProfiles/models.py index 82e1980f3..22fd30559 100644 --- a/python/cactus/transpiler/ModelProfiles/models.py +++ b/python/cactus/transpiler/ModelProfiles/models.py @@ -37,5 +37,5 @@ class ModelProfile: inference_type: dict[str, InferencePattern] cache_type: Cache files: Files - fusion_fields: tuple[str, ...] #(gemma, generic) - features: tuple[str, ...] #String of what major ops (attention, RMSNorm, ) + fusion_fields: tuple[str, ...] + features: tuple[str, ...] From 3333bb1fd75dd5a4c29fd20b9351b97852562bc8 Mon Sep 17 00:00:00 2001 From: Parkirat Date: Fri, 17 Jul 2026 11:15:09 -0400 Subject: [PATCH 25/84] Removed redundant objects and refactored class structure --- .../transpiler/ModelProfiles/combinations.py | 12 ++ .../transpiler/ModelProfiles/components.py | 47 +---- .../cactus/transpiler/ModelProfiles/models.py | 12 +- .../transpiler/ModelProfiles/profiles.py | 24 +-- .../cactus/transpiler/ModelProfiles/routes.py | 185 +++--------------- 5 files changed, 52 insertions(+), 228 deletions(-) create mode 100644 python/cactus/transpiler/ModelProfiles/combinations.py diff --git a/python/cactus/transpiler/ModelProfiles/combinations.py b/python/cactus/transpiler/ModelProfiles/combinations.py new file mode 100644 index 000000000..fa53ea279 --- /dev/null +++ b/python/cactus/transpiler/ModelProfiles/combinations.py @@ -0,0 +1,12 @@ +from .models import Combinations + +# The input and output strings are component names from components.py. +TEXT_EMBED_TO_DECODER = Combinations(input=("text_embed"), output="text_decoder") +DECODER_TO_LM_HEAD = Combinations(input=("text_decoder"), output="lm_head") +TEXT_VISION_TO_TOKEN_MERGE = Combinations(input=("text_embed", "vision_encoder"), output="token_merge") +TEXT_AUDIO_TO_TOKEN_MERGE = Combinations(input=("text_embed", "audio_encoder"), output="token_merge") +TEXT_VISION_AUDIO_TO_TOKEN_MERGE = Combinations(input=("text_embed", "vision_encoder", "audio_encoder"), output="token_merge") +TOKEN_MERGE_TO_DECODER = Combinations(input=("token_merge"), output="text_decoder") +AUDIO_TEXT_TO_DECODER = Combinations(input=("audio_encoder", "text_embed"), output="text_decoder") +AUDIO_ENCODER_TO_ASR_DECODER = Combinations(input=("audio_encoder"), output="asr_decoder") +AUDIO_ASR_DECODER_TO_ASR_HEAD = Combinations(input=("audio_encoder", "asr_decoder"), output="asr_head") diff --git a/python/cactus/transpiler/ModelProfiles/components.py b/python/cactus/transpiler/ModelProfiles/components.py index 3dc551eae..72fd7dc5c 100644 --- a/python/cactus/transpiler/ModelProfiles/components.py +++ b/python/cactus/transpiler/ModelProfiles/components.py @@ -3,45 +3,14 @@ from .models import Components -TEXT_EMBED = Components( - name="text_embed", - patterns=("embed_tokens", "embedding"), -) - -VISION_ENCODER = Components( - name="vision_encoder", - patterns=("vision", "image"), -) - -AUDIO_ENCODER = Components( - name="audio_encoder", - patterns=("audio",), -) - -ASR_DECODER = Components( - name="asr_decoder", - patterns=("decoder", "tdt_decoder", "rnnt_decoder", "prediction"), -) - -ASR_HEAD = Components( - name="asr_head", - patterns=("joint", "tdt_head", "rnnt_head", "ctc", "log_probs"), -) - -TOKEN_MERGE = Components( - name="token_merge", - patterns=("multi_modal", "multimodal", "projector", "merge"), -) - -TEXT_DECODER = Components( - name="text_decoder", - patterns=("language_model", "text", "decoder", "model.layers"), -) - -LM_HEAD = Components( - name="lm_head", - patterns=("lm_head", "logits"), -) +TEXT_EMBED = Components(name="text_embed", patterns=("embed_tokens", "embedding")) +VISION_ENCODER = Components(name="vision_encoder", patterns=("vision", "image")) +AUDIO_ENCODER = Components(name="audio_encoder", patterns=("audio",)) +ASR_DECODER = Components(name="asr_decoder", patterns=("decoder", "tdt_decoder", "rnnt_decoder", "prediction")) +ASR_HEAD = Components(name="asr_head", patterns=("joint", "tdt_head", "rnnt_head", "ctc", "log_probs")) +TOKEN_MERGE = Components(name="token_merge", patterns=("multi_modal", "multimodal", "projector", "merge")) +TEXT_DECODER = Components(name="text_decoder", patterns=("language_model", "text", "decoder", "model.layers")) +LM_HEAD = Components(name="lm_head", patterns=("lm_head", "logits"),) TEXT_COMPONENTS = { diff --git a/python/cactus/transpiler/ModelProfiles/models.py b/python/cactus/transpiler/ModelProfiles/models.py index 22fd30559..210d6cb21 100644 --- a/python/cactus/transpiler/ModelProfiles/models.py +++ b/python/cactus/transpiler/ModelProfiles/models.py @@ -12,18 +12,10 @@ class Combinations: input:tuple[str, ...] output:str -@dataclass(slots=True) -class Route: - route:tuple[Combinations, ...] - @dataclass(slots=True) class InferencePattern: name:str - route:Route - -@dataclass(slots=True) -class Cache: - types:tuple[str, ...] + route:tuple[Combinations, ...] @dataclass(slots=True) class Files: @@ -35,7 +27,7 @@ class ModelProfile: model_profiles:str components: dict[str, Components] inference_type: dict[str, InferencePattern] - cache_type: Cache + cache_type: tuple[str, ...] files: Files fusion_fields: tuple[str, ...] features: tuple[str, ...] diff --git a/python/cactus/transpiler/ModelProfiles/profiles.py b/python/cactus/transpiler/ModelProfiles/profiles.py index e1a9b2ce0..17682f36c 100644 --- a/python/cactus/transpiler/ModelProfiles/profiles.py +++ b/python/cactus/transpiler/ModelProfiles/profiles.py @@ -1,5 +1,3 @@ -from __future__ import annotations - from .components import ( MULTIMODAL_COMPONENTS, SPEECH_SEQ2SEQ_COMPONENTS, @@ -7,7 +5,7 @@ TEXT_COMPONENTS, VISION_LANGUAGE_COMPONENTS, ) -from .models import Cache, Files, ModelProfile +from .models import Files, ModelProfile from .routes import ( GEMMA4_INFERENCE_PATTERNS, LFM_VLM_INFERENCE_PATTERNS, @@ -21,9 +19,7 @@ model_profiles="gemma4_e2b", components=MULTIMODAL_COMPONENTS, inference_type=GEMMA4_INFERENCE_PATTERNS, - cache_type=Cache( - types=("attention_kv", "sliding_window_attention_kv"), - ), + cache_type=("attention_kv", "sliding_window_attention_kv"), files=Files( required=("config.json",), optional=( @@ -62,9 +58,7 @@ model_profiles="whisper", components=SPEECH_SEQ2SEQ_COMPONENTS, inference_type=WHISPER_INFERENCE_PATTERNS, - cache_type=Cache( - types=("cross_attention_kv", "decoder_attention_kv"), - ), + cache_type=("cross_attention_kv", "decoder_attention_kv"), files=Files( required=("config.json",), optional=( @@ -96,9 +90,7 @@ model_profiles="parakeet", components=SPEECH_TRANSCRIBER_COMPONENTS, inference_type=PARAKEET_INFERENCE_PATTERNS, - cache_type=Cache( - types=(), - ), + cache_type=(), files=Files( required=("config.json",), optional=( @@ -128,9 +120,7 @@ model_profiles="lfm_vlm", components=VISION_LANGUAGE_COMPONENTS, inference_type=LFM_VLM_INFERENCE_PATTERNS, - cache_type=Cache( - types=("attention_kv",), - ), + cache_type=("attention_kv",), files=Files( required=("config.json",), optional=( @@ -164,9 +154,7 @@ model_profiles="qwen2_5_0_5b", components=TEXT_COMPONENTS, inference_type=QWEN2_5_INFERENCE_PATTERNS, - cache_type=Cache( - types=("attention_kv",), - ), + cache_type=("attention_kv",), files=Files( required=("config.json",), optional=( diff --git a/python/cactus/transpiler/ModelProfiles/routes.py b/python/cactus/transpiler/ModelProfiles/routes.py index fb59970e5..355f4b352 100644 --- a/python/cactus/transpiler/ModelProfiles/routes.py +++ b/python/cactus/transpiler/ModelProfiles/routes.py @@ -1,164 +1,27 @@ -from __future__ import annotations - -from .models import Combinations, InferencePattern, Route - - -TEXT_PREFILL_ROUTE = Route( - route=( - Combinations( - input=("text_embed",), - output="text_decoder", - ), - Combinations( - input=("text_decoder",), - output="lm_head", - ), - ) -) - -TEXT_VISION_PREFILL_ROUTE = Route( - route=( - Combinations( - input=("text_embed", "vision_encoder"), - output="token_merge", - ), - Combinations( - input=("token_merge",), - output="text_decoder", - ), - Combinations( - input=("text_decoder",), - output="lm_head", - ), - ) -) - -TEXT_AUDIO_PREFILL_ROUTE = Route( - route=( - Combinations( - input=("text_embed", "audio_encoder"), - output="token_merge", - ), - Combinations( - input=("token_merge",), - output="text_decoder", - ), - Combinations( - input=("text_decoder",), - output="lm_head", - ), - ) -) - -TEXT_VISION_AUDIO_PREFILL_ROUTE = Route( - route=( - Combinations( - input=("text_embed", "vision_encoder", "audio_encoder"), - output="token_merge", - ), - Combinations( - input=("token_merge",), - output="text_decoder", - ), - Combinations( - input=("text_decoder",), - output="lm_head", - ), - ) -) - -TEXT_DECODE_ROUTE = Route( - route=( - Combinations( - input=("text_embed",), - output="text_decoder", - ), - Combinations( - input=("text_decoder",), - output="lm_head", - ), - ) -) - -SPEECH_SEQ2SEQ_PREFILL_ROUTE = Route( - route=( - Combinations( - input=("audio_encoder", "text_embed"), - output="text_decoder", - ), - Combinations( - input=("text_decoder",), - output="lm_head", - ), - ) -) - -SPEECH_SEQ2SEQ_DECODE_ROUTE = Route( - route=( - Combinations( - input=("audio_encoder", "text_embed"), - output="text_decoder", - ), - Combinations( - input=("text_decoder",), - output="lm_head", - ), - ) -) - -SPEECH_TRANSCRIBE_ROUTE = Route( - route=( - Combinations( - input=("audio_encoder",), - output="asr_decoder", - ), - Combinations( - input=("audio_encoder", "asr_decoder"), - output="asr_head", - ), - ) -) - - -TEXT_PREFILL = InferencePattern( - name="text_prefill", - route=TEXT_PREFILL_ROUTE, -) - -TEXT_VISION_PREFILL = InferencePattern( - name="text_vision_prefill", - route=TEXT_VISION_PREFILL_ROUTE, -) - -TEXT_AUDIO_PREFILL = InferencePattern( - name="text_audio_prefill", - route=TEXT_AUDIO_PREFILL_ROUTE, -) - -TEXT_VISION_AUDIO_PREFILL = InferencePattern( - name="text_vision_audio_prefill", - route=TEXT_VISION_AUDIO_PREFILL_ROUTE, -) - -TEXT_DECODE = InferencePattern( - name="text_decode", - route=TEXT_DECODE_ROUTE, -) - -SPEECH_SEQ2SEQ_PREFILL = InferencePattern( - name="speech_seq2seq_prefill", - route=SPEECH_SEQ2SEQ_PREFILL_ROUTE, -) - -SPEECH_SEQ2SEQ_DECODE = InferencePattern( - name="speech_seq2seq_decode", - route=SPEECH_SEQ2SEQ_DECODE_ROUTE, -) - -SPEECH_TRANSCRIBE = InferencePattern( - name="speech_transcribe", - route=SPEECH_TRANSCRIBE_ROUTE, -) +from . import combinations +from .models import InferencePattern + + +TEXT_PREFILL_ROUTE = ( + combinations.TEXT_EMBED_TO_DECODER, + combinations.DECODER_TO_LM_HEAD, +) + +TEXT_VISION_PREFILL_ROUTE = (combinations.TEXT_VISION_TO_TOKEN_MERGE, combinations.TOKEN_MERGE_TO_DECODER, combinations.DECODER_TO_LM_HEAD) +TEXT_AUDIO_PREFILL_ROUTE = (combinations.TEXT_AUDIO_TO_TOKEN_MERGE, combinations.TOKEN_MERGE_TO_DECODER, combinations.DECODER_TO_LM_HEAD) +TEXT_VISION_AUDIO_PREFILL_ROUTE = (combinations.TEXT_VISION_AUDIO_TO_TOKEN_MERGE, combinations.TOKEN_MERGE_TO_DECODER, combinations.DECODER_TO_LM_HEAD) +TEXT_DECODE_ROUTE = (combinations.TEXT_EMBED_TO_DECODER, combinations.DECODER_TO_LM_HEAD) +SPEECH_SEQ2SEQ_PREFILL_ROUTE = (combinations.AUDIO_TEXT_TO_DECODER, combinations.DECODER_TO_LM_HEAD) +SPEECH_SEQ2SEQ_DECODE_ROUTE = (combinations.AUDIO_TEXT_TO_DECODER, combinations.DECODER_TO_LM_HEAD) +SPEECH_TRANSCRIBE_ROUTE = (combinations.AUDIO_ENCODER_TO_ASR_DECODER, combinations.AUDIO_ASR_DECODER_TO_ASR_HEAD) +TEXT_PREFILL = InferencePattern(name="text_prefill", route=TEXT_PREFILL_ROUTE,) +TEXT_VISION_PREFILL = InferencePattern(name="text_vision_prefill", route=TEXT_VISION_PREFILL_ROUTE) +TEXT_AUDIO_PREFILL = InferencePattern(name="text_audio_prefill", route=TEXT_AUDIO_PREFILL_ROUTE) +TEXT_VISION_AUDIO_PREFILL = InferencePattern(name="text_vision_audio_prefill", route=TEXT_VISION_AUDIO_PREFILL_ROUTE) +TEXT_DECODE = InferencePattern(name="text_decode", route=TEXT_DECODE_ROUTE,) +SPEECH_SEQ2SEQ_PREFILL = InferencePattern(name="speech_seq2seq_prefill", route=SPEECH_SEQ2SEQ_PREFILL_ROUTE,) +SPEECH_SEQ2SEQ_DECODE = InferencePattern(name="speech_seq2seq_decode", route=SPEECH_SEQ2SEQ_DECODE_ROUTE) +SPEECH_TRANSCRIBE = InferencePattern(name="speech_transcribe", route=SPEECH_TRANSCRIBE_ROUTE) GEMMA4_INFERENCE_PATTERNS = { From a6311c531084e6f337913e38e7aef8cfe68690b3 Mon Sep 17 00:00:00 2001 From: Parkirat Date: Fri, 17 Jul 2026 11:18:33 -0400 Subject: [PATCH 26/84] Removed redundant Files class + modified existing model profiles --- .../cactus/transpiler/ModelProfiles/models.py | 8 +--- .../transpiler/ModelProfiles/profiles.py | 48 +++---------------- 2 files changed, 7 insertions(+), 49 deletions(-) diff --git a/python/cactus/transpiler/ModelProfiles/models.py b/python/cactus/transpiler/ModelProfiles/models.py index 210d6cb21..ce10f4e49 100644 --- a/python/cactus/transpiler/ModelProfiles/models.py +++ b/python/cactus/transpiler/ModelProfiles/models.py @@ -1,5 +1,4 @@ from dataclasses import dataclass -from typing import Optional @dataclass(slots=True) @@ -17,17 +16,12 @@ class InferencePattern: name:str route:tuple[Combinations, ...] -@dataclass(slots=True) -class Files: - required: tuple[str, ...] - optional: tuple[str, ...] - @dataclass(slots=True) class ModelProfile: model_profiles:str components: dict[str, Components] inference_type: dict[str, InferencePattern] cache_type: tuple[str, ...] - files: Files + files: tuple[str, ...] fusion_fields: tuple[str, ...] features: tuple[str, ...] diff --git a/python/cactus/transpiler/ModelProfiles/profiles.py b/python/cactus/transpiler/ModelProfiles/profiles.py index 17682f36c..803541ce0 100644 --- a/python/cactus/transpiler/ModelProfiles/profiles.py +++ b/python/cactus/transpiler/ModelProfiles/profiles.py @@ -5,7 +5,7 @@ TEXT_COMPONENTS, VISION_LANGUAGE_COMPONENTS, ) -from .models import Files, ModelProfile +from .models import ModelProfile from .routes import ( GEMMA4_INFERENCE_PATTERNS, LFM_VLM_INFERENCE_PATTERNS, @@ -20,14 +20,7 @@ components=MULTIMODAL_COMPONENTS, inference_type=GEMMA4_INFERENCE_PATTERNS, cache_type=("attention_kv", "sliding_window_attention_kv"), - files=Files( - required=("config.json",), - optional=( - "generation_config.json", - "processor_config.json", - "tokenizer_config.json", - ), - ), + files=("config.json", "generation_config.json", "processor_config.json", "tokenizer_config.json"), fusion_fields=( "generic", "embedding", @@ -59,14 +52,7 @@ components=SPEECH_SEQ2SEQ_COMPONENTS, inference_type=WHISPER_INFERENCE_PATTERNS, cache_type=("cross_attention_kv", "decoder_attention_kv"), - files=Files( - required=("config.json",), - optional=( - "generation_config.json", - "preprocessor_config.json", - "tokenizer_config.json", - ), - ), + files=("config.json", "generation_config.json", "preprocessor_config.json", "tokenizer_config.json"), fusion_fields=( "generic", "audio", @@ -91,13 +77,7 @@ components=SPEECH_TRANSCRIBER_COMPONENTS, inference_type=PARAKEET_INFERENCE_PATTERNS, cache_type=(), - files=Files( - required=("config.json",), - optional=( - "preprocessor_config.json", - "tokenizer_config.json", - ), - ), + files=("config.json", "preprocessor_config.json", "tokenizer_config.json"), fusion_fields=( "generic", "audio", @@ -121,14 +101,7 @@ components=VISION_LANGUAGE_COMPONENTS, inference_type=LFM_VLM_INFERENCE_PATTERNS, cache_type=("attention_kv",), - files=Files( - required=("config.json",), - optional=( - "generation_config.json", - "processor_config.json", - "tokenizer_config.json", - ), - ), + files=("config.json", "generation_config.json", "processor_config.json", "tokenizer_config.json"), fusion_fields=( "generic", "embedding", @@ -155,16 +128,7 @@ components=TEXT_COMPONENTS, inference_type=QWEN2_5_INFERENCE_PATTERNS, cache_type=("attention_kv",), - files=Files( - required=("config.json",), - optional=( - "generation_config.json", - "tokenizer.json", - "tokenizer_config.json", - "vocab.json", - "merges.txt", - ), - ), + files=("config.json", "generation_config.json", "tokenizer.json", "tokenizer_config.json", "vocab.json", "merges.txt"), fusion_fields=( "generic", "embedding", From 391b6fdbad444679681bd2faf04f12c722b3d124 Mon Sep 17 00:00:00 2001 From: Parkirat Date: Fri, 17 Jul 2026 11:51:39 -0400 Subject: [PATCH 27/84] Added supported modalities to model profiles class --- python/cactus/transpiler/ModelProfiles/models.py | 1 + python/cactus/transpiler/ModelProfiles/profiles.py | 7 +++++++ 2 files changed, 8 insertions(+) diff --git a/python/cactus/transpiler/ModelProfiles/models.py b/python/cactus/transpiler/ModelProfiles/models.py index ce10f4e49..b9f26d365 100644 --- a/python/cactus/transpiler/ModelProfiles/models.py +++ b/python/cactus/transpiler/ModelProfiles/models.py @@ -25,3 +25,4 @@ class ModelProfile: files: tuple[str, ...] fusion_fields: tuple[str, ...] features: tuple[str, ...] + supported_modalties: tuple[str, ...] diff --git a/python/cactus/transpiler/ModelProfiles/profiles.py b/python/cactus/transpiler/ModelProfiles/profiles.py index 803541ce0..6e3ca8069 100644 --- a/python/cactus/transpiler/ModelProfiles/profiles.py +++ b/python/cactus/transpiler/ModelProfiles/profiles.py @@ -45,6 +45,7 @@ "kv_cache", "multimodal_token_merge", ), + supported_modalties=("audio", "vision", "text") ) WHISPER_PROFILE = ModelProfile( @@ -70,6 +71,7 @@ "conv", "kv_cache", ), + supported_modalties=("audio", "text"), ) PARAKEET_PROFILE = ModelProfile( @@ -94,6 +96,7 @@ "conv", "attention", ), + supported_modalties=("audio",), ) LFM_VLM_PROFILE = ModelProfile( @@ -121,6 +124,7 @@ "kv_cache", "multimodal_token_merge", ), + supported_modalties=("vision", "text"), ) QWEN2_5_0_5B_PROFILE = ModelProfile( @@ -147,6 +151,7 @@ "mlp", "kv_cache", ), + supported_modalties=("text",), ) ALL_PROFILES = ( @@ -156,3 +161,5 @@ LFM_VLM_PROFILE, QWEN2_5_0_5B_PROFILE, ) + +MODEL_ID_MAP = {} From 5898e7e320abf2ce92e112e2c56950ec5c24eb84 Mon Sep 17 00:00:00 2001 From: Parkirat Date: Fri, 17 Jul 2026 16:04:48 -0400 Subject: [PATCH 28/84] Progress to convert --- python/cactus/transpiler/Converter/convert.py | 10 +++ python/cactus/transpiler/Converter/inputs.py | 0 python/cactus/transpiler/Converter/models.py | 83 ++++++++++++++++++- .../cactus/transpiler/ModelProfiles/models.py | 2 +- 4 files changed, 91 insertions(+), 4 deletions(-) create mode 100644 python/cactus/transpiler/Converter/convert.py create mode 100644 python/cactus/transpiler/Converter/inputs.py diff --git a/python/cactus/transpiler/Converter/convert.py b/python/cactus/transpiler/Converter/convert.py new file mode 100644 index 000000000..3c1f1c40a --- /dev/null +++ b/python/cactus/transpiler/Converter/convert.py @@ -0,0 +1,10 @@ +from . import models +from ..ModelProfiles import profiles as MP +from ..ModelProfiles import models as MP_Models + + +def convert(model_id:str, input_modalities:tuple[str,...], custom_profile:MP_Models.ModelProfile = None) -> models.LayerMap | None: + if custom_profile is not None: + return None + + \ No newline at end of file diff --git a/python/cactus/transpiler/Converter/inputs.py b/python/cactus/transpiler/Converter/inputs.py new file mode 100644 index 000000000..e69de29bb diff --git a/python/cactus/transpiler/Converter/models.py b/python/cactus/transpiler/Converter/models.py index 55f0ebade..12220d81c 100644 --- a/python/cactus/transpiler/Converter/models.py +++ b/python/cactus/transpiler/Converter/models.py @@ -1,10 +1,16 @@ -from __future__ import annotations - import json +import os from dataclasses import dataclass, asdict +from pathlib import Path from typing import Any from pydantic import BaseModel import torch +from ..ModelProfiles import models as MP_Models + +modality_input_path = { + "vision": "/Users/sandhup/Documents/personal/cactus/python/cactus/assets/test_monkey.png", + "audio" : "/Users/sandhup/Documents/personal/cactus/python/cactus/assets/test.wav" +} @dataclass(slots=True) class Input: @@ -16,9 +22,15 @@ class Input: @dataclass(slots=True) class Model: name: str + model_profile: MP_Models.ModelProfile input: Input - def export(self, input: Input) -> LayerMap: + @classmethod + def from_profile(mp: MP_Models.ModelProfile, input_modalities: tuple[str,...], model_id: str) -> "Model": + return create_model(mp, input_modalities, model_id) + + + def export(self, input: Input) -> "LayerMap": export_(model=self, input=input) @@ -163,6 +175,71 @@ def extract_module_stack(node: torch.fx.Node) -> Any | None: return out + +def load_files(mp: MP_Models.ModelProfile, model_id: str) -> dict[str, str]: + from huggingface_hub import hf_hub_download + from huggingface_hub.errors import EntryNotFoundError, RepositoryNotFoundError + from ..ModelProfiles import profiles as MP_Profiles + + default_model_ids = { + "gemma4_e2b": "google/gemma-4-E2B", + "whisper": "openai/whisper-tiny", + "parakeet": "nvidia/parakeet-tdt-0.6b-v3", + "lfm_vlm": "LiquidAI/LFM2-VL-3B", + "qwen2_5_0_5b": "Qwen/Qwen2.5-0.5B", + } + + repo_id = ( + model_id + or MP_Profiles.MODEL_ID_MAP.get(mp.model_profiles) + or default_model_ids.get(mp.model_profiles) + or mp.model_profiles + ) + output_dir = Path(__file__).resolve().parent / "jsons" / mp.model_profiles + output_dir.mkdir(parents=True, exist_ok=True) + + token = os.environ.get("HF_TOKEN") + downloaded_files: dict[str, str] = {} + + for file in mp.files: + try: + downloaded_path = hf_hub_download( + repo_id=repo_id, + filename=file, + local_dir=output_dir, + token=token, + ) + except EntryNotFoundError: + if file == "config.json": + raise + + print(f"Skipping missing optional file for {repo_id}: {file}") + continue + except RepositoryNotFoundError: + raise ValueError(f"Could not find HuggingFace repository: {repo_id}") from None + + downloaded_files[file] = str(downloaded_path) + + return downloaded_files + + +def build_input(mp: MP_Models.ModelProfile, input_modalties: tuple[str, ...]) -> Input | None: + if not all(modality in mp.supported_modalties for modality in input_modalties): + print("Requesting unsupported modalities") + return None + + + + + + + +def create_model(mp: MP_Models.ModelProfile, input_modalities: tuple[str, ...], model_id: str) -> Model: + name_ = model_id + b + + + def export_(model: Model, input: Input) -> LayerMap: exported = torch.export.export(model.model, args=input.args, kwargs=input.kwargs, strict=False,).run_decompositions() records = [LayerRecord.from_node(i, node) for i, node in enumerate(exported.graph.nodes)] diff --git a/python/cactus/transpiler/ModelProfiles/models.py b/python/cactus/transpiler/ModelProfiles/models.py index b9f26d365..beb036ae9 100644 --- a/python/cactus/transpiler/ModelProfiles/models.py +++ b/python/cactus/transpiler/ModelProfiles/models.py @@ -25,4 +25,4 @@ class ModelProfile: files: tuple[str, ...] fusion_fields: tuple[str, ...] features: tuple[str, ...] - supported_modalties: tuple[str, ...] + supported_modalties: tuple[str, ...] \ No newline at end of file From 8d86c94fa92b9a023db7b232464ed09a5b4a1254 Mon Sep 17 00:00:00 2001 From: Parkirat Sandhu Date: Sun, 19 Jul 2026 11:53:28 -0400 Subject: [PATCH 29/84] Load files function --- python/cactus/transpiler/Converter/models.py | 27 ++++++++++++++++++++ 1 file changed, 27 insertions(+) diff --git a/python/cactus/transpiler/Converter/models.py b/python/cactus/transpiler/Converter/models.py index 12220d81c..c9a7a3341 100644 --- a/python/cactus/transpiler/Converter/models.py +++ b/python/cactus/transpiler/Converter/models.py @@ -6,6 +6,16 @@ from pydantic import BaseModel import torch from ..ModelProfiles import models as MP_Models +from ..ModelProfiles import profiles as MP_Profiles +from huggingface_hub import hf_hub_download + +default_model_ids = { + "gemma4_e2b": "google/gemma-4-E2B", + "whisper": "openai/whisper-tiny", + "parakeet": "nvidia/parakeet-tdt-0.6b-v3", + "lfm_vlm": "LiquidAI/LFM2-VL-3B", + "qwen2_5_0_5b": "Qwen/Qwen2.5-0.5B", +} modality_input_path = { "vision": "/Users/sandhup/Documents/personal/cactus/python/cactus/assets/test_monkey.png", @@ -223,6 +233,23 @@ def load_files(mp: MP_Models.ModelProfile, model_id: str) -> dict[str, str]: return downloaded_files +def load_files(mp: MP_Models.ModelProfile, model_id: str) -> dict[str, str]: + output_dir = Path(__file__).resolve().parent / "jsons" / mp.model_profiles + output_dir.mkdir(parents=True, exist_ok=True) + token = os.environ.get("HF_TOKEN") + downloaded_files: dict[str, str] = {} + + for file in mp.files: + try: + downloaded_path = hf_hub_download(repo_id=model_id, filename=file, local_dir=output_dir, token=token) + except Exception as e: + print(f"Error downloading {file} from {model_id}: {e}") + continue + + downloaded_files[file] = str(downloaded_path) + + return downloaded_files + def build_input(mp: MP_Models.ModelProfile, input_modalties: tuple[str, ...]) -> Input | None: if not all(modality in mp.supported_modalties for modality in input_modalties): print("Requesting unsupported modalities") From a0530e134e21784d1da8079a7482b924372d8521 Mon Sep 17 00:00:00 2001 From: Parkirat Sandhu Date: Sun, 19 Jul 2026 11:55:48 -0400 Subject: [PATCH 30/84] Load files function --- python/cactus/transpiler/Converter/models.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/python/cactus/transpiler/Converter/models.py b/python/cactus/transpiler/Converter/models.py index c9a7a3341..17d995abe 100644 --- a/python/cactus/transpiler/Converter/models.py +++ b/python/cactus/transpiler/Converter/models.py @@ -238,7 +238,7 @@ def load_files(mp: MP_Models.ModelProfile, model_id: str) -> dict[str, str]: output_dir.mkdir(parents=True, exist_ok=True) token = os.environ.get("HF_TOKEN") downloaded_files: dict[str, str] = {} - + for file in mp.files: try: downloaded_path = hf_hub_download(repo_id=model_id, filename=file, local_dir=output_dir, token=token) @@ -247,7 +247,6 @@ def load_files(mp: MP_Models.ModelProfile, model_id: str) -> dict[str, str]: continue downloaded_files[file] = str(downloaded_path) - return downloaded_files def build_input(mp: MP_Models.ModelProfile, input_modalties: tuple[str, ...]) -> Input | None: @@ -255,9 +254,10 @@ def build_input(mp: MP_Models.ModelProfile, input_modalties: tuple[str, ...]) -> print("Requesting unsupported modalities") return None - + return Input(args=(), kwargs={}) + From 1b6f44dcb50168cafd75200d76cf388fab257eae Mon Sep 17 00:00:00 2001 From: Parkirat Date: Mon, 20 Jul 2026 13:58:55 -0400 Subject: [PATCH 31/84] Added input loading utilities --- .../Converter/{inputs.py => constants.py} | 0 .../transpiler/Converter/input_utils.py | 134 +++++++++++++++++ python/cactus/transpiler/Converter/models.py | 136 ++++++++---------- .../cactus/transpiler/ModelProfiles/routes.py | 6 +- 4 files changed, 196 insertions(+), 80 deletions(-) rename python/cactus/transpiler/Converter/{inputs.py => constants.py} (100%) create mode 100644 python/cactus/transpiler/Converter/input_utils.py diff --git a/python/cactus/transpiler/Converter/inputs.py b/python/cactus/transpiler/Converter/constants.py similarity index 100% rename from python/cactus/transpiler/Converter/inputs.py rename to python/cactus/transpiler/Converter/constants.py diff --git a/python/cactus/transpiler/Converter/input_utils.py b/python/cactus/transpiler/Converter/input_utils.py new file mode 100644 index 000000000..2fb23566b --- /dev/null +++ b/python/cactus/transpiler/Converter/input_utils.py @@ -0,0 +1,134 @@ +from ..ModelProfiles import models as MP_Models +import torch +import os +from pathlib import Path +from huggingface_hub import hf_hub_download +import json +from typing import Any + +#Locations to inputs for +modality_input_path = { + "vision": "/Users/sandhup/Documents/personal/cactus/python/cactus/assets/test_monkey.png", + "audio" : "/Users/sandhup/Documents/personal/cactus/python/cactus/assets/test.wav" +} + + +#Downloads the file from hf +def load_files(mp: MP_Models.ModelProfile, model_id: str) -> dict[str, dict[str, Any]]: + output_dir = Path(__file__).resolve().parent / "jsons" / mp.model_profiles + output_dir.mkdir(parents=True, exist_ok=True) + token = os.environ.get("HF_TOKEN") + downloaded_files: dict[str, dict[str, Any]] = {} + + for file in mp.files: + try: + downloaded_path = hf_hub_download(repo_id=model_id, filename=file, local_dir=output_dir, token=token) + except Exception as e: + print(f"Error downloading {file} from {model_id}: {e}") + continue + + if file.endswith(".json"): + with open(downloaded_path, "r", encoding="utf-8") as f: + downloaded_files[file] = json.load(f) + + downloaded_files[file] = str(downloaded_path) + return downloaded_files + + +#TODO: Clean up all functions below this point +#Searches through configs to find desired keys, and their values +def _find_config_value(configs: dict[str, dict[str, Any]], keys: tuple[str, ...]) -> Any | None: + search_sections = ( + (), + ("text_config",), + ("vision_config",), + ("audio_config",), + ("encoder",), + ("decoder",), + ("model_config",), + ) + + for config in configs.values(): + for section_path in search_sections: + section = config + + for section_name in section_path: + if not isinstance(section, dict): + section = None + break + section = section.get(section_name) + + if not isinstance(section, dict): + continue + + for key in keys: + if key in section and section[key] is not None: + return section[key] + + return None + + +def _first_int(configs: dict[str, dict[str, Any]], keys: tuple[str, ...], default: int) -> int: + value = _find_config_value(configs, keys) + if isinstance(value, bool): + return default + if isinstance(value, int): + return value + if isinstance(value, float): + return int(value) + if isinstance(value, str) and value.isdigit(): + return int(value) + return default + + +def _build_input_ids(batch_size: int, sequence_length: int) -> torch.Tensor: + tokens = torch.arange(sequence_length, dtype=torch.long) + 1 + return tokens.unsqueeze(0).repeat(batch_size, 1) + + +def _build_attention_mask(batch_size: int, sequence_length: int) -> torch.Tensor: + return torch.ones((batch_size, sequence_length), dtype=torch.long) + + +def _build_pixel_values(batch_size: int, configs: dict[str, dict[str, Any]]) -> torch.Tensor: + size = _find_config_value(configs, ("size", "image_size")) + channels = _first_int(configs, ("num_channels", "image_num_channels"), 3) + height = 224 + width = 224 + + if isinstance(size, int): + height = size + width = size + elif isinstance(size, dict): + height = int(size.get("height") or size.get("shortest_edge") or height) + width = int(size.get("width") or size.get("shortest_edge") or width) + + return torch.zeros((batch_size, channels, height, width), dtype=torch.float32) + + +def _build_input_features(batch_size: int, configs: dict[str, dict[str, Any]]) -> torch.Tensor: + feature_bins = _first_int(configs, ("feature_size", "num_mel_bins", "n_mels"), 80) + max_source_positions = _first_int(configs, ("max_source_positions",), 1500) + feature_frames = _first_int(configs, ("nb_max_frames", "max_length", "feature_frames"), max_source_positions * 2) + return torch.zeros((batch_size, feature_bins, feature_frames), dtype=torch.float32) + + +def _build_past_key_values(configs: dict[str, dict[str, Any]], batch_size: int, past_sequence_length: int) -> tuple[tuple[torch.Tensor, torch.Tensor], ...]: + num_layers = _first_int(configs, ("num_hidden_layers", "n_layer", "num_layers", "decoder_layers"), 1) + num_attention_heads = _first_int(configs, ("num_attention_heads", "n_head", "decoder_attention_heads"), 1) + num_key_value_heads = _first_int(configs, ("num_key_value_heads", "num_kv_heads"), num_attention_heads) + hidden_size = _first_int(configs, ("hidden_size", "n_embd", "d_model"), 64) + head_dim = _first_int(configs, ("head_dim", "kv_channels"), max(1, hidden_size // max(1, num_attention_heads))) + cache_shape = (batch_size, num_key_value_heads, past_sequence_length, head_dim) + + return tuple( + ( + torch.zeros(cache_shape, dtype=torch.float32), + torch.zeros(cache_shape, dtype=torch.float32), + ) + for _ in range(num_layers) + ) + + +def _decoder_start_token_id(configs: dict[str, dict[str, Any]]) -> int: + return _first_int(configs, ("decoder_start_token_id", "bos_token_id", "pad_token_id"), 1) \ No newline at end of file diff --git a/python/cactus/transpiler/Converter/models.py b/python/cactus/transpiler/Converter/models.py index 17d995abe..368e50263 100644 --- a/python/cactus/transpiler/Converter/models.py +++ b/python/cactus/transpiler/Converter/models.py @@ -1,13 +1,12 @@ import json import os from dataclasses import dataclass, asdict -from pathlib import Path from typing import Any from pydantic import BaseModel import torch +from . import input_utils as IU from ..ModelProfiles import models as MP_Models from ..ModelProfiles import profiles as MP_Profiles -from huggingface_hub import hf_hub_download default_model_ids = { "gemma4_e2b": "google/gemma-4-E2B", @@ -17,10 +16,6 @@ "qwen2_5_0_5b": "Qwen/Qwen2.5-0.5B", } -modality_input_path = { - "vision": "/Users/sandhup/Documents/personal/cactus/python/cactus/assets/test_monkey.png", - "audio" : "/Users/sandhup/Documents/personal/cactus/python/cactus/assets/test.wav" -} @dataclass(slots=True) class Input: @@ -185,85 +180,76 @@ def extract_module_stack(node: torch.fx.Node) -> Any | None: return out - -def load_files(mp: MP_Models.ModelProfile, model_id: str) -> dict[str, str]: - from huggingface_hub import hf_hub_download - from huggingface_hub.errors import EntryNotFoundError, RepositoryNotFoundError - from ..ModelProfiles import profiles as MP_Profiles - - default_model_ids = { - "gemma4_e2b": "google/gemma-4-E2B", - "whisper": "openai/whisper-tiny", - "parakeet": "nvidia/parakeet-tdt-0.6b-v3", - "lfm_vlm": "LiquidAI/LFM2-VL-3B", - "qwen2_5_0_5b": "Qwen/Qwen2.5-0.5B", - } - - repo_id = ( - model_id - or MP_Profiles.MODEL_ID_MAP.get(mp.model_profiles) - or default_model_ids.get(mp.model_profiles) - or mp.model_profiles - ) - output_dir = Path(__file__).resolve().parent / "jsons" / mp.model_profiles - output_dir.mkdir(parents=True, exist_ok=True) - - token = os.environ.get("HF_TOKEN") - downloaded_files: dict[str, str] = {} - - for file in mp.files: - try: - downloaded_path = hf_hub_download( - repo_id=repo_id, - filename=file, - local_dir=output_dir, - token=token, - ) - except EntryNotFoundError: - if file == "config.json": - raise - - print(f"Skipping missing optional file for {repo_id}: {file}") - continue - except RepositoryNotFoundError: - raise ValueError(f"Could not find HuggingFace repository: {repo_id}") from None - - downloaded_files[file] = str(downloaded_path) - - return downloaded_files - - -def load_files(mp: MP_Models.ModelProfile, model_id: str) -> dict[str, str]: - output_dir = Path(__file__).resolve().parent / "jsons" / mp.model_profiles - output_dir.mkdir(parents=True, exist_ok=True) - token = os.environ.get("HF_TOKEN") - downloaded_files: dict[str, str] = {} - - for file in mp.files: - try: - downloaded_path = hf_hub_download(repo_id=model_id, filename=file, local_dir=output_dir, token=token) - except Exception as e: - print(f"Error downloading {file} from {model_id}: {e}") - continue - - downloaded_files[file] = str(downloaded_path) - return downloaded_files - -def build_input(mp: MP_Models.ModelProfile, input_modalties: tuple[str, ...]) -> Input | None: +#TODO: Clean up this function +def build_input( + mp: MP_Models.ModelProfile, + input_modalties: tuple[str, ...], + model_id: str | None = None, + inference_mode: str = "prefill_no_cache", +) -> Input | None: if not all(modality in mp.supported_modalties for modality in input_modalties): print("Requesting unsupported modalities") return None - - - return Input(args=(), kwargs={}) + loaded_files = IU.load_files(mp, model_id) + + configs = IU._load_json_files(loaded_files or {}) + modalities = tuple(input_modalties) + batch_size = 1 + sequence_length = 8 + past_sequence_length = 8 + kwargs: dict[str, Any] = {} + + if "text" in modalities: + if inference_mode == "decode_with_cache": + sequence_length = 1 + + kwargs["input_ids"] = IU._build_input_ids(batch_size, sequence_length) + kwargs["attention_mask"] = IU._build_attention_mask(batch_size, sequence_length) + + if "vision" in modalities: + kwargs["pixel_values"] = IU._build_pixel_values(batch_size, configs) + + if "audio" in modalities: + kwargs["input_features"] = IU._build_input_features(batch_size, configs) + + if "audio" in modalities and "text" in modalities and "vision" not in modalities and "speech" in inference_mode: + kwargs.pop("input_ids", None) + kwargs.pop("attention_mask", None) + kwargs["decoder_input_ids"] = torch.full( + (batch_size, 1), + IU._decoder_start_token_id(configs), + dtype=torch.long, + ) + + if "audio" in modalities and "text" not in modalities: + input_features = kwargs["input_features"] + kwargs["attention_mask"] = IU._build_attention_mask(batch_size, input_features.shape[-1]) + + if inference_mode in ("prefill_with_cache", "decode_with_cache"): + total_sequence_length = past_sequence_length + sequence_length + kwargs["attention_mask"] = IU._build_attention_mask(batch_size, total_sequence_length) + kwargs["past_key_values"] = IU._build_past_key_values(configs, batch_size, past_sequence_length) + kwargs["cache_position"] = torch.arange( + past_sequence_length, + past_sequence_length + sequence_length, + dtype=torch.long, + ) + kwargs["use_cache"] = True + + return Input( + args=(), + kwargs=kwargs, + modalities=modalities, + inference_mode=inference_mode, + ) def create_model(mp: MP_Models.ModelProfile, input_modalities: tuple[str, ...], model_id: str) -> Model: name_ = model_id - b + input_ = build_input(mp, input_modalities) diff --git a/python/cactus/transpiler/ModelProfiles/routes.py b/python/cactus/transpiler/ModelProfiles/routes.py index 355f4b352..d6f9dcb34 100644 --- a/python/cactus/transpiler/ModelProfiles/routes.py +++ b/python/cactus/transpiler/ModelProfiles/routes.py @@ -2,11 +2,7 @@ from .models import InferencePattern -TEXT_PREFILL_ROUTE = ( - combinations.TEXT_EMBED_TO_DECODER, - combinations.DECODER_TO_LM_HEAD, -) - +TEXT_PREFILL_ROUTE = (combinations.TEXT_EMBED_TO_DECODER, combinations.DECODER_TO_LM_HEAD) TEXT_VISION_PREFILL_ROUTE = (combinations.TEXT_VISION_TO_TOKEN_MERGE, combinations.TOKEN_MERGE_TO_DECODER, combinations.DECODER_TO_LM_HEAD) TEXT_AUDIO_PREFILL_ROUTE = (combinations.TEXT_AUDIO_TO_TOKEN_MERGE, combinations.TOKEN_MERGE_TO_DECODER, combinations.DECODER_TO_LM_HEAD) TEXT_VISION_AUDIO_PREFILL_ROUTE = (combinations.TEXT_VISION_AUDIO_TO_TOKEN_MERGE, combinations.TOKEN_MERGE_TO_DECODER, combinations.DECODER_TO_LM_HEAD) From c0ff8a0babcf234e1f2bbffd3ea02aded38554ff Mon Sep 17 00:00:00 2001 From: Parkirat Date: Mon, 20 Jul 2026 16:38:43 -0400 Subject: [PATCH 32/84] Brute-force code for converted jsons --- python/cactus/transpiler/Converter/convert.py | 83 ++- .../transpiler/Converter/input_utils.py | 185 ++++- python/cactus/transpiler/Converter/models.py | 648 +++++++++++++++++- python/cactus/transpiler/main.py | 22 +- 4 files changed, 899 insertions(+), 39 deletions(-) diff --git a/python/cactus/transpiler/Converter/convert.py b/python/cactus/transpiler/Converter/convert.py index 3c1f1c40a..025a0a370 100644 --- a/python/cactus/transpiler/Converter/convert.py +++ b/python/cactus/transpiler/Converter/convert.py @@ -1,10 +1,81 @@ +from __future__ import annotations + +from pathlib import Path + from . import models -from ..ModelProfiles import profiles as MP from ..ModelProfiles import models as MP_Models +from ..ModelProfiles import profiles as MP + + +#Builds a lookup table from model profile names to profile objects. +def _profile_by_name() -> dict[str, MP_Models.ModelProfile]: + return {profile.model_profiles: profile for profile in MP.ALL_PROFILES} + + +#Selects the best known model profile for a Hugging Face model id. +def _infer_profile(model_id: str) -> MP_Models.ModelProfile: + model_id_lower = model_id.lower() + profiles_by_name = _profile_by_name() + + for profile_name, mapped_model_id in MP.MODEL_ID_MAP.items(): + if model_id_lower == mapped_model_id.lower() and profile_name in profiles_by_name: + return profiles_by_name[profile_name] + + for profile_name, mapped_model_id in models.default_model_ids.items(): + if model_id_lower == mapped_model_id.lower() and profile_name in profiles_by_name: + return profiles_by_name[profile_name] + + if "gemma" in model_id_lower: + return MP.GEMMA4_E2B_PROFILE + if "whisper" in model_id_lower: + return MP.WHISPER_PROFILE + if "parakeet" in model_id_lower: + return MP.PARAKEET_PROFILE + if "lfm" in model_id_lower and ("vl" in model_id_lower or "vision" in model_id_lower): + return MP.LFM_VLM_PROFILE + if "qwen2.5" in model_id_lower or "qwen2_5" in model_id_lower: + return MP.QWEN2_5_0_5B_PROFILE + + raise ValueError(f"Could not infer model profile for {model_id}. Pass custom_profile instead.") + + +#Resolves where the exported LayerMap JSON should be written. +def _output_json_path(output_path: str | Path | None = None) -> Path: + if output_path is None: + path = Path(__file__).resolve().parent / "jsons" / "output.json" + else: + path = Path(output_path) + if path.suffix != ".json": + path = path / "output.json" + + path.parent.mkdir(parents=True, exist_ok=True) + return path + + +#Runs the full converter pipeline and writes the exported LayerMap JSON. +def convert( + model_id: str, + input_modalities: tuple[str, ...], + custom_profile: MP_Models.ModelProfile | None = None, + inference_mode: str = "prefill_no_cache", + output_path: str | Path | None = None, +) -> models.LayerMap: + profile = custom_profile or _infer_profile(model_id) + model = models.create_model( + mp=profile, + input_modalities=input_modalities, + model_id=model_id, + inference_mode=inference_mode, + ) + layer_map = model.export(model.input) + + output_file = _output_json_path(output_path) + output_file.write_text(layer_map.model_dump_json(indent=4), encoding="utf-8") + return layer_map -def convert(model_id:str, input_modalities:tuple[str,...], custom_profile:MP_Models.ModelProfile = None) -> models.LayerMap | None: - if custom_profile is not None: - return None - - \ No newline at end of file +if __name__ == "__main__": + convert( + model_id="google/gemma-4-E2B", + input_modalities=("text", "vision", "audio"), + ) diff --git a/python/cactus/transpiler/Converter/input_utils.py b/python/cactus/transpiler/Converter/input_utils.py index 2fb23566b..6e309ca59 100644 --- a/python/cactus/transpiler/Converter/input_utils.py +++ b/python/cactus/transpiler/Converter/input_utils.py @@ -5,6 +5,7 @@ from huggingface_hub import hf_hub_download import json from typing import Any +import numpy as np #Locations to inputs for modality_input_path = { @@ -13,12 +14,12 @@ } -#Downloads the file from hf -def load_files(mp: MP_Models.ModelProfile, model_id: str) -> dict[str, dict[str, Any]]: +#Downloads profile-declared files from Hugging Face into the local converter jsons folder. +def load_files(mp: MP_Models.ModelProfile, model_id: str) -> dict[str, str]: output_dir = Path(__file__).resolve().parent / "jsons" / mp.model_profiles output_dir.mkdir(parents=True, exist_ok=True) token = os.environ.get("HF_TOKEN") - downloaded_files: dict[str, dict[str, Any]] = {} + downloaded_files: dict[str, str] = {} for file in mp.files: try: @@ -27,15 +28,24 @@ def load_files(mp: MP_Models.ModelProfile, model_id: str) -> dict[str, dict[str, print(f"Error downloading {file} from {model_id}: {e}") continue - if file.endswith(".json"): - with open(downloaded_path, "r", encoding="utf-8") as f: - downloaded_files[file] = json.load(f) - downloaded_files[file] = str(downloaded_path) return downloaded_files -#TODO: Clean up all functions below this point +#Loads downloaded JSON files into dictionaries keyed by filename. +def _load_json_files(files: dict[str, str]) -> dict[str, dict[str, Any]]: + configs: dict[str, dict[str, Any]] = {} + + for file_name, file_path in files.items(): + if not file_name.endswith(".json"): + continue + + with open(file_path, "r", encoding="utf-8") as f: + configs[file_name] = json.load(f) + + return configs + + #Searches through configs to find desired keys, and their values def _find_config_value(configs: dict[str, dict[str, Any]], keys: tuple[str, ...]) -> Any | None: search_sections = ( @@ -68,6 +78,7 @@ def _find_config_value(configs: dict[str, dict[str, Any]], keys: tuple[str, ...] return None +#Returns the first config value that can safely be interpreted as an int. def _first_int(configs: dict[str, dict[str, Any]], keys: tuple[str, ...], default: int) -> int: value = _find_config_value(configs, keys) if isinstance(value, bool): @@ -81,15 +92,18 @@ def _first_int(configs: dict[str, dict[str, Any]], keys: tuple[str, ...], defaul return default +#Builds simple token ids for generic text-only synthetic exports. def _build_input_ids(batch_size: int, sequence_length: int) -> torch.Tensor: tokens = torch.arange(sequence_length, dtype=torch.long) + 1 return tokens.unsqueeze(0).repeat(batch_size, 1) +#Builds an all-ones attention mask for synthetic exports. def _build_attention_mask(batch_size: int, sequence_length: int) -> torch.Tensor: return torch.ones((batch_size, sequence_length), dtype=torch.long) +#Builds synthetic image pixel tensors from config-derived image dimensions. def _build_pixel_values(batch_size: int, configs: dict[str, dict[str, Any]]) -> torch.Tensor: size = _find_config_value(configs, ("size", "image_size")) channels = _first_int(configs, ("num_channels", "image_num_channels"), 3) @@ -106,6 +120,7 @@ def _build_pixel_values(batch_size: int, configs: dict[str, dict[str, Any]]) -> return torch.zeros((batch_size, channels, height, width), dtype=torch.float32) +#Builds synthetic audio feature tensors from config-derived audio dimensions. def _build_input_features(batch_size: int, configs: dict[str, dict[str, Any]]) -> torch.Tensor: feature_bins = _first_int(configs, ("feature_size", "num_mel_bins", "n_mels"), 80) max_source_positions = _first_int(configs, ("max_source_positions",), 1500) @@ -113,6 +128,7 @@ def _build_input_features(batch_size: int, configs: dict[str, dict[str, Any]]) - return torch.zeros((batch_size, feature_bins, feature_frames), dtype=torch.float32) +#Builds generic tuple-style KV cache tensors for models that still accept legacy cache inputs. def _build_past_key_values(configs: dict[str, dict[str, Any]], batch_size: int, past_sequence_length: int) -> tuple[tuple[torch.Tensor, torch.Tensor], ...]: num_layers = _first_int(configs, ("num_hidden_layers", "n_layer", "num_layers", "decoder_layers"), 1) num_attention_heads = _first_int(configs, ("num_attention_heads", "n_head", "decoder_attention_heads"), 1) @@ -130,5 +146,156 @@ def _build_past_key_values(configs: dict[str, dict[str, Any]], batch_size: int, ) +#Finds a decoder start token id for seq2seq-style synthetic inputs. def _decoder_start_token_id(configs: dict[str, dict[str, Any]]) -> int: - return _first_int(configs, ("decoder_start_token_id", "bos_token_id", "pad_token_id"), 1) \ No newline at end of file + return _first_int(configs, ("decoder_start_token_id", "bos_token_id", "pad_token_id"), 1) + + +#Loads the local image asset used for processor-backed multimodal exports. +def _load_image_asset(): + from PIL import Image + + return Image.open(modality_input_path["vision"]).convert("RGB") + + +#Loads and normalizes the local audio asset used for processor-backed multimodal exports. +def _load_audio_asset() -> np.ndarray: + from scipy.io import wavfile + + _sample_rate, audio = wavfile.read(modality_input_path["audio"]) + + if audio.ndim > 1: + audio = audio.mean(axis=1) + + if np.issubdtype(audio.dtype, np.integer): + audio = audio.astype(np.float32) / np.iinfo(audio.dtype).max + else: + audio = audio.astype(np.float32) + + return audio + + +#Removes Hugging Face processor type metadata before constructing processor components directly. +def _without_type_key(config: dict[str, Any], type_key: str) -> dict[str, Any]: + return {key: value for key, value in config.items() if key != type_key} + + +#Finds a previously cached Hugging Face file without requiring network access. +def _find_cached_hf_file(model_id: str, filename: str) -> str | None: + cache_model_id = model_id.replace("/", "--") + cache_root = Path.home() / ".cache" / "huggingface" / "hub" / f"models--{cache_model_id}" + + if not cache_root.exists(): + return None + + for candidate in cache_root.glob(f"snapshots/*/{filename}"): + if candidate.is_file(): + return str(candidate) + + return None + + +#Gemma-specific: manually builds Gemma4Processor when AutoProcessor cannot be loaded cleanly offline. +def _build_gemma4_processor(model_id: str, configs: dict[str, dict[str, Any]]): + from transformers import PreTrainedTokenizerFast + from transformers.models.gemma4 import ( + Gemma4AudioFeatureExtractor, + Gemma4ImageProcessor, + Gemma4Processor, + Gemma4VideoProcessor, + ) + + processor_config = configs.get("processor_config.json") + if processor_config is None or processor_config.get("processor_class") != "Gemma4Processor": + return None + + token = os.environ.get("HF_TOKEN") + tokenizer_path = _find_cached_hf_file(model_id, "tokenizer.json") + if tokenizer_path is None: + try: + tokenizer_path = hf_hub_download(repo_id=model_id, filename="tokenizer.json", token=token, local_files_only=True) + except Exception: + tokenizer_path = hf_hub_download(repo_id=model_id, filename="tokenizer.json", token=token) + + tokenizer_config = configs.get("tokenizer_config.json", {}) + tokenizer = PreTrainedTokenizerFast( + tokenizer_file=tokenizer_path, + bos_token=tokenizer_config.get("bos_token", ""), + eos_token=tokenizer_config.get("eos_token", ""), + unk_token=tokenizer_config.get("unk_token", ""), + pad_token=tokenizer_config.get("pad_token", ""), + ) + + for token_name in ( + "image_token", + "audio_token", + "boi_token", + "eoi_token", + "boa_token", + "eoa_token", + ): + token_value = tokenizer_config.get(token_name) + if token_value is not None: + setattr(tokenizer, token_name, token_value) + + tokenizer.image_token_id = tokenizer.convert_tokens_to_ids(tokenizer.image_token) + tokenizer.audio_token_id = tokenizer.convert_tokens_to_ids(tokenizer.audio_token) + image_processor = Gemma4ImageProcessor( + **_without_type_key(processor_config.get("image_processor", {}), "image_processor_type") + ) + feature_extractor = Gemma4AudioFeatureExtractor( + **_without_type_key(processor_config.get("feature_extractor", {}), "feature_extractor_type") + ) + video_processor = Gemma4VideoProcessor( + **_without_type_key(processor_config.get("video_processor", {}), "video_processor_type") + ) + + return Gemma4Processor( + feature_extractor=feature_extractor, + image_processor=image_processor, + tokenizer=tokenizer, + video_processor=video_processor, + image_seq_length=processor_config.get("image_seq_length", 280), + audio_seq_length=processor_config.get("audio_seq_length", 750), + audio_ms_per_token=processor_config.get("audio_ms_per_token", 40), + ) + + +#Builds real processor-backed tensor kwargs for multimodal models. +def build_processor_kwargs( + model_id: str, + input_modalities: tuple[str, ...], + configs: dict[str, dict[str, Any]] | None = None, +) -> dict[str, Any]: + from transformers import AutoProcessor + + configs = configs or {} + token = os.environ.get("HF_TOKEN") + processor_kwargs: dict[str, Any] = {"trust_remote_code": True, "local_files_only": True} + if token: + processor_kwargs["token"] = token + + processor = _build_gemma4_processor(model_id, configs) + if processor is None: + processor = AutoProcessor.from_pretrained(model_id, **processor_kwargs) + prompt_parts = [] + call_kwargs: dict[str, Any] = {"return_tensors": "pt"} + + if "vision" in input_modalities: + image_token = getattr(processor, "image_token", None) + if image_token is not None: + prompt_parts.append(image_token) + call_kwargs["images"] = _load_image_asset() + + if "audio" in input_modalities: + audio_token = getattr(processor, "audio_token", None) + if audio_token is not None: + prompt_parts.append(audio_token) + call_kwargs["audio"] = _load_audio_asset() + + if "text" in input_modalities: + prompt_parts.append("Describe this input.") + + call_kwargs["text"] = " ".join(prompt_parts) if prompt_parts else "Hello" + processed = processor(**call_kwargs) + return {key: value for key, value in dict(processed).items() if isinstance(value, torch.Tensor)} diff --git a/python/cactus/transpiler/Converter/models.py b/python/cactus/transpiler/Converter/models.py index 368e50263..3c7ea9c54 100644 --- a/python/cactus/transpiler/Converter/models.py +++ b/python/cactus/transpiler/Converter/models.py @@ -1,6 +1,7 @@ import json import os -from dataclasses import dataclass, asdict +import inspect +from dataclasses import dataclass from typing import Any from pydantic import BaseModel import torch @@ -8,6 +9,8 @@ from ..ModelProfiles import models as MP_Models from ..ModelProfiles import profiles as MP_Profiles +from transformers import AutoConfig, AutoModel, AutoModelForCausalLM, AutoModelForCTC, AutoModelForImageTextToText, AutoModelForSeq2SeqLM, AutoModelForSpeechSeq2Seq + default_model_ids = { "gemma4_e2b": "google/gemma-4-E2B", "whisper": "openai/whisper-tiny", @@ -17,6 +20,7 @@ } +#Stores the concrete args/kwargs and mode used for one export run. @dataclass(slots=True) class Input: args: tuple @@ -24,42 +28,209 @@ class Input: modalities: tuple[str, ...] inference_mode: str +#Bundles the loaded HF model with its profile and prepared export input. @dataclass(slots=True) class Model: name: str model_profile: MP_Models.ModelProfile input: Input + model: Any + #Creates a loaded converter model from a model profile. @classmethod def from_profile(mp: MP_Models.ModelProfile, input_modalities: tuple[str,...], model_id: str) -> "Model": return create_model(mp, input_modalities, model_id) + #Exports this model using the provided prepared input. def export(self, input: Input) -> "LayerMap": - export_(model=self, input=input) + return export_(model=self, input=input) +#Stores tensor metadata in a JSON-safe shape/dtype form. @dataclass(slots=True) class TensorInstance: - shape: list[int] + shape: list[Any] dtype: str + #Extracts JSON-safe metadata from real, fake, or symbolic tensors. @classmethod def from_tensor(cls, x: torch.Tensor) -> "TensorInstance": - return cls(shape=list(x.shape), dtype=str(x.dtype)) + return cls(shape=[jsonable_shape_dim(dim) for dim in x.shape], dtype=str(x.dtype)) +#Stores Python slice metadata in a JSON-safe form. @dataclass(slots=True) class Slice: start: Any stop: Any step: Any + #Converts a Python slice into a JSON-safe metadata record. @classmethod def from_slice(cls, x: slice) -> "Slice": return cls(start=jsonable(x.start), stop=jsonable(x.stop), step=jsonable(x.step)) +#Describes one KV cache layer's tensor shapes and cache behavior. +@dataclass(slots=True) +class CacheLayerSpec: + index: int + layer_type: str + key_shape: tuple[int, ...] + value_shape: tuple[int, ...] + sliding_window: int | None = None + + +#Describes all flat KV tensors needed to rebuild an HF DynamicCache. @dataclass(slots=True) +class CacheSpec: + layers: tuple[CacheLayerSpec, ...] + config: Any + past_sequence_length: int + dtype: torch.dtype + device: torch.device + + #Infers cache tensor shapes from the loaded model config and attention modules. + @classmethod + def from_model( + cls, + model: torch.nn.Module, + batch_size: int, + past_sequence_length: int, + ) -> "CacheSpec": + config = getattr(model, "config", None) + text_config = get_text_config(config) + dtype, device = model_dtype_and_device(model) + num_key_value_heads = int(getattr(text_config, "num_key_value_heads", getattr(text_config, "num_attention_heads", 1))) + num_attention_heads = int(getattr(text_config, "num_attention_heads", num_key_value_heads)) + hidden_size = int(getattr(text_config, "hidden_size", num_attention_heads)) + head_dim = int(getattr(text_config, "head_dim", max(1, hidden_size // max(1, num_attention_heads)))) + sliding_window = getattr(text_config, "sliding_window", None) or getattr(text_config, "attention_chunk_size", None) + decoder_layers = model_decoder_layers(model) + layers: list[CacheLayerSpec] = [] + + for index, layer_type in enumerate(cache_layer_types(text_config)): + cache_sequence_length = past_sequence_length + layer_sliding_window = None + layer_num_key_value_heads = num_key_value_heads + layer_head_dim = head_dim + + if index < len(decoder_layers): + attention = attention_module(decoder_layers[index]) + layer_head_dim = int(getattr(attention, "head_dim", layer_head_dim)) + key_out_features = linear_out_features(getattr(attention, "k_proj", None)) + + if key_out_features is not None: + layer_num_key_value_heads = max(1, key_out_features // max(1, layer_head_dim)) + + if layer_type in ("sliding_attention", "chunked_attention") and sliding_window is not None: + layer_sliding_window = int(sliding_window) + cache_sequence_length = min(past_sequence_length, max(1, layer_sliding_window - 1)) + + cache_shape = (batch_size, layer_num_key_value_heads, cache_sequence_length, layer_head_dim) + layers.append( + CacheLayerSpec( + index=index, + layer_type=str(layer_type), + key_shape=cache_shape, + value_shape=cache_shape, + sliding_window=layer_sliding_window, + ) + ) + + return cls( + layers=tuple(layers), + config=config, + past_sequence_length=past_sequence_length, + dtype=dtype, + device=device, + ) + + #Creates zero-filled flat KV tensors that match this cache spec. + def empty_tensors(self) -> tuple[torch.Tensor, ...]: + tensors: list[torch.Tensor] = [] + + for layer in self.layers: + tensors.append(torch.zeros(layer.key_shape, dtype=self.dtype, device=self.device)) + tensors.append(torch.zeros(layer.value_shape, dtype=self.dtype, device=self.device)) + + return tuple(tensors) + + #Rebuilds an HF DynamicCache from flat exported tensor inputs. + def to_dynamic_cache(self, flat_tensors: tuple[torch.Tensor, ...]): + from transformers.cache_utils import DynamicCache + + if len(flat_tensors) != len(self.layers) * 2: + raise ValueError(f"Expected {len(self.layers) * 2} cache tensors, got {len(flat_tensors)}") + + cache = DynamicCache(config=self.config) + + for layer_spec, flat_index in zip(self.layers, range(0, len(flat_tensors), 2)): + layer = cache.layers[layer_spec.index] + key = flat_tensors[flat_index] + value = flat_tensors[flat_index + 1] + layer.keys = key + layer.values = value + layer.dtype = key.dtype + layer.device = key.device + layer.is_initialized = True + + if hasattr(layer, "cumulative_length"): + layer.cumulative_length = self.past_sequence_length + + return cache + + +#Wraps HF cache-mode forwards so torch.export only sees tensor inputs/outputs. +class CacheExportWrapper(torch.nn.Module): + #Stores the wrapped model and cache spec used to tensorize cache state. + def __init__(self, model: torch.nn.Module, cache_spec: CacheSpec, mode: str): + super().__init__() + self.model = model + self.cache_spec = cache_spec + self.mode = mode + + #Runs the HF model with DynamicCache internally and returns logits plus flat cache tensors. + def forward( + self, + input_ids=None, + pixel_values=None, + pixel_values_videos=None, + input_features=None, + attention_mask=None, + input_features_mask=None, + position_ids=None, + image_position_ids=None, + video_position_ids=None, + cache_position=None, + past_key_values=None, + mm_token_type_ids=None, + inputs_embeds=None, + ): + cache = self.cache_spec.to_dynamic_cache(past_key_values) if past_key_values is not None else None + model_kwargs = { + "input_ids": input_ids, + "pixel_values": pixel_values, + "pixel_values_videos": pixel_values_videos, + "input_features": input_features, + "attention_mask": attention_mask, + "input_features_mask": input_features_mask, + "position_ids": position_ids, + "image_position_ids": image_position_ids, + "video_position_ids": video_position_ids, + "cache_position": cache_position, + "past_key_values": cache, + "mm_token_type_ids": mm_token_type_ids, + "inputs_embeds": inputs_embeds, + "use_cache": True, + } + model_kwargs = {key: value for key, value in model_kwargs.items() if value is not None} + outputs = self.model(**filter_forward_kwargs(self.model, model_kwargs)) + + return (primary_model_output(outputs), *flatten_dynamic_cache(outputs.past_key_values)) + + +#Serializable record for one exported FX graph node. class LayerRecord(BaseModel): index: int name: str @@ -71,6 +242,7 @@ class LayerRecord(BaseModel): tensor_output_meta: Any | None module_stack: Any | None + #Builds a LayerRecord from a torch.fx.Node. @classmethod def from_node(cls, num: int, x: torch.fx.Node) -> "LayerRecord": return cls( @@ -87,7 +259,7 @@ def from_node(cls, num: int, x: torch.fx.Node) -> "LayerRecord": -@dataclass(slots=True) +#Serializable top-level export IR container. class LayerMap(BaseModel): model_name: str task: str @@ -95,24 +267,29 @@ class LayerMap(BaseModel): range_constants: str nodes: list[LayerRecord] + #Builds a LayerMap from a torch.export ExportedProgram and serialized nodes. @classmethod def from_data(cls, x: torch.export.ExportedProgram, name: str, model_task: str, nodes_list: list[LayerRecord]) -> "LayerMap": return cls(model_name=name, task=model_task, graph_signature=repr(x.graph_signature), range_constants=repr(x.range_constraints), nodes=nodes_list) - def model_dump(self) -> dict[str, Any]: - return asdict(self) +"""#####################################Model Utils#####################################""" - def model_dump_json(self, indent: int | None = None) -> str: - return json.dumps(self.model_dump(), indent=indent) +#Converts symbolic shape dims to JSON-safe values without forcing torch guards. +def jsonable_shape_dim(x: Any) -> Any: + if type(x) is int: + return x + + return str(x) -"""#####################################Model Utils#####################################""" +#Recursively converts FX/export metadata into JSON-safe Python values. def jsonable(x: Any) -> Any: if isinstance(x, torch.fx.Node): return {"node": x.name} if isinstance(x, torch.Tensor): - return asdict(TensorInstance.from_tensor(x)) + tensor = TensorInstance.from_tensor(x) + return {"shape": tensor.shape, "dtype": tensor.dtype} if isinstance(x, torch.Size): return list(x) @@ -124,7 +301,8 @@ def jsonable(x: Any) -> Any: return str(x) if isinstance(x, slice): - return asdict(Slice.from_slice(x)) + slice_ = Slice.from_slice(x) + return {"start": slice_.start, "stop": slice_.stop, "step": slice_.step} if isinstance(x, range): return list(x) @@ -141,6 +319,146 @@ def jsonable(x: Any) -> Any: return repr(x) +#Returns the text/decoder config for models that wrap language configs inside multimodal configs. +def get_text_config(config: Any) -> Any: + if config is not None and hasattr(config, "get_text_config"): + try: + return config.get_text_config(decoder=True) + except TypeError: + return config.get_text_config() + + return config + + +#Finds the dtype and device of the loaded model parameters. +def model_dtype_and_device(model: torch.nn.Module) -> tuple[torch.dtype, torch.device]: + try: + param = next(model.parameters()) + return param.dtype, param.device + except StopIteration: + return torch.float32, torch.device("cpu") + + +#Finds decoder layer modules across common HF model wrapper layouts. +def model_decoder_layers(model: torch.nn.Module) -> tuple[Any, ...]: + candidates = ( + ("model", "language_model", "layers"), + ("language_model", "layers"), + ("model", "decoder", "layers"), + ("decoder", "layers"), + ("model", "layers"), + ("layers",), + ) + + for path in candidates: + value: Any = model + + for attr in path: + value = getattr(value, attr, None) + if value is None: + break + + if value is not None: + return tuple(value) + + return () + + +#Finds the attention module inside one decoder layer across common naming conventions. +def attention_module(layer: Any) -> Any: + for attr in ("self_attn", "attention", "attn"): + value = getattr(layer, attr, None) + if value is not None: + return value + + return layer + + +#Returns a linear-like module's output feature count when available. +def linear_out_features(module: Any) -> int | None: + if module is None: + return None + + if hasattr(module, "out_features"): + return int(module.out_features) + + weight = getattr(module, "weight", None) + if weight is None and hasattr(module, "linear"): + weight = getattr(module.linear, "weight", None) + + if weight is None: + return None + + return int(weight.shape[0]) + + +#Infers which decoder layers actually own KV cache entries. +def cache_layer_types(text_config: Any) -> tuple[str, ...]: + num_hidden_layers = int(getattr(text_config, "num_hidden_layers", 1)) + layer_types = getattr(text_config, "layer_types", None) + + if layer_types is None: + sliding_window = getattr(text_config, "sliding_window", None) or getattr(text_config, "attention_chunk_size", None) + default_layer_type = "sliding_attention" if sliding_window is not None else "full_attention" + layer_types = tuple(default_layer_type for _ in range(num_hidden_layers)) + else: + layer_types = tuple(str(layer_type) for layer_type in layer_types) + + num_kv_shared_layers = int(getattr(text_config, "num_kv_shared_layers", 0) or 0) + if num_kv_shared_layers > 0: + layer_types = layer_types[: -num_kv_shared_layers] + + return layer_types + + +#Selects the main tensor output from a model output object. +def primary_model_output(outputs: Any) -> torch.Tensor: + if hasattr(outputs, "logits") and outputs.logits is not None: + return outputs.logits + + if hasattr(outputs, "last_hidden_state") and outputs.last_hidden_state is not None: + return outputs.last_hidden_state + + if isinstance(outputs, (tuple, list)) and len(outputs) > 0: + return outputs[0] + + raise ValueError("Could not find a tensor output to export") + + +#Flattens an HF DynamicCache into key/value tensor outputs. +def flatten_dynamic_cache(cache: Any) -> tuple[torch.Tensor, ...]: + if cache is None: + return () + + flat: list[torch.Tensor] = [] + + for layer in cache.layers: + key = getattr(layer, "keys", None) + value = getattr(layer, "values", None) + + if key is None or value is None: + continue + + flat.append(key) + flat.append(value) + + return tuple(flat) + + +#Drops kwargs unsupported by a module's forward signature. +def filter_forward_kwargs(model: torch.nn.Module, kwargs: dict[str, Any]) -> dict[str, Any]: + try: + signature = inspect.signature(model.forward) + except (TypeError, ValueError): + return kwargs + + if any(param.kind == inspect.Parameter.VAR_KEYWORD for param in signature.parameters.values()): + return kwargs + + return {key: value for key, value in kwargs.items() if key in signature.parameters} + + +#Normalizes torch operator targets into stable ATen-style names. def aten_name(target: Any) -> str: schema = getattr(target, "_schema", None) @@ -155,6 +473,7 @@ def aten_name(target: Any) -> str: return str(target) +#Extracts tensor metadata from an exported FX node. def extract_tensor_meta(node: torch.fx.Node) -> Any | None: if "val" not in node.meta: return None @@ -162,6 +481,7 @@ def extract_tensor_meta(node: torch.fx.Node) -> Any | None: return jsonable(node.meta["val"]) +#Extracts module stack metadata from an exported FX node. def extract_module_stack(node: torch.fx.Node) -> Any | None: stack = node.meta.get("nn_module_stack", None) if stack is None: @@ -180,7 +500,7 @@ def extract_module_stack(node: torch.fx.Node) -> Any | None: return out -#TODO: Clean up this function +#Builds representative model inputs for the requested modalities and inference mode. def build_input( mp: MP_Models.ModelProfile, input_modalties: tuple[str, ...], @@ -200,6 +520,19 @@ def build_input( past_sequence_length = 8 kwargs: dict[str, Any] = {} + #Gemma-specific: use the real Gemma4 processor to create multimodal placeholder ids and position tensors. + if model_id is not None and "gemma" in model_id.lower() and any(m in modalities for m in ("vision", "audio")): + try: + kwargs = IU.build_processor_kwargs(model_id, modalities, configs) + return Input( + args=(), + kwargs=kwargs, + modalities=modalities, + inference_mode=inference_mode, + ) + except Exception as e: + print(f"Processor input build failed for {model_id}, falling back to synthetic inputs: {e}") + if "text" in modalities: if inference_mode == "decode_with_cache": sequence_length = 1 @@ -245,15 +578,294 @@ def build_input( ) +#Applies mode-specific input rewrites after the real model/config is loaded. +def prepare_input_for_export(model: torch.nn.Module, input: Input) -> Input: + if input.inference_mode == "decode_with_cache": + return build_decode_with_cache_input(model, input) + return input -def create_model(mp: MP_Models.ModelProfile, input_modalities: tuple[str, ...], model_id: str) -> Model: + +#Builds one-token decode inputs and flat KV tensors for cache-mode exports. +def build_decode_with_cache_input(model: torch.nn.Module, input: Input) -> Input: + kwargs = dict(input.kwargs) + token_key = "input_ids" if "input_ids" in kwargs else "decoder_input_ids" + + if token_key not in kwargs: + raise ValueError("decode_with_cache requires input_ids or decoder_input_ids") + + token_ids = kwargs[token_key] + batch_size = int(token_ids.shape[0]) + past_sequence_length = int(token_ids.shape[1]) + cache_spec = CacheSpec.from_model( + model=model, + batch_size=batch_size, + past_sequence_length=past_sequence_length, + ) + + kwargs[token_key] = token_ids[:, -1:].clone() + kwargs["attention_mask"] = torch.ones( + (batch_size, past_sequence_length + 1), + dtype=torch.long, + device=token_ids.device, + ) + kwargs["cache_position"] = torch.arange( + past_sequence_length, + past_sequence_length + 1, + dtype=torch.long, + device=token_ids.device, + ) + kwargs["past_key_values"] = cache_spec.empty_tensors() + + for multimodal_key in ( + "pixel_values", + "pixel_values_videos", + "input_features", + "input_features_mask", + "image_position_ids", + "video_position_ids", + "mm_token_type_ids", + ): + kwargs.pop(multimodal_key, None) + + return Input( + args=input.args, + kwargs=kwargs, + modalities=input.modalities, + inference_mode=input.inference_mode, + ) + + +#Loads the best HF model class for the requested model profile. +def load_model(model_id: str, mp: MP_Models.ModelProfile | None = None) -> torch.nn.Module: + + + token = os.environ.get("HF_TOKEN") + load_kwargs: dict[str, Any] = {"trust_remote_code": True} + if token: + load_kwargs["token"] = token + + AutoConfig.from_pretrained(model_id, **load_kwargs) + + candidate_classes: list[Any] = [] + features = set(mp.features if mp is not None else ()) + modalities = set(mp.supported_modalties if mp is not None else ()) + is_gemma4_profile = mp is not None and mp.model_profiles == "gemma4_e2b" + + if "vision" in modalities and "text" in modalities: + candidate_classes.extend((AutoModelForImageTextToText, AutoModelForCausalLM)) + #Gemma-specific: avoid falling back to plain AutoModel, which loads Gemma4 weights with wrong prefixes. + if not is_gemma4_profile: + candidate_classes.append(AutoModel) + elif "audio" in modalities and "text" in modalities: + candidate_classes.extend((AutoModelForSpeechSeq2Seq, AutoModelForSeq2SeqLM, AutoModel)) + elif "audio" in modalities: + candidate_classes.extend((AutoModelForCTC, AutoModel)) + elif "text" in modalities or "causal_lm" in features: + candidate_classes.extend((AutoModelForCausalLM, AutoModel)) + else: + candidate_classes.append(AutoModel) + + last_error: Exception | None = None + seen_classes: set[Any] = set() + + for model_class in candidate_classes: + if model_class in seen_classes: + continue + seen_classes.add(model_class) + + for attempt_kwargs in _model_load_kwargs(load_kwargs): + try: + model = model_class.from_pretrained(model_id, **attempt_kwargs) + model.eval() + + configure_model_for_export(model, should_use_cache=False) + + #Gemma-specific: patch Gemma4 audio masking so torch.export can trace it. + if mp is not None and mp.model_profiles == "gemma4_e2b" and "audio" in mp.supported_modalties: + patch_gemma4_audio_mask_for_export() + + return model + except Exception as e: + last_error = e + + raise RuntimeError(f"Unable to load model {model_id}") from last_error + + +#Returns normal load kwargs plus a local-files-only retry for cached HF models. +def _model_load_kwargs(load_kwargs: dict[str, Any]) -> tuple[dict[str, Any], ...]: + if load_kwargs.get("local_files_only") is True: + return (load_kwargs,) + + return ( + load_kwargs, + {**load_kwargs, "local_files_only": True}, + ) + + +#Sets an attribute on a config object and common nested config objects. +def _set_config_attr_recursive(config: Any, attr: str, value: Any, seen: set[int] | None = None) -> None: + if config is None: + return + + seen = seen or set() + config_id = id(config) + if config_id in seen: + return + + seen.add(config_id) + + if hasattr(config, attr): + setattr(config, attr, value) + + for child_name in ( + "text_config", + "vision_config", + "audio_config", + "encoder", + "decoder", + "model_config", + ): + _set_config_attr_recursive(getattr(config, child_name, None), attr, value, seen) + + +#Configures cache-related model settings before torch.export. +def configure_model_for_export(model: torch.nn.Module, should_use_cache: bool) -> None: + for config_name in ("config", "generation_config"): + config = getattr(model, config_name, None) + _set_config_attr_recursive(config, "use_cache", should_use_cache) + + +#Checks whether a module forward can accept a specific keyword. +def _forward_accepts_kwarg(model: torch.nn.Module, key: str) -> bool: + try: + signature = inspect.signature(model.forward) + except (TypeError, ValueError): + return False + + if key in signature.parameters: + return True + + return any(param.kind == inspect.Parameter.VAR_KEYWORD for param in signature.parameters.values()) + + +#Builds the kwargs actually passed into torch.export. +def _export_kwargs(model: torch.nn.Module, input: Input) -> dict[str, Any]: + kwargs = filter_forward_kwargs(model, dict(input.kwargs)) + should_use_cache = input.inference_mode in ("prefill_with_cache", "decode_with_cache") + + if _forward_accepts_kwarg(model, "use_cache"): + kwargs["use_cache"] = should_use_cache + + return kwargs + + +#Gemma-specific: replaces Gemma4's audio bidirectional mask helper with an exportable tensor implementation. +def patch_gemma4_audio_mask_for_export() -> None: + import transformers.models.gemma4.modeling_gemma4 as gemma4_modeling + + #Gemma-specific: computes Gemma4 audio tower local bidirectional attention masks without vmap/proxy issues. + def exportable_bidirectional_mask(config, inputs_embeds, attention_mask=None, and_mask_function=None, **kwargs): + batch_size, seq_len = inputs_embeds.shape[:2] + device = inputs_embeds.device + q_idx = torch.arange(seq_len, device=device).view(1, 1, seq_len, 1) + kv_idx = torch.arange(seq_len, device=device).view(1, 1, 1, seq_len) + + left_window_size = getattr(config, "attention_context_left", seq_len) - 1 + right_window_size = getattr(config, "attention_context_right", 0) + distance = q_idx - kv_idx + left_mask = (distance >= 0) & (distance < left_window_size) + right_mask = (distance < 0) & (-distance < right_window_size) + mask = left_mask | right_mask + + if attention_mask is not None: + mask = mask & attention_mask[:, None, None, :].bool() + + return mask.expand(batch_size, 1, seq_len, seq_len) + + gemma4_modeling.create_bidirectional_mask = exportable_bidirectional_mask + + +#Builds a loaded model bundle from profile, modalities, model id, and inference mode. +def create_model( + mp: MP_Models.ModelProfile, + input_modalities: tuple[str, ...], + model_id: str, + inference_mode: str = "prefill_no_cache", +) -> Model: name_ = model_id - input_ = build_input(mp, input_modalities) + input_ = build_input(mp, input_modalities, model_id=model_id, inference_mode=inference_mode) + + if input_ is None: + raise ValueError(f"Could not build input for modalities {input_modalities}") + + loaded_model = load_model(model_id, mp) + input_ = prepare_input_for_export(loaded_model, input_) + return Model(name=name_, model_profile=mp, input=input_, model=loaded_model) +#Selects the raw model or cache tensorization wrapper for export. +def build_export_model(model: torch.nn.Module, input: Input) -> torch.nn.Module: + if input.inference_mode not in ("prefill_with_cache", "decode_with_cache"): + return model + batch_size = infer_batch_size(input.kwargs) + past_sequence_length = infer_past_sequence_length(input) + cache_spec = CacheSpec.from_model( + model=model, + batch_size=batch_size, + past_sequence_length=past_sequence_length, + ) + return CacheExportWrapper(model=model, cache_spec=cache_spec, mode=input.inference_mode) + + +#Infers batch size from the first tensor-shaped input. +def infer_batch_size(kwargs: dict[str, Any]) -> int: + for value in kwargs.values(): + if isinstance(value, torch.Tensor) and value.ndim > 0: + return int(value.shape[0]) + + return 1 + + +#Infers the past/prompt sequence length represented by the export input. +def infer_past_sequence_length(input: Input) -> int: + if input.inference_mode == "decode_with_cache" and "cache_position" in input.kwargs: + return int(input.kwargs["cache_position"][0].item()) + + if "input_ids" in input.kwargs: + return int(input.kwargs["input_ids"].shape[1]) + + if "decoder_input_ids" in input.kwargs: + return int(input.kwargs["decoder_input_ids"].shape[1]) + + if "attention_mask" in input.kwargs: + return int(input.kwargs["attention_mask"].shape[1]) + + return 0 + + + +#Exports the prepared model and serializes its FX graph into a LayerMap. def export_(model: Model, input: Input) -> LayerMap: - exported = torch.export.export(model.model, args=input.args, kwargs=input.kwargs, strict=False,).run_decompositions() - records = [LayerRecord.from_node(i, node) for i, node in enumerate(exported.graph.nodes)] - return LayerMap.from_exported_model(model=model, input=input, exported_model=exported, nodes=records) + should_use_cache = input.inference_mode in ("prefill_with_cache", "decode_with_cache") + configure_model_for_export(model.model, should_use_cache=should_use_cache) + export_model = build_export_model(model.model, input) + export_kwargs = _export_kwargs(export_model, input) + + with torch.no_grad(): + exported = torch.export.export( + export_model, + args=input.args, + kwargs=export_kwargs, + strict=False, + ).run_decompositions() + + graph = exported.graph_module.graph if hasattr(exported, "graph_module") else exported.graph + records = [LayerRecord.from_node(i, node) for i, node in enumerate(graph.nodes)] + return LayerMap.from_data( + x=exported, + name=model.name, + model_task=input.inference_mode, + nodes_list=records, + ) diff --git a/python/cactus/transpiler/main.py b/python/cactus/transpiler/main.py index b36194922..339f5ebae 100644 --- a/python/cactus/transpiler/main.py +++ b/python/cactus/transpiler/main.py @@ -1,7 +1,17 @@ -from dataclasses import dataclass +import sys +from pathlib import Path -@dataclass(slots=True) -class ModelProfile: - model_family:str - component_graphs:tuple[str] - export_modes:ep \ No newline at end of file +sys.path.append(str(Path(__file__).resolve().parents[1])) + +from transpiler.Converter.convert import convert + + +def main() -> None: + convert( + model_id="google/gemma-4-E2B", + input_modalities=("text", "vision", "audio"), + ) + + +if __name__ == "__main__": + main() From 26e53ba47c55b60fa1f86e6385ef36149bf26b50 Mon Sep 17 00:00:00 2001 From: Parkirat Date: Tue, 21 Jul 2026 11:59:58 -0400 Subject: [PATCH 33/84] Modifications to ModelProfiles classes to add additional cache and input meta-data --- python/cactus/transpiler/Converter/convert.py | 8 +-- python/cactus/transpiler/Converter/models.py | 40 ++++------- .../cactus/transpiler/ModelProfiles/models.py | 20 +++--- .../transpiler/ModelProfiles/profiles.py | 70 ++++++------------- 4 files changed, 49 insertions(+), 89 deletions(-) diff --git a/python/cactus/transpiler/Converter/convert.py b/python/cactus/transpiler/Converter/convert.py index 025a0a370..275f09740 100644 --- a/python/cactus/transpiler/Converter/convert.py +++ b/python/cactus/transpiler/Converter/convert.py @@ -1,5 +1,3 @@ -from __future__ import annotations - from pathlib import Path from . import models @@ -21,9 +19,9 @@ def _infer_profile(model_id: str) -> MP_Models.ModelProfile: if model_id_lower == mapped_model_id.lower() and profile_name in profiles_by_name: return profiles_by_name[profile_name] - for profile_name, mapped_model_id in models.default_model_ids.items(): - if model_id_lower == mapped_model_id.lower() and profile_name in profiles_by_name: - return profiles_by_name[profile_name] + for mapped_model_id, profile in models.default_model_ids.items(): + if model_id_lower == mapped_model_id.lower(): + return profile if "gemma" in model_id_lower: return MP.GEMMA4_E2B_PROFILE diff --git a/python/cactus/transpiler/Converter/models.py b/python/cactus/transpiler/Converter/models.py index 3c7ea9c54..4441bce0a 100644 --- a/python/cactus/transpiler/Converter/models.py +++ b/python/cactus/transpiler/Converter/models.py @@ -11,16 +11,15 @@ from transformers import AutoConfig, AutoModel, AutoModelForCausalLM, AutoModelForCTC, AutoModelForImageTextToText, AutoModelForSeq2SeqLM, AutoModelForSpeechSeq2Seq -default_model_ids = { - "gemma4_e2b": "google/gemma-4-E2B", - "whisper": "openai/whisper-tiny", - "parakeet": "nvidia/parakeet-tdt-0.6b-v3", - "lfm_vlm": "LiquidAI/LFM2-VL-3B", - "qwen2_5_0_5b": "Qwen/Qwen2.5-0.5B", +default_model_ids: dict[str, MP_Models.ModelProfile] = { + "google/gemma-4-E2B": MP_Profiles.GEMMA4_E2B_PROFILE, + "openai/whisper-tiny": MP_Profiles.WHISPER_PROFILE, + "nvidia/parakeet-tdt-0.6b-v3": MP_Profiles.PARAKEET_PROFILE, + "LiquidAI/LFM2-VL-3B": MP_Profiles.LFM_VLM_PROFILE, + "Qwen/Qwen2.5-0.5B": MP_Profiles.QWEN2_5_0_5B_PROFILE, } -#Stores the concrete args/kwargs and mode used for one export run. @dataclass(slots=True) class Input: args: tuple @@ -28,7 +27,6 @@ class Input: modalities: tuple[str, ...] inference_mode: str -#Bundles the loaded HF model with its profile and prepared export input. @dataclass(slots=True) class Model: name: str @@ -36,42 +34,32 @@ class Model: input: Input model: Any - #Creates a loaded converter model from a model profile. @classmethod def from_profile(mp: MP_Models.ModelProfile, input_modalities: tuple[str,...], model_id: str) -> "Model": return create_model(mp, input_modalities, model_id) - - #Exports this model using the provided prepared input. def export(self, input: Input) -> "LayerMap": return export_(model=self, input=input) - -#Stores tensor metadata in a JSON-safe shape/dtype form. @dataclass(slots=True) class TensorInstance: shape: list[Any] dtype: str - #Extracts JSON-safe metadata from real, fake, or symbolic tensors. @classmethod def from_tensor(cls, x: torch.Tensor) -> "TensorInstance": return cls(shape=[jsonable_shape_dim(dim) for dim in x.shape], dtype=str(x.dtype)) -#Stores Python slice metadata in a JSON-safe form. @dataclass(slots=True) class Slice: start: Any stop: Any step: Any - #Converts a Python slice into a JSON-safe metadata record. @classmethod def from_slice(cls, x: slice) -> "Slice": return cls(start=jsonable(x.start), stop=jsonable(x.stop), step=jsonable(x.step)) - -#Describes one KV cache layer's tensor shapes and cache behavior. @dataclass(slots=True) class CacheLayerSpec: index: int @@ -80,8 +68,6 @@ class CacheLayerSpec: value_shape: tuple[int, ...] sliding_window: int | None = None - -#Describes all flat KV tensors needed to rebuild an HF DynamicCache. @dataclass(slots=True) class CacheSpec: layers: tuple[CacheLayerSpec, ...] @@ -274,7 +260,7 @@ def from_data(cls, x: torch.export.ExportedProgram, name: str, model_task: str, """#####################################Model Utils#####################################""" -#Converts symbolic shape dims to JSON-safe values without forcing torch guards. +#Converts symbolic shape dims to JSON-safe values without forcing torch guards. x def jsonable_shape_dim(x: Any) -> Any: if type(x) is int: return x @@ -282,7 +268,7 @@ def jsonable_shape_dim(x: Any) -> Any: return str(x) -#Recursively converts FX/export metadata into JSON-safe Python values. +#Recursively converts FX/export metadata into JSON-safe Python values. x def jsonable(x: Any) -> Any: if isinstance(x, torch.fx.Node): return {"node": x.name} @@ -648,20 +634,20 @@ def load_model(model_id: str, mp: MP_Models.ModelProfile | None = None) -> torch AutoConfig.from_pretrained(model_id, **load_kwargs) candidate_classes: list[Any] = [] - features = set(mp.features if mp is not None else ()) modalities = set(mp.supported_modalties if mp is not None else ()) + load_strategy = mp.load_strategy if mp is not None else "" is_gemma4_profile = mp is not None and mp.model_profiles == "gemma4_e2b" - if "vision" in modalities and "text" in modalities: + if load_strategy == "image_text_to_text" or ("vision" in modalities and "text" in modalities): candidate_classes.extend((AutoModelForImageTextToText, AutoModelForCausalLM)) #Gemma-specific: avoid falling back to plain AutoModel, which loads Gemma4 weights with wrong prefixes. if not is_gemma4_profile: candidate_classes.append(AutoModel) - elif "audio" in modalities and "text" in modalities: + elif load_strategy == "speech_seq2seq" or ("audio" in modalities and "text" in modalities): candidate_classes.extend((AutoModelForSpeechSeq2Seq, AutoModelForSeq2SeqLM, AutoModel)) - elif "audio" in modalities: + elif load_strategy == "ctc" or "audio" in modalities: candidate_classes.extend((AutoModelForCTC, AutoModel)) - elif "text" in modalities or "causal_lm" in features: + elif load_strategy == "causal_lm" or "text" in modalities: candidate_classes.extend((AutoModelForCausalLM, AutoModel)) else: candidate_classes.append(AutoModel) diff --git a/python/cactus/transpiler/ModelProfiles/models.py b/python/cactus/transpiler/ModelProfiles/models.py index beb036ae9..805659107 100644 --- a/python/cactus/transpiler/ModelProfiles/models.py +++ b/python/cactus/transpiler/ModelProfiles/models.py @@ -18,11 +18,15 @@ class InferencePattern: @dataclass(slots=True) class ModelProfile: - model_profiles:str - components: dict[str, Components] - inference_type: dict[str, InferencePattern] - cache_type: tuple[str, ...] - files: tuple[str, ...] - fusion_fields: tuple[str, ...] - features: tuple[str, ...] - supported_modalties: tuple[str, ...] \ No newline at end of file + model_profiles:str #Model family for which this profile is valid (Theoretically can be shared across model families that have similar architectures) + components: dict[str, Components] #Individual components (vision tower, audio tower, token merge, etc.) this model will be split into + inference_type: dict[str, InferencePattern] #Specific operations mapped to their tuple of components ordered by execution + cache_type: tuple[str, ...] #Which layers/operations to expect KV cache for + cache_policy: tuple[str, ...] #How to handle the cache during different phases of inference + files: tuple[str, ...] #Config and other necessary files + fusion_fields: tuple[str, ...] #Specific fusion groups to consider during fusion process + supported_modalties: tuple[str, ...] #Input modalities supported by model + input_strategy: str #Specifies what functions/procedures to use when generating sample input for torch.export + export_patches: tuple[str, ...] #Specifies what patches/masks to apply to sample inputs + load_strategy: str #Specifies which Hugging Face AutoModel class/loading path to prefer + diff --git a/python/cactus/transpiler/ModelProfiles/profiles.py b/python/cactus/transpiler/ModelProfiles/profiles.py index 6e3ca8069..584326fe2 100644 --- a/python/cactus/transpiler/ModelProfiles/profiles.py +++ b/python/cactus/transpiler/ModelProfiles/profiles.py @@ -20,6 +20,7 @@ components=MULTIMODAL_COMPONENTS, inference_type=GEMMA4_INFERENCE_PATTERNS, cache_type=("attention_kv", "sliding_window_attention_kv"), + cache_policy=("dynamic_cache", "drop_multimodal_on_decode", "shared_kv_layers"), files=("config.json", "generation_config.json", "processor_config.json", "tokenizer_config.json"), fusion_fields=( "generic", @@ -33,19 +34,10 @@ "gemma4_mlp", "linear", ), - features=( - "text", - "vision", - "audio", - "attention", - "sliding_window_attention", - "rope", - "rms_norm", - "conv", - "kv_cache", - "multimodal_token_merge", - ), - supported_modalties=("audio", "vision", "text") + supported_modalties=("audio", "vision", "text"), + input_strategy="manual_gemma4_processor", + export_patches=("gemma4_audio_mask",), + load_strategy="image_text_to_text", ) WHISPER_PROFILE = ModelProfile( @@ -53,6 +45,7 @@ components=SPEECH_SEQ2SEQ_COMPONENTS, inference_type=WHISPER_INFERENCE_PATTERNS, cache_type=("cross_attention_kv", "decoder_attention_kv"), + cache_policy=("dynamic_cache", "encoder_decoder_cache"), files=("config.json", "generation_config.json", "preprocessor_config.json", "tokenizer_config.json"), fusion_fields=( "generic", @@ -62,16 +55,10 @@ "whisper_mlp", "linear", ), - features=( - "audio", - "speech_to_text", - "encoder_decoder", - "attention", - "cross_attention", - "conv", - "kv_cache", - ), supported_modalties=("audio", "text"), + input_strategy="processor", + export_patches=(), + load_strategy="speech_seq2seq", ) PARAKEET_PROFILE = ModelProfile( @@ -79,6 +66,7 @@ components=SPEECH_TRANSCRIBER_COMPONENTS, inference_type=PARAKEET_INFERENCE_PATTERNS, cache_type=(), + cache_policy=(), files=("config.json", "preprocessor_config.json", "tokenizer_config.json"), fusion_fields=( "generic", @@ -88,15 +76,10 @@ "conv", "linear", ), - features=( - "audio", - "speech_to_text", - "conformer", - "tdt", - "conv", - "attention", - ), supported_modalties=("audio",), + input_strategy="processor", + export_patches=(), + load_strategy="ctc", ) LFM_VLM_PROFILE = ModelProfile( @@ -104,6 +87,7 @@ components=VISION_LANGUAGE_COMPONENTS, inference_type=LFM_VLM_INFERENCE_PATTERNS, cache_type=("attention_kv",), + cache_policy=("dynamic_cache", "drop_multimodal_on_decode"), files=("config.json", "generation_config.json", "processor_config.json", "tokenizer_config.json"), fusion_fields=( "generic", @@ -114,17 +98,10 @@ "lfm_mlp", "linear", ), - features=( - "text", - "vision", - "vision_language", - "attention", - "rope", - "rms_norm", - "kv_cache", - "multimodal_token_merge", - ), supported_modalties=("vision", "text"), + input_strategy="processor", + export_patches=(), + load_strategy="image_text_to_text", ) QWEN2_5_0_5B_PROFILE = ModelProfile( @@ -132,6 +109,7 @@ components=TEXT_COMPONENTS, inference_type=QWEN2_5_INFERENCE_PATTERNS, cache_type=("attention_kv",), + cache_policy=("dynamic_cache",), files=("config.json", "generation_config.json", "tokenizer.json", "tokenizer_config.json", "vocab.json", "merges.txt"), fusion_fields=( "generic", @@ -142,16 +120,10 @@ "qwen2_5_mlp", "linear", ), - features=( - "text", - "causal_lm", - "attention", - "rope", - "rms_norm", - "mlp", - "kv_cache", - ), supported_modalties=("text",), + input_strategy="synthetic", + export_patches=(), + load_strategy="causal_lm", ) ALL_PROFILES = ( From 21fa95a18f8512a5f84604bc850f48e2c6c99af8 Mon Sep 17 00:00:00 2001 From: Parkirat Date: Tue, 21 Jul 2026 15:52:34 -0400 Subject: [PATCH 34/84] Streamlining convertert logic + code; still need to clean up --- python/cactus/transpiler/Converter/convert.py | 41 +- .../transpiler/Converter/input_utils.py | 199 +++++++--- python/cactus/transpiler/Converter/models.py | 355 +++++++----------- .../transpiler/ModelProfiles/profiles.py | 26 +- 4 files changed, 304 insertions(+), 317 deletions(-) diff --git a/python/cactus/transpiler/Converter/convert.py b/python/cactus/transpiler/Converter/convert.py index 275f09740..a0d8e4894 100644 --- a/python/cactus/transpiler/Converter/convert.py +++ b/python/cactus/transpiler/Converter/convert.py @@ -5,34 +5,17 @@ from ..ModelProfiles import profiles as MP -#Builds a lookup table from model profile names to profile objects. -def _profile_by_name() -> dict[str, MP_Models.ModelProfile]: - return {profile.model_profiles: profile for profile in MP.ALL_PROFILES} - - #Selects the best known model profile for a Hugging Face model id. def _infer_profile(model_id: str) -> MP_Models.ModelProfile: model_id_lower = model_id.lower() - profiles_by_name = _profile_by_name() - - for profile_name, mapped_model_id in MP.MODEL_ID_MAP.items(): - if model_id_lower == mapped_model_id.lower() and profile_name in profiles_by_name: - return profiles_by_name[profile_name] - for mapped_model_id, profile in models.default_model_ids.items(): + for mapped_model_id, profile in MP.MODEL_ID_MAP.items(): if model_id_lower == mapped_model_id.lower(): return profile - if "gemma" in model_id_lower: - return MP.GEMMA4_E2B_PROFILE - if "whisper" in model_id_lower: - return MP.WHISPER_PROFILE - if "parakeet" in model_id_lower: - return MP.PARAKEET_PROFILE - if "lfm" in model_id_lower and ("vl" in model_id_lower or "vision" in model_id_lower): - return MP.LFM_VLM_PROFILE - if "qwen2.5" in model_id_lower or "qwen2_5" in model_id_lower: - return MP.QWEN2_5_0_5B_PROFILE + for required_terms, profile in MP.MODEL_ID_HINTS: + if all(term in model_id_lower for term in required_terms): + return profile raise ValueError(f"Could not infer model profile for {model_id}. Pass custom_profile instead.") @@ -51,22 +34,10 @@ def _output_json_path(output_path: str | Path | None = None) -> Path: #Runs the full converter pipeline and writes the exported LayerMap JSON. -def convert( - model_id: str, - input_modalities: tuple[str, ...], - custom_profile: MP_Models.ModelProfile | None = None, - inference_mode: str = "prefill_no_cache", - output_path: str | Path | None = None, -) -> models.LayerMap: +def convert(model_id: str, input_modalities: tuple[str, ...], custom_profile: MP_Models.ModelProfile | None = None, inference_mode: str = "prefill_no_cache", output_path: str | Path | None = None) -> models.LayerMap: profile = custom_profile or _infer_profile(model_id) - model = models.create_model( - mp=profile, - input_modalities=input_modalities, - model_id=model_id, - inference_mode=inference_mode, - ) + model = models.create_model(mp=profile, input_modalities=input_modalities, model_id=model_id, inference_mode=inference_mode) layer_map = model.export(model.input) - output_file = _output_json_path(output_path) output_file.write_text(layer_map.model_dump_json(indent=4), encoding="utf-8") return layer_map diff --git a/python/cactus/transpiler/Converter/input_utils.py b/python/cactus/transpiler/Converter/input_utils.py index 6e309ca59..fcf7352d9 100644 --- a/python/cactus/transpiler/Converter/input_utils.py +++ b/python/cactus/transpiler/Converter/input_utils.py @@ -1,65 +1,85 @@ -from ..ModelProfiles import models as MP_Models -import torch +import json import os from pathlib import Path -from huggingface_hub import hf_hub_download -import json from typing import Any + +from huggingface_hub import hf_hub_download import numpy as np +import torch + +from ..ModelProfiles import models as MP_Models + -#Locations to inputs for -modality_input_path = { - "vision": "/Users/sandhup/Documents/personal/cactus/python/cactus/assets/test_monkey.png", - "audio" : "/Users/sandhup/Documents/personal/cactus/python/cactus/assets/test.wav" +DEFAULT_BATCH_SIZE = 1 +DEFAULT_SEQUENCE_LENGTH = 8 +DEFAULT_PAST_SEQUENCE_LENGTH = 8 +SYNTHETIC_INPUT_STRATEGY = "synthetic" +PROCESSOR_INPUT_STRATEGY = "processor" +PREFILL_WITH_CACHE_MODE = "prefill_with_cache" +DECODE_WITH_CACHE_MODE = "decode_with_cache" +CONVERTER_JSON_DIR = Path(__file__).resolve().parent / "jsons" +ASSET_DIR = Path(__file__).resolve().parents[2] / "assets" + +CONFIG_SEARCH_SECTIONS = ( + (), + ("text_config",), + ("vision_config",), + ("audio_config",), + ("encoder",), + ("decoder",), + ("model_config",), +) + +MULTIMODAL_KEYS = ( + "pixel_values", + "pixel_values_videos", + "input_features", + "input_features_mask", + "image_position_ids", + "video_position_ids", + "mm_token_type_ids", +) + +#Locations to local sample inputs used by processor-backed multimodal exports. +MODALITY_INPUT_PATH = { + "vision": ASSET_DIR / "test_monkey.png", + "audio": ASSET_DIR / "test.wav", } -#Downloads profile-declared files from Hugging Face into the local converter jsons folder. -def load_files(mp: MP_Models.ModelProfile, model_id: str) -> dict[str, str]: - output_dir = Path(__file__).resolve().parent / "jsons" / mp.model_profiles +#Downloads profile-declared JSON files from Hugging Face and loads them. +def load_configs(mp: MP_Models.ModelProfile, model_id: str | None) -> dict[str, dict[str, Any]]: + if model_id is None: + return {} + + output_dir = CONVERTER_JSON_DIR / mp.model_profiles output_dir.mkdir(parents=True, exist_ok=True) token = os.environ.get("HF_TOKEN") - downloaded_files: dict[str, str] = {} + configs: dict[str, dict[str, Any]] = {} - for file in mp.files: - try: - downloaded_path = hf_hub_download(repo_id=model_id, filename=file, local_dir=output_dir, token=token) - except Exception as e: - print(f"Error downloading {file} from {model_id}: {e}") + for filename in mp.files: + if not filename.endswith(".json"): continue - - downloaded_files[file] = str(downloaded_path) - return downloaded_files - -#Loads downloaded JSON files into dictionaries keyed by filename. -def _load_json_files(files: dict[str, str]) -> dict[str, dict[str, Any]]: - configs: dict[str, dict[str, Any]] = {} + local_path = output_dir / filename - for file_name, file_path in files.items(): - if not file_name.endswith(".json"): - continue + if not local_path.exists(): + try: + local_path = Path(hf_hub_download(repo_id=model_id, filename=filename, local_dir=output_dir, token=token)) + except Exception as e: + print(f"Error downloading {filename} from {model_id}: {e}") + continue - with open(file_path, "r", encoding="utf-8") as f: - configs[file_name] = json.load(f) + with local_path.open("r", encoding="utf-8") as f: + configs[filename] = json.load(f) return configs #Searches through configs to find desired keys, and their values def _find_config_value(configs: dict[str, dict[str, Any]], keys: tuple[str, ...]) -> Any | None: - search_sections = ( - (), - ("text_config",), - ("vision_config",), - ("audio_config",), - ("encoder",), - ("decoder",), - ("model_config",), - ) - for config in configs.values(): - for section_path in search_sections: + for section_path in CONFIG_SEARCH_SECTIONS: section = config for section_name in section_path: @@ -151,18 +171,72 @@ def _decoder_start_token_id(configs: dict[str, dict[str, Any]]) -> int: return _first_int(configs, ("decoder_start_token_id", "bos_token_id", "pad_token_id"), 1) +#Builds config-shaped synthetic kwargs when a profile does not use a processor. +def build_synthetic_kwargs( + modalities: tuple[str, ...], + configs: dict[str, dict[str, Any]], + inference_mode: str, + uses_dynamic_cache: bool, +) -> dict[str, Any]: + batch_size = DEFAULT_BATCH_SIZE + sequence_length = DEFAULT_SEQUENCE_LENGTH + past_sequence_length = DEFAULT_PAST_SEQUENCE_LENGTH + kwargs: dict[str, Any] = {} + + if "text" in modalities: + if inference_mode == DECODE_WITH_CACHE_MODE and not uses_dynamic_cache: + sequence_length = 1 + + kwargs["input_ids"] = _build_input_ids(batch_size, sequence_length) + kwargs["attention_mask"] = _build_attention_mask(batch_size, sequence_length) + + if "vision" in modalities: + kwargs["pixel_values"] = _build_pixel_values(batch_size, configs) + + if "audio" in modalities: + kwargs["input_features"] = _build_input_features(batch_size, configs) + + if "audio" in modalities and "text" in modalities and "vision" not in modalities and "speech" in inference_mode: + kwargs.pop("input_ids", None) + kwargs.pop("attention_mask", None) + kwargs["decoder_input_ids"] = torch.full( + (batch_size, 1), + _decoder_start_token_id(configs), + dtype=torch.long, + ) + + if "audio" in modalities and "text" not in modalities: + input_features = kwargs["input_features"] + kwargs["attention_mask"] = _build_attention_mask(batch_size, input_features.shape[-1]) + + if inference_mode == DECODE_WITH_CACHE_MODE and not uses_dynamic_cache: + total_sequence_length = past_sequence_length + sequence_length + kwargs["attention_mask"] = _build_attention_mask(batch_size, total_sequence_length) + kwargs["past_key_values"] = _build_past_key_values(configs, batch_size, past_sequence_length) + kwargs["cache_position"] = torch.arange( + past_sequence_length, + past_sequence_length + sequence_length, + dtype=torch.long, + ) + kwargs["use_cache"] = True + elif inference_mode == PREFILL_WITH_CACHE_MODE: + kwargs["use_cache"] = True + + return kwargs + + #Loads the local image asset used for processor-backed multimodal exports. def _load_image_asset(): from PIL import Image - return Image.open(modality_input_path["vision"]).convert("RGB") + return Image.open(MODALITY_INPUT_PATH["vision"]).convert("RGB") #Loads and normalizes the local audio asset used for processor-backed multimodal exports. def _load_audio_asset() -> np.ndarray: from scipy.io import wavfile - _sample_rate, audio = wavfile.read(modality_input_path["audio"]) + _sample_rate, audio = wavfile.read(MODALITY_INPUT_PATH["audio"]) if audio.ndim > 1: audio = audio.mean(axis=1) @@ -261,23 +335,52 @@ def _build_gemma4_processor(model_id: str, configs: dict[str, dict[str, Any]]): ) +#Loads AutoProcessor from local HF cache first, then falls back to normal HF loading. +def _load_auto_processor(model_id: str, processor_kwargs: dict[str, Any]): + from transformers import AutoProcessor + + last_error: Exception | None = None + + for attempt_kwargs in hf_load_kwargs(processor_kwargs): + try: + return AutoProcessor.from_pretrained(model_id, **attempt_kwargs) + except Exception as e: + last_error = e + + raise RuntimeError(f"Unable to load processor for {model_id}") from last_error + + +#Returns local-cache-first HF loading kwargs plus a normal network-capable fallback. +def hf_load_kwargs(load_kwargs: dict[str, Any]) -> tuple[dict[str, Any], ...]: + if load_kwargs.get("local_files_only") is True: + return (load_kwargs,) + + return ( + {**load_kwargs, "local_files_only": True}, + load_kwargs, + ) + + #Builds real processor-backed tensor kwargs for multimodal models. def build_processor_kwargs( model_id: str, input_modalities: tuple[str, ...], configs: dict[str, dict[str, Any]] | None = None, + input_strategy: str = PROCESSOR_INPUT_STRATEGY, ) -> dict[str, Any]: - from transformers import AutoProcessor - configs = configs or {} token = os.environ.get("HF_TOKEN") - processor_kwargs: dict[str, Any] = {"trust_remote_code": True, "local_files_only": True} + processor_kwargs: dict[str, Any] = {"trust_remote_code": True} if token: processor_kwargs["token"] = token - processor = _build_gemma4_processor(model_id, configs) - if processor is None: - processor = AutoProcessor.from_pretrained(model_id, **processor_kwargs) + if input_strategy == PROCESSOR_INPUT_STRATEGY: + processor = _load_auto_processor(model_id, processor_kwargs) + elif input_strategy == "manual_gemma4_processor": + processor = _build_gemma4_processor(model_id, configs) or _load_auto_processor(model_id, processor_kwargs) + else: + raise ValueError(f"Unknown input strategy: {input_strategy}") + prompt_parts = [] call_kwargs: dict[str, Any] = {"return_tensors": "pt"} diff --git a/python/cactus/transpiler/Converter/models.py b/python/cactus/transpiler/Converter/models.py index 4441bce0a..ba5eaf062 100644 --- a/python/cactus/transpiler/Converter/models.py +++ b/python/cactus/transpiler/Converter/models.py @@ -7,18 +7,21 @@ import torch from . import input_utils as IU from ..ModelProfiles import models as MP_Models -from ..ModelProfiles import profiles as MP_Profiles -from transformers import AutoConfig, AutoModel, AutoModelForCausalLM, AutoModelForCTC, AutoModelForImageTextToText, AutoModelForSeq2SeqLM, AutoModelForSpeechSeq2Seq +from transformers import AutoModel, AutoModelForCausalLM, AutoModelForCTC, AutoModelForImageTextToText, AutoModelForSeq2SeqLM, AutoModelForSpeechSeq2Seq -default_model_ids: dict[str, MP_Models.ModelProfile] = { - "google/gemma-4-E2B": MP_Profiles.GEMMA4_E2B_PROFILE, - "openai/whisper-tiny": MP_Profiles.WHISPER_PROFILE, - "nvidia/parakeet-tdt-0.6b-v3": MP_Profiles.PARAKEET_PROFILE, - "LiquidAI/LFM2-VL-3B": MP_Profiles.LFM_VLM_PROFILE, - "Qwen/Qwen2.5-0.5B": MP_Profiles.QWEN2_5_0_5B_PROFILE, +LOAD_STRATEGIES = { + "image_text_to_text": (AutoModelForImageTextToText, AutoModelForCausalLM, AutoModel), + "image_text_to_text_strict": (AutoModelForImageTextToText, AutoModelForCausalLM), + "speech_seq2seq": (AutoModelForSpeechSeq2Seq, AutoModelForSeq2SeqLM, AutoModel), + "ctc": (AutoModelForCTC, AutoModel), + "causal_lm": (AutoModelForCausalLM, AutoModel), } +CACHE_INFERENCE_MODES = {IU.PREFILL_WITH_CACHE_MODE, IU.DECODE_WITH_CACHE_MODE} +DYNAMIC_CACHE_POLICY = "dynamic_cache" +DROP_MULTIMODAL_ON_DECODE_POLICY = "drop_multimodal_on_decode" + @dataclass(slots=True) class Input: @@ -27,6 +30,7 @@ class Input: modalities: tuple[str, ...] inference_mode: str + @dataclass(slots=True) class Model: name: str @@ -34,10 +38,6 @@ class Model: input: Input model: Any - @classmethod - def from_profile(mp: MP_Models.ModelProfile, input_modalities: tuple[str,...], model_id: str) -> "Model": - return create_model(mp, input_modalities, model_id) - def export(self, input: Input) -> "LayerMap": return export_(model=self, input=input) @@ -50,6 +50,7 @@ class TensorInstance: def from_tensor(cls, x: torch.Tensor) -> "TensorInstance": return cls(shape=[jsonable_shape_dim(dim) for dim in x.shape], dtype=str(x.dtype)) + @dataclass(slots=True) class Slice: start: Any @@ -60,6 +61,7 @@ class Slice: def from_slice(cls, x: slice) -> "Slice": return cls(start=jsonable(x.start), stop=jsonable(x.stop), step=jsonable(x.step)) + @dataclass(slots=True) class CacheLayerSpec: index: int @@ -68,25 +70,18 @@ class CacheLayerSpec: value_shape: tuple[int, ...] sliding_window: int | None = None + @dataclass(slots=True) class CacheSpec: layers: tuple[CacheLayerSpec, ...] config: Any past_sequence_length: int - dtype: torch.dtype - device: torch.device #Infers cache tensor shapes from the loaded model config and attention modules. @classmethod - def from_model( - cls, - model: torch.nn.Module, - batch_size: int, - past_sequence_length: int, - ) -> "CacheSpec": + def from_model(cls, model: torch.nn.Module, batch_size: int, past_sequence_length: int) -> "CacheSpec": config = getattr(model, "config", None) text_config = get_text_config(config) - dtype, device = model_dtype_and_device(model) num_key_value_heads = int(getattr(text_config, "num_key_value_heads", getattr(text_config, "num_attention_heads", 1))) num_attention_heads = int(getattr(text_config, "num_attention_heads", num_key_value_heads)) hidden_size = int(getattr(text_config, "hidden_size", num_attention_heads)) @@ -128,17 +123,16 @@ def from_model( layers=tuple(layers), config=config, past_sequence_length=past_sequence_length, - dtype=dtype, - device=device, ) #Creates zero-filled flat KV tensors that match this cache spec. - def empty_tensors(self) -> tuple[torch.Tensor, ...]: + def empty_tensors(self, dtype: torch.dtype = torch.float32, device: torch.device | None = None) -> tuple[torch.Tensor, ...]: tensors: list[torch.Tensor] = [] + device = device or torch.device("cpu") for layer in self.layers: - tensors.append(torch.zeros(layer.key_shape, dtype=self.dtype, device=self.device)) - tensors.append(torch.zeros(layer.value_shape, dtype=self.dtype, device=self.device)) + tensors.append(torch.zeros(layer.key_shape, dtype=dtype, device=device)) + tensors.append(torch.zeros(layer.value_shape, dtype=dtype, device=device)) return tuple(tensors) @@ -151,7 +145,8 @@ def to_dynamic_cache(self, flat_tensors: tuple[torch.Tensor, ...]): cache = DynamicCache(config=self.config) - for layer_spec, flat_index in zip(self.layers, range(0, len(flat_tensors), 2)): + for index, layer_spec in enumerate(self.layers): + flat_index = index * 2 layer = cache.layers[layer_spec.index] key = flat_tensors[flat_index] value = flat_tensors[flat_index + 1] @@ -170,11 +165,10 @@ def to_dynamic_cache(self, flat_tensors: tuple[torch.Tensor, ...]): #Wraps HF cache-mode forwards so torch.export only sees tensor inputs/outputs. class CacheExportWrapper(torch.nn.Module): #Stores the wrapped model and cache spec used to tensorize cache state. - def __init__(self, model: torch.nn.Module, cache_spec: CacheSpec, mode: str): + def __init__(self, model: torch.nn.Module, cache_spec: CacheSpec): super().__init__() self.model = model self.cache_spec = cache_spec - self.mode = mode #Runs the HF model with DynamicCache internally and returns logits plus flat cache tensors. def forward( @@ -213,7 +207,7 @@ def forward( model_kwargs = {key: value for key, value in model_kwargs.items() if value is not None} outputs = self.model(**filter_forward_kwargs(self.model, model_kwargs)) - return (primary_model_output(outputs), *flatten_dynamic_cache(outputs.past_key_values)) + return (primary_model_output(outputs), *flatten_dynamic_cache(getattr(outputs, "past_key_values", None))) #Serializable record for one exported FX graph node. @@ -231,19 +225,7 @@ class LayerRecord(BaseModel): #Builds a LayerRecord from a torch.fx.Node. @classmethod def from_node(cls, num: int, x: torch.fx.Node) -> "LayerRecord": - return cls( - index=num, - name=str(x.name), - node_type=str(x.op), - target=str(x.target), - args=jsonable(x.args), - kwargs=jsonable(x.kwargs), - users=[user.name for user in x.users], - tensor_output_meta=extract_tensor_meta(x), - module_stack=extract_module_stack(x), - ) - - + return cls(index=num, name=str(x.name), node_type=str(x.op), target=str(x.target), args=jsonable(x.args), kwargs=jsonable(x.kwargs), users=[user.name for user in x.users], tensor_output_meta=extract_tensor_meta(x), module_stack=extract_module_stack(x)) #Serializable top-level export IR container. class LayerMap(BaseModel): @@ -258,11 +240,11 @@ class LayerMap(BaseModel): def from_data(cls, x: torch.export.ExportedProgram, name: str, model_task: str, nodes_list: list[LayerRecord]) -> "LayerMap": return cls(model_name=name, task=model_task, graph_signature=repr(x.graph_signature), range_constants=repr(x.range_constraints), nodes=nodes_list) -"""#####################################Model Utils#####################################""" +###################################################### Model utility helpers!!!!! ########################################################################. #Converts symbolic shape dims to JSON-safe values without forcing torch guards. x def jsonable_shape_dim(x: Any) -> Any: - if type(x) is int: + if isinstance(x, int): return x return str(x) @@ -316,13 +298,13 @@ def get_text_config(config: Any) -> Any: return config -#Finds the dtype and device of the loaded model parameters. -def model_dtype_and_device(model: torch.nn.Module) -> tuple[torch.dtype, torch.device]: +#Finds the dtype of the loaded model parameters. +def model_dtype(model: torch.nn.Module) -> torch.dtype: try: param = next(model.parameters()) - return param.dtype, param.device + return param.dtype except StopIteration: - return torch.float32, torch.device("cpu") + return torch.float32 #Finds decoder layer modules across common HF model wrapper layouts. @@ -399,6 +381,12 @@ def cache_layer_types(text_config: Any) -> tuple[str, ...]: #Selects the main tensor output from a model output object. def primary_model_output(outputs: Any) -> torch.Tensor: + if isinstance(outputs, dict): + for key in ("logits", "last_hidden_state"): + output = outputs.get(key) + if output is not None: + return output + if hasattr(outputs, "logits") and outputs.logits is not None: return outputs.logits @@ -417,8 +405,20 @@ def flatten_dynamic_cache(cache: Any) -> tuple[torch.Tensor, ...]: return () flat: list[torch.Tensor] = [] + layers = getattr(cache, "layers", None) + + if layers is None and isinstance(cache, (tuple, list)): + for entry in cache: + if isinstance(entry, (tuple, list)) and len(entry) >= 2: + key, value = entry[:2] + if isinstance(key, torch.Tensor) and isinstance(value, torch.Tensor): + flat.extend((key, value)) + return tuple(flat) + + if layers is None: + return () - for layer in cache.layers: + for layer in layers: key = getattr(layer, "keys", None) value = getattr(layer, "values", None) @@ -444,21 +444,6 @@ def filter_forward_kwargs(model: torch.nn.Module, kwargs: dict[str, Any]) -> dic return {key: value for key, value in kwargs.items() if key in signature.parameters} -#Normalizes torch operator targets into stable ATen-style names. -def aten_name(target: Any) -> str: - schema = getattr(target, "_schema", None) - - if schema is not None and "::" in schema.name: - namespace, op = schema.name.split("::", 1) - overload = schema.overload_name if schema.overload_name else "default" - return f"{namespace}.{op}.{overload}" - - if hasattr(target, "name"): - return target.name - - return str(target) - - #Extracts tensor metadata from an exported FX node. def extract_tensor_meta(node: torch.fx.Node) -> Any | None: if "val" not in node.meta: @@ -479,7 +464,13 @@ def extract_module_stack(node: torch.fx.Node) -> Any | None: if isinstance(value, tuple) and len(value) >= 2: module_path = value[0] module_type = value[1] - out.append({"key": str(key), "module_path": str(module_path), "module_type": getattr(module_type, "__name__", str(module_type))}) + out.append( + { + "key": str(key), + "module_path": str(module_path), + "module_type": getattr(module_type, "__name__", str(module_type)), + } + ) else: out.append({"key": str(key), "value": repr(value)}) @@ -489,96 +480,58 @@ def extract_module_stack(node: torch.fx.Node) -> Any | None: #Builds representative model inputs for the requested modalities and inference mode. def build_input( mp: MP_Models.ModelProfile, - input_modalties: tuple[str, ...], + input_modalities: tuple[str, ...], model_id: str | None = None, inference_mode: str = "prefill_no_cache", ) -> Input | None: - if not all(modality in mp.supported_modalties for modality in input_modalties): + if not all(modality in mp.supported_modalties for modality in input_modalities): print("Requesting unsupported modalities") return None - loaded_files = IU.load_files(mp, model_id) + configs = IU.load_configs(mp, model_id) + modalities = tuple(input_modalities) + uses_dynamic_cache = DYNAMIC_CACHE_POLICY in mp.cache_policy - configs = IU._load_json_files(loaded_files or {}) - modalities = tuple(input_modalties) - batch_size = 1 - sequence_length = 8 - past_sequence_length = 8 - kwargs: dict[str, Any] = {} - - #Gemma-specific: use the real Gemma4 processor to create multimodal placeholder ids and position tensors. - if model_id is not None and "gemma" in model_id.lower() and any(m in modalities for m in ("vision", "audio")): + if model_id is not None and mp.input_strategy != IU.SYNTHETIC_INPUT_STRATEGY: try: - kwargs = IU.build_processor_kwargs(model_id, modalities, configs) return Input( args=(), - kwargs=kwargs, + kwargs=IU.build_processor_kwargs( + model_id=model_id, + input_modalities=modalities, + configs=configs, + input_strategy=mp.input_strategy, + ), modalities=modalities, inference_mode=inference_mode, ) except Exception as e: - print(f"Processor input build failed for {model_id}, falling back to synthetic inputs: {e}") - - if "text" in modalities: - if inference_mode == "decode_with_cache": - sequence_length = 1 - - kwargs["input_ids"] = IU._build_input_ids(batch_size, sequence_length) - kwargs["attention_mask"] = IU._build_attention_mask(batch_size, sequence_length) - - if "vision" in modalities: - kwargs["pixel_values"] = IU._build_pixel_values(batch_size, configs) - - if "audio" in modalities: - kwargs["input_features"] = IU._build_input_features(batch_size, configs) - - if "audio" in modalities and "text" in modalities and "vision" not in modalities and "speech" in inference_mode: - kwargs.pop("input_ids", None) - kwargs.pop("attention_mask", None) - kwargs["decoder_input_ids"] = torch.full( - (batch_size, 1), - IU._decoder_start_token_id(configs), - dtype=torch.long, - ) - - if "audio" in modalities and "text" not in modalities: - input_features = kwargs["input_features"] - kwargs["attention_mask"] = IU._build_attention_mask(batch_size, input_features.shape[-1]) - - if inference_mode in ("prefill_with_cache", "decode_with_cache"): - total_sequence_length = past_sequence_length + sequence_length - kwargs["attention_mask"] = IU._build_attention_mask(batch_size, total_sequence_length) - kwargs["past_key_values"] = IU._build_past_key_values(configs, batch_size, past_sequence_length) - kwargs["cache_position"] = torch.arange( - past_sequence_length, - past_sequence_length + sequence_length, - dtype=torch.long, - ) - kwargs["use_cache"] = True + print(f"{mp.input_strategy} input build failed for {model_id}, falling back to synthetic inputs: {e}") return Input( args=(), - kwargs=kwargs, + kwargs=IU.build_synthetic_kwargs( + modalities=modalities, + configs=configs, + inference_mode=inference_mode, + uses_dynamic_cache=uses_dynamic_cache, + ), modalities=modalities, inference_mode=inference_mode, ) -#Applies mode-specific input rewrites after the real model/config is loaded. -def prepare_input_for_export(model: torch.nn.Module, input: Input) -> Input: - if input.inference_mode == "decode_with_cache": - return build_decode_with_cache_input(model, input) - - return input - - #Builds one-token decode inputs and flat KV tensors for cache-mode exports. -def build_decode_with_cache_input(model: torch.nn.Module, input: Input) -> Input: +def build_decode_with_cache_input( + model: torch.nn.Module, + input: Input, + drop_multimodal: bool, +) -> Input: kwargs = dict(input.kwargs) token_key = "input_ids" if "input_ids" in kwargs else "decoder_input_ids" if token_key not in kwargs: - raise ValueError("decode_with_cache requires input_ids or decoder_input_ids") + raise ValueError(f"{IU.DECODE_WITH_CACHE_MODE} requires input_ids or decoder_input_ids") token_ids = kwargs[token_key] batch_size = int(token_ids.shape[0]) @@ -601,18 +554,14 @@ def build_decode_with_cache_input(model: torch.nn.Module, input: Input) -> Input dtype=torch.long, device=token_ids.device, ) - kwargs["past_key_values"] = cache_spec.empty_tensors() - - for multimodal_key in ( - "pixel_values", - "pixel_values_videos", - "input_features", - "input_features_mask", - "image_position_ids", - "video_position_ids", - "mm_token_type_ids", - ): - kwargs.pop(multimodal_key, None) + kwargs["past_key_values"] = cache_spec.empty_tensors( + dtype=model_dtype(model), + device=token_ids.device, + ) + + if drop_multimodal: + for multimodal_key in IU.MULTIMODAL_KEYS: + kwargs.pop(multimodal_key, None) return Input( args=input.args, @@ -624,33 +573,17 @@ def build_decode_with_cache_input(model: torch.nn.Module, input: Input) -> Input #Loads the best HF model class for the requested model profile. def load_model(model_id: str, mp: MP_Models.ModelProfile | None = None) -> torch.nn.Module: - - token = os.environ.get("HF_TOKEN") load_kwargs: dict[str, Any] = {"trust_remote_code": True} if token: load_kwargs["token"] = token - AutoConfig.from_pretrained(model_id, **load_kwargs) - - candidate_classes: list[Any] = [] - modalities = set(mp.supported_modalties if mp is not None else ()) - load_strategy = mp.load_strategy if mp is not None else "" - is_gemma4_profile = mp is not None and mp.model_profiles == "gemma4_e2b" - - if load_strategy == "image_text_to_text" or ("vision" in modalities and "text" in modalities): - candidate_classes.extend((AutoModelForImageTextToText, AutoModelForCausalLM)) - #Gemma-specific: avoid falling back to plain AutoModel, which loads Gemma4 weights with wrong prefixes. - if not is_gemma4_profile: - candidate_classes.append(AutoModel) - elif load_strategy == "speech_seq2seq" or ("audio" in modalities and "text" in modalities): - candidate_classes.extend((AutoModelForSpeechSeq2Seq, AutoModelForSeq2SeqLM, AutoModel)) - elif load_strategy == "ctc" or "audio" in modalities: - candidate_classes.extend((AutoModelForCTC, AutoModel)) - elif load_strategy == "causal_lm" or "text" in modalities: - candidate_classes.extend((AutoModelForCausalLM, AutoModel)) + if mp is None: + candidate_classes = (AutoModel,) + export_patches = () else: - candidate_classes.append(AutoModel) + candidate_classes = LOAD_STRATEGIES.get(mp.load_strategy, (AutoModel,)) + export_patches = mp.export_patches last_error: Exception | None = None seen_classes: set[Any] = set() @@ -660,16 +593,16 @@ def load_model(model_id: str, mp: MP_Models.ModelProfile | None = None) -> torch continue seen_classes.add(model_class) - for attempt_kwargs in _model_load_kwargs(load_kwargs): + for attempt_kwargs in IU.hf_load_kwargs(load_kwargs): try: model = model_class.from_pretrained(model_id, **attempt_kwargs) model.eval() - configure_model_for_export(model, should_use_cache=False) - - #Gemma-specific: patch Gemma4 audio masking so torch.export can trace it. - if mp is not None and mp.model_profiles == "gemma4_e2b" and "audio" in mp.supported_modalties: - patch_gemma4_audio_mask_for_export() + for patch in export_patches: + patch_fn = EXPORT_PATCHES.get(patch) + if patch_fn is None: + raise ValueError(f"Unknown export patch: {patch}") + patch_fn() return model except Exception as e: @@ -678,17 +611,6 @@ def load_model(model_id: str, mp: MP_Models.ModelProfile | None = None) -> torch raise RuntimeError(f"Unable to load model {model_id}") from last_error -#Returns normal load kwargs plus a local-files-only retry for cached HF models. -def _model_load_kwargs(load_kwargs: dict[str, Any]) -> tuple[dict[str, Any], ...]: - if load_kwargs.get("local_files_only") is True: - return (load_kwargs,) - - return ( - load_kwargs, - {**load_kwargs, "local_files_only": True}, - ) - - #Sets an attribute on a config object and common nested config objects. def _set_config_attr_recursive(config: Any, attr: str, value: Any, seen: set[int] | None = None) -> None: if config is None: @@ -722,30 +644,6 @@ def configure_model_for_export(model: torch.nn.Module, should_use_cache: bool) - _set_config_attr_recursive(config, "use_cache", should_use_cache) -#Checks whether a module forward can accept a specific keyword. -def _forward_accepts_kwarg(model: torch.nn.Module, key: str) -> bool: - try: - signature = inspect.signature(model.forward) - except (TypeError, ValueError): - return False - - if key in signature.parameters: - return True - - return any(param.kind == inspect.Parameter.VAR_KEYWORD for param in signature.parameters.values()) - - -#Builds the kwargs actually passed into torch.export. -def _export_kwargs(model: torch.nn.Module, input: Input) -> dict[str, Any]: - kwargs = filter_forward_kwargs(model, dict(input.kwargs)) - should_use_cache = input.inference_mode in ("prefill_with_cache", "decode_with_cache") - - if _forward_accepts_kwarg(model, "use_cache"): - kwargs["use_cache"] = should_use_cache - - return kwargs - - #Gemma-specific: replaces Gemma4's audio bidirectional mask helper with an exportable tensor implementation. def patch_gemma4_audio_mask_for_export() -> None: import transformers.models.gemma4.modeling_gemma4 as gemma4_modeling @@ -772,6 +670,11 @@ def exportable_bidirectional_mask(config, inputs_embeds, attention_mask=None, an gemma4_modeling.create_bidirectional_mask = exportable_bidirectional_mask +EXPORT_PATCHES = { + "gemma4_audio_mask": patch_gemma4_audio_mask_for_export, +} + + #Builds a loaded model bundle from profile, modalities, model id, and inference mode. def create_model( mp: MP_Models.ModelProfile, @@ -779,30 +682,20 @@ def create_model( model_id: str, inference_mode: str = "prefill_no_cache", ) -> Model: - name_ = model_id input_ = build_input(mp, input_modalities, model_id=model_id, inference_mode=inference_mode) if input_ is None: raise ValueError(f"Could not build input for modalities {input_modalities}") loaded_model = load_model(model_id, mp) - input_ = prepare_input_for_export(loaded_model, input_) - return Model(name=name_, model_profile=mp, input=input_, model=loaded_model) + if input_.inference_mode == IU.DECODE_WITH_CACHE_MODE and DYNAMIC_CACHE_POLICY in mp.cache_policy: + input_ = build_decode_with_cache_input( + model=loaded_model, + input=input_, + drop_multimodal=DROP_MULTIMODAL_ON_DECODE_POLICY in mp.cache_policy, + ) - -#Selects the raw model or cache tensorization wrapper for export. -def build_export_model(model: torch.nn.Module, input: Input) -> torch.nn.Module: - if input.inference_mode not in ("prefill_with_cache", "decode_with_cache"): - return model - - batch_size = infer_batch_size(input.kwargs) - past_sequence_length = infer_past_sequence_length(input) - cache_spec = CacheSpec.from_model( - model=model, - batch_size=batch_size, - past_sequence_length=past_sequence_length, - ) - return CacheExportWrapper(model=model, cache_spec=cache_spec, mode=input.inference_mode) + return Model(name=model_id, model_profile=mp, input=input_, model=loaded_model) #Infers batch size from the first tensor-shaped input. @@ -816,7 +709,7 @@ def infer_batch_size(kwargs: dict[str, Any]) -> int: #Infers the past/prompt sequence length represented by the export input. def infer_past_sequence_length(input: Input) -> int: - if input.inference_mode == "decode_with_cache" and "cache_position" in input.kwargs: + if input.inference_mode == IU.DECODE_WITH_CACHE_MODE and "cache_position" in input.kwargs: return int(input.kwargs["cache_position"][0].item()) if "input_ids" in input.kwargs: @@ -831,13 +724,25 @@ def infer_past_sequence_length(input: Input) -> int: return 0 - #Exports the prepared model and serializes its FX graph into a LayerMap. def export_(model: Model, input: Input) -> LayerMap: - should_use_cache = input.inference_mode in ("prefill_with_cache", "decode_with_cache") + should_use_cache = input.inference_mode in CACHE_INFERENCE_MODES configure_model_for_export(model.model, should_use_cache=should_use_cache) - export_model = build_export_model(model.model, input) - export_kwargs = _export_kwargs(export_model, input) + export_model = model.model + + if should_use_cache and DYNAMIC_CACHE_POLICY in model.model_profile.cache_policy: + export_model = CacheExportWrapper( + model=model.model, + cache_spec=CacheSpec.from_model( + model=model.model, + batch_size=infer_batch_size(input.kwargs), + past_sequence_length=infer_past_sequence_length(input), + ), + ) + + export_kwargs = dict(input.kwargs) + export_kwargs["use_cache"] = should_use_cache + export_kwargs = filter_forward_kwargs(export_model, export_kwargs) with torch.no_grad(): exported = torch.export.export( diff --git a/python/cactus/transpiler/ModelProfiles/profiles.py b/python/cactus/transpiler/ModelProfiles/profiles.py index 584326fe2..b03c1760a 100644 --- a/python/cactus/transpiler/ModelProfiles/profiles.py +++ b/python/cactus/transpiler/ModelProfiles/profiles.py @@ -37,7 +37,7 @@ supported_modalties=("audio", "vision", "text"), input_strategy="manual_gemma4_processor", export_patches=("gemma4_audio_mask",), - load_strategy="image_text_to_text", + load_strategy="image_text_to_text_strict", ) WHISPER_PROFILE = ModelProfile( @@ -126,12 +126,20 @@ load_strategy="causal_lm", ) -ALL_PROFILES = ( - GEMMA4_E2B_PROFILE, - WHISPER_PROFILE, - PARAKEET_PROFILE, - LFM_VLM_PROFILE, - QWEN2_5_0_5B_PROFILE, -) +MODEL_ID_MAP = { + "google/gemma-4-E2B": GEMMA4_E2B_PROFILE, + "openai/whisper-tiny": WHISPER_PROFILE, + "nvidia/parakeet-tdt-0.6b-v3": PARAKEET_PROFILE, + "LiquidAI/LFM2-VL-3B": LFM_VLM_PROFILE, + "Qwen/Qwen2.5-0.5B": QWEN2_5_0_5B_PROFILE, +} -MODEL_ID_MAP = {} +MODEL_ID_HINTS = ( + (("gemma",), GEMMA4_E2B_PROFILE), + (("whisper",), WHISPER_PROFILE), + (("parakeet",), PARAKEET_PROFILE), + (("lfm", "vl"), LFM_VLM_PROFILE), + (("lfm", "vision"), LFM_VLM_PROFILE), + (("qwen2.5",), QWEN2_5_0_5B_PROFILE), + (("qwen2_5",), QWEN2_5_0_5B_PROFILE), +) From c886f1cc31182987382a98803e6667e5b327f2f2 Mon Sep 17 00:00:00 2001 From: Parkirat Date: Tue, 21 Jul 2026 16:54:07 -0400 Subject: [PATCH 35/84] Cleaned out convert.py --- python/cactus/transpiler/Converter/convert.py | 43 +++---------------- python/cactus/transpiler/Fusions/models.py | 20 +++++++++ 2 files changed, 25 insertions(+), 38 deletions(-) diff --git a/python/cactus/transpiler/Converter/convert.py b/python/cactus/transpiler/Converter/convert.py index a0d8e4894..e2e8e1060 100644 --- a/python/cactus/transpiler/Converter/convert.py +++ b/python/cactus/transpiler/Converter/convert.py @@ -2,49 +2,16 @@ from . import models from ..ModelProfiles import models as MP_Models -from ..ModelProfiles import profiles as MP - - -#Selects the best known model profile for a Hugging Face model id. -def _infer_profile(model_id: str) -> MP_Models.ModelProfile: - model_id_lower = model_id.lower() - - for mapped_model_id, profile in MP.MODEL_ID_MAP.items(): - if model_id_lower == mapped_model_id.lower(): - return profile - - for required_terms, profile in MP.MODEL_ID_HINTS: - if all(term in model_id_lower for term in required_terms): - return profile - - raise ValueError(f"Could not infer model profile for {model_id}. Pass custom_profile instead.") - - -#Resolves where the exported LayerMap JSON should be written. -def _output_json_path(output_path: str | Path | None = None) -> Path: - if output_path is None: - path = Path(__file__).resolve().parent / "jsons" / "output.json" - else: - path = Path(output_path) - if path.suffix != ".json": - path = path / "output.json" - - path.parent.mkdir(parents=True, exist_ok=True) - return path - +from ..ModelProfiles import profiles as MP_Profiles #Runs the full converter pipeline and writes the exported LayerMap JSON. -def convert(model_id: str, input_modalities: tuple[str, ...], custom_profile: MP_Models.ModelProfile | None = None, inference_mode: str = "prefill_no_cache", output_path: str | Path | None = None) -> models.LayerMap: - profile = custom_profile or _infer_profile(model_id) +def convert(model_id: str, input_modalities: tuple[str, ...], output_path: str, custom_profile: MP_Models.ModelProfile | None = None, inference_mode: str = "prefill_no_cache") -> models.LayerMap: + profile = custom_profile or MP_Profiles.MODEL_ID_MAP[model_id] model = models.create_model(mp=profile, input_modalities=input_modalities, model_id=model_id, inference_mode=inference_mode) layer_map = model.export(model.input) - output_file = _output_json_path(output_path) - output_file.write_text(layer_map.model_dump_json(indent=4), encoding="utf-8") + Path(output_path).write_text(layer_map.model_dump_json(indent=4), encoding="utf-8") return layer_map if __name__ == "__main__": - convert( - model_id="google/gemma-4-E2B", - input_modalities=("text", "vision", "audio"), - ) + convert(model_id="google/gemma-4-E2B", input_modalities=("text", "vision", "audio")) diff --git a/python/cactus/transpiler/Fusions/models.py b/python/cactus/transpiler/Fusions/models.py index e69de29bb..33fb58c8b 100644 --- a/python/cactus/transpiler/Fusions/models.py +++ b/python/cactus/transpiler/Fusions/models.py @@ -0,0 +1,20 @@ +from dataclasses import dataclass +from typing import Any + + + + +@dataclass(slots=True) +class SynthNode: + attrs: dict[str, Any] + +@dataclass(slots=True) +class Route: + pass + + + +@dataclass(slots=True) +class SynthGraph: + source_node: SynthNode + route: Route \ No newline at end of file From 0bf9dc7641e2697d7b3d71b6b17015909c152241 Mon Sep 17 00:00:00 2001 From: Parkirat Date: Wed, 22 Jul 2026 16:42:49 -0400 Subject: [PATCH 36/84] Reorganizing + optimizing converter --- .../cactus/transpiler/Converter/constants.py | 19 + .../transpiler/Converter/input_utils.py | 125 ++++++ python/cactus/transpiler/Converter/models.py | 419 ++++++------------ 3 files changed, 286 insertions(+), 277 deletions(-) diff --git a/python/cactus/transpiler/Converter/constants.py b/python/cactus/transpiler/Converter/constants.py index e69de29bb..8b1d7ebcf 100644 --- a/python/cactus/transpiler/Converter/constants.py +++ b/python/cactus/transpiler/Converter/constants.py @@ -0,0 +1,19 @@ +from transformers import AutoModel, AutoModelForCausalLM, AutoModelForCTC, AutoModelForImageTextToText, AutoModelForSeq2SeqLM, AutoModelForSpeechSeq2Seq +from . import input_utils as IU + +#Names of component configs in config.txt (put in constants to allow future expansion to tuple to support different names) +TEXT_CONFIG_NAME = "text_config" +VISION_CONFIG_NAME = "vision_config" +AUDIO_CONFIG_NAME = "audio_config" + +LOAD_STRATEGIES = { + "image_text_to_text": (AutoModelForImageTextToText, AutoModelForCausalLM, AutoModel), + "image_text_to_text_strict": (AutoModelForImageTextToText, AutoModelForCausalLM), + "speech_seq2seq": (AutoModelForSpeechSeq2Seq, AutoModelForSeq2SeqLM, AutoModel), + "ctc": (AutoModelForCTC, AutoModel), + "causal_lm": (AutoModelForCausalLM, AutoModel), +} + +CACHE_INFERENCE_MODES = {IU.PREFILL_WITH_CACHE_MODE, IU.DECODE_WITH_CACHE_MODE} +DYNAMIC_CACHE_POLICY = IU.DYNAMIC_CACHE_POLICY +DROP_MULTIMODAL_ON_DECODE_POLICY = "drop_multimodal_on_decode" diff --git a/python/cactus/transpiler/Converter/input_utils.py b/python/cactus/transpiler/Converter/input_utils.py index fcf7352d9..b83219ea7 100644 --- a/python/cactus/transpiler/Converter/input_utils.py +++ b/python/cactus/transpiler/Converter/input_utils.py @@ -17,6 +17,7 @@ PROCESSOR_INPUT_STRATEGY = "processor" PREFILL_WITH_CACHE_MODE = "prefill_with_cache" DECODE_WITH_CACHE_MODE = "decode_with_cache" +DYNAMIC_CACHE_POLICY = "dynamic_cache" CONVERTER_JSON_DIR = Path(__file__).resolve().parent / "jsons" ASSET_DIR = Path(__file__).resolve().parents[2] / "assets" @@ -225,6 +226,130 @@ def build_synthetic_kwargs( return kwargs +#Builds representative model inputs for the requested modalities and inference mode. X +def build_input( + mp: MP_Models.ModelProfile, + input_modalities: tuple[str, ...], + input_cls: Any, + model_id: str | None = None, + inference_mode: str = "prefill_no_cache", +) -> Any | None: + if not all(modality in mp.supported_modalties for modality in input_modalities): + print("Requesting unsupported modalities") + return None + + configs = load_configs(mp, model_id) + modalities = tuple(input_modalities) + uses_dynamic_cache = DYNAMIC_CACHE_POLICY in mp.cache_policy + + if model_id is not None and mp.input_strategy != SYNTHETIC_INPUT_STRATEGY: + try: + return input_cls( + args=(), + kwargs=build_processor_kwargs( + model_id=model_id, + input_modalities=modalities, + configs=configs, + input_strategy=mp.input_strategy, + ), + modalities=modalities, + inference_mode=inference_mode, + ) + except Exception as e: + print(f"{mp.input_strategy} input build failed for {model_id}, falling back to synthetic inputs: {e}") + + return input_cls( + args=(), + kwargs=build_synthetic_kwargs( + modalities=modalities, + configs=configs, + inference_mode=inference_mode, + uses_dynamic_cache=uses_dynamic_cache, + ), + modalities=modalities, + inference_mode=inference_mode, + ) + + +#Builds one-token decode inputs and flat KV tensors for cache-mode exports. X +def build_decode_with_cache_input( + model: torch.nn.Module, + input_: Any, + input_cls: Any, + cache_spec_cls: Any, + model_dtype_fn: Any, + drop_multimodal: bool, +) -> Any: + kwargs = dict(input_.kwargs) + token_key = "input_ids" if "input_ids" in kwargs else "decoder_input_ids" + + if token_key not in kwargs: + raise ValueError(f"{DECODE_WITH_CACHE_MODE} requires input_ids or decoder_input_ids") + + token_ids = kwargs[token_key] + batch_size = int(token_ids.shape[0]) + past_sequence_length = int(token_ids.shape[1]) + cache_spec = cache_spec_cls.from_model( + model=model, + batch_size=batch_size, + past_sequence_length=past_sequence_length, + ) + + kwargs[token_key] = token_ids[:, -1:].clone() + kwargs["attention_mask"] = torch.ones( + (batch_size, past_sequence_length + 1), + dtype=torch.long, + device=token_ids.device, + ) + kwargs["cache_position"] = torch.arange( + past_sequence_length, + past_sequence_length + 1, + dtype=torch.long, + device=token_ids.device, + ) + kwargs["past_key_values"] = cache_spec.empty_tensors( + dtype=model_dtype_fn(model), + device=token_ids.device, + ) + + if drop_multimodal: + for multimodal_key in MULTIMODAL_KEYS: + kwargs.pop(multimodal_key, None) + + return input_cls( + args=input_.args, + kwargs=kwargs, + modalities=input_.modalities, + inference_mode=input_.inference_mode, + ) + + +#Infers batch size from the first tensor-shaped input. X +def infer_batch_size(kwargs: dict[str, Any]) -> int: + for value in kwargs.values(): + if isinstance(value, torch.Tensor) and value.ndim > 0: + return int(value.shape[0]) + + return 1 + + +#Infers the past/prompt sequence length represented by the export input. X +def infer_past_sequence_length(input_: Any) -> int: + if input_.inference_mode == DECODE_WITH_CACHE_MODE and "cache_position" in input_.kwargs: + return int(input_.kwargs["cache_position"][0].item()) + + if "input_ids" in input_.kwargs: + return int(input_.kwargs["input_ids"].shape[1]) + + if "decoder_input_ids" in input_.kwargs: + return int(input_.kwargs["decoder_input_ids"].shape[1]) + + if "attention_mask" in input_.kwargs: + return int(input_.kwargs["attention_mask"].shape[1]) + + return 0 + + #Loads the local image asset used for processor-backed multimodal exports. def _load_image_asset(): from PIL import Image diff --git a/python/cactus/transpiler/Converter/models.py b/python/cactus/transpiler/Converter/models.py index ba5eaf062..d66f98f14 100644 --- a/python/cactus/transpiler/Converter/models.py +++ b/python/cactus/transpiler/Converter/models.py @@ -1,26 +1,14 @@ -import json import os import inspect from dataclasses import dataclass from typing import Any from pydantic import BaseModel import torch +from transformers import AutoModel + from . import input_utils as IU from ..ModelProfiles import models as MP_Models - -from transformers import AutoModel, AutoModelForCausalLM, AutoModelForCTC, AutoModelForImageTextToText, AutoModelForSeq2SeqLM, AutoModelForSpeechSeq2Seq - -LOAD_STRATEGIES = { - "image_text_to_text": (AutoModelForImageTextToText, AutoModelForCausalLM, AutoModel), - "image_text_to_text_strict": (AutoModelForImageTextToText, AutoModelForCausalLM), - "speech_seq2seq": (AutoModelForSpeechSeq2Seq, AutoModelForSeq2SeqLM, AutoModel), - "ctc": (AutoModelForCTC, AutoModel), - "causal_lm": (AutoModelForCausalLM, AutoModel), -} - -CACHE_INFERENCE_MODES = {IU.PREFILL_WITH_CACHE_MODE, IU.DECODE_WITH_CACHE_MODE} -DYNAMIC_CACHE_POLICY = "dynamic_cache" -DROP_MULTIMODAL_ON_DECODE_POLICY = "drop_multimodal_on_decode" +from . import constants @dataclass(slots=True) @@ -30,7 +18,6 @@ class Input: modalities: tuple[str, ...] inference_mode: str - @dataclass(slots=True) class Model: name: str @@ -41,26 +28,6 @@ class Model: def export(self, input: Input) -> "LayerMap": return export_(model=self, input=input) -@dataclass(slots=True) -class TensorInstance: - shape: list[Any] - dtype: str - - @classmethod - def from_tensor(cls, x: torch.Tensor) -> "TensorInstance": - return cls(shape=[jsonable_shape_dim(dim) for dim in x.shape], dtype=str(x.dtype)) - - -@dataclass(slots=True) -class Slice: - start: Any - stop: Any - step: Any - - @classmethod - def from_slice(cls, x: slice) -> "Slice": - return cls(start=jsonable(x.start), stop=jsonable(x.stop), step=jsonable(x.step)) - @dataclass(slots=True) class CacheLayerSpec: @@ -77,100 +44,29 @@ class CacheSpec: config: Any past_sequence_length: int - #Infers cache tensor shapes from the loaded model config and attention modules. + #Delegates cache tensor shape inference to the utility helper below. X @classmethod def from_model(cls, model: torch.nn.Module, batch_size: int, past_sequence_length: int) -> "CacheSpec": - config = getattr(model, "config", None) - text_config = get_text_config(config) - num_key_value_heads = int(getattr(text_config, "num_key_value_heads", getattr(text_config, "num_attention_heads", 1))) - num_attention_heads = int(getattr(text_config, "num_attention_heads", num_key_value_heads)) - hidden_size = int(getattr(text_config, "hidden_size", num_attention_heads)) - head_dim = int(getattr(text_config, "head_dim", max(1, hidden_size // max(1, num_attention_heads)))) - sliding_window = getattr(text_config, "sliding_window", None) or getattr(text_config, "attention_chunk_size", None) - decoder_layers = model_decoder_layers(model) - layers: list[CacheLayerSpec] = [] - - for index, layer_type in enumerate(cache_layer_types(text_config)): - cache_sequence_length = past_sequence_length - layer_sliding_window = None - layer_num_key_value_heads = num_key_value_heads - layer_head_dim = head_dim - - if index < len(decoder_layers): - attention = attention_module(decoder_layers[index]) - layer_head_dim = int(getattr(attention, "head_dim", layer_head_dim)) - key_out_features = linear_out_features(getattr(attention, "k_proj", None)) - - if key_out_features is not None: - layer_num_key_value_heads = max(1, key_out_features // max(1, layer_head_dim)) - - if layer_type in ("sliding_attention", "chunked_attention") and sliding_window is not None: - layer_sliding_window = int(sliding_window) - cache_sequence_length = min(past_sequence_length, max(1, layer_sliding_window - 1)) - - cache_shape = (batch_size, layer_num_key_value_heads, cache_sequence_length, layer_head_dim) - layers.append( - CacheLayerSpec( - index=index, - layer_type=str(layer_type), - key_shape=cache_shape, - value_shape=cache_shape, - sliding_window=layer_sliding_window, - ) - ) - - return cls( - layers=tuple(layers), - config=config, - past_sequence_length=past_sequence_length, - ) + return cache_spec_from_model(cls, model, batch_size, past_sequence_length) - #Creates zero-filled flat KV tensors that match this cache spec. + #Delegates creation of zero-filled flat KV tensors to the utility helper below. X def empty_tensors(self, dtype: torch.dtype = torch.float32, device: torch.device | None = None) -> tuple[torch.Tensor, ...]: - tensors: list[torch.Tensor] = [] - device = device or torch.device("cpu") - - for layer in self.layers: - tensors.append(torch.zeros(layer.key_shape, dtype=dtype, device=device)) - tensors.append(torch.zeros(layer.value_shape, dtype=dtype, device=device)) - - return tuple(tensors) + return cache_spec_empty_tensors(self, dtype=dtype, device=device) - #Rebuilds an HF DynamicCache from flat exported tensor inputs. + #Delegates rebuilding an HF DynamicCache from flat exported tensor inputs. X def to_dynamic_cache(self, flat_tensors: tuple[torch.Tensor, ...]): - from transformers.cache_utils import DynamicCache + return cache_spec_to_dynamic_cache(self, flat_tensors) - if len(flat_tensors) != len(self.layers) * 2: - raise ValueError(f"Expected {len(self.layers) * 2} cache tensors, got {len(flat_tensors)}") - cache = DynamicCache(config=self.config) - - for index, layer_spec in enumerate(self.layers): - flat_index = index * 2 - layer = cache.layers[layer_spec.index] - key = flat_tensors[flat_index] - value = flat_tensors[flat_index + 1] - layer.keys = key - layer.values = value - layer.dtype = key.dtype - layer.device = key.device - layer.is_initialized = True - - if hasattr(layer, "cumulative_length"): - layer.cumulative_length = self.past_sequence_length - - return cache - - -#Wraps HF cache-mode forwards so torch.export only sees tensor inputs/outputs. +#Wraps HF cache-mode forwards so torch.export only sees tensor inputs/outputs. X class CacheExportWrapper(torch.nn.Module): - #Stores the wrapped model and cache spec used to tensorize cache state. + #Stores the wrapped model and cache spec used to tensorize cache state. X def __init__(self, model: torch.nn.Module, cache_spec: CacheSpec): super().__init__() self.model = model self.cache_spec = cache_spec - #Runs the HF model with DynamicCache internally and returns logits plus flat cache tensors. + #Runs the HF model with DynamicCache internally and returns logits plus flat cache tensors. X def forward( self, input_ids=None, @@ -242,22 +138,13 @@ def from_data(cls, x: torch.export.ExportedProgram, name: str, model_task: str, ###################################################### Model utility helpers!!!!! ########################################################################. -#Converts symbolic shape dims to JSON-safe values without forcing torch guards. x -def jsonable_shape_dim(x: Any) -> Any: - if isinstance(x, int): - return x - - return str(x) - - -#Recursively converts FX/export metadata into JSON-safe Python values. x +#Recursively converts FX/export metadata into JSON-safe Python values. def jsonable(x: Any) -> Any: if isinstance(x, torch.fx.Node): return {"node": x.name} if isinstance(x, torch.Tensor): - tensor = TensorInstance.from_tensor(x) - return {"shape": tensor.shape, "dtype": tensor.dtype} + return {"shape": x if isinstance(x, int) else str(x), "dtype": x.dtype} if isinstance(x, torch.Size): return list(x) @@ -269,8 +156,7 @@ def jsonable(x: Any) -> Any: return str(x) if isinstance(x, slice): - slice_ = Slice.from_slice(x) - return {"start": slice_.start, "stop": slice_.stop, "step": slice_.step} + return {"start":jsonable(x.start), "stop":jsonable(x.stop), "step":jsonable(x.step)} if isinstance(x, range): return list(x) @@ -287,7 +173,7 @@ def jsonable(x: Any) -> Any: return repr(x) -#Returns the text/decoder config for models that wrap language configs inside multimodal configs. +#Returns the text/decoder config for models that wrap language configs inside multimodal configs. X def get_text_config(config: Any) -> Any: if config is not None and hasattr(config, "get_text_config"): try: @@ -298,7 +184,7 @@ def get_text_config(config: Any) -> Any: return config -#Finds the dtype of the loaded model parameters. +#Finds the dtype of the loaded model parameters. X def model_dtype(model: torch.nn.Module) -> torch.dtype: try: param = next(model.parameters()) @@ -307,7 +193,7 @@ def model_dtype(model: torch.nn.Module) -> torch.dtype: return torch.float32 -#Finds decoder layer modules across common HF model wrapper layouts. +#Finds decoder layer modules across common HF model wrapper layouts. X def model_decoder_layers(model: torch.nn.Module) -> tuple[Any, ...]: candidates = ( ("model", "language_model", "layers"), @@ -332,7 +218,7 @@ def model_decoder_layers(model: torch.nn.Module) -> tuple[Any, ...]: return () -#Finds the attention module inside one decoder layer across common naming conventions. +#Finds the attention module inside one decoder layer across common naming conventions. X def attention_module(layer: Any) -> Any: for attr in ("self_attn", "attention", "attn"): value = getattr(layer, attr, None) @@ -342,7 +228,7 @@ def attention_module(layer: Any) -> Any: return layer -#Returns a linear-like module's output feature count when available. +#Returns a linear-like module's output feature count when available. X def linear_out_features(module: Any) -> int | None: if module is None: return None @@ -360,7 +246,7 @@ def linear_out_features(module: Any) -> int | None: return int(weight.shape[0]) -#Infers which decoder layers actually own KV cache entries. +#Infers which decoder layers actually own KV cache entries. X def cache_layer_types(text_config: Any) -> tuple[str, ...]: num_hidden_layers = int(getattr(text_config, "num_hidden_layers", 1)) layer_types = getattr(text_config, "layer_types", None) @@ -379,7 +265,97 @@ def cache_layer_types(text_config: Any) -> tuple[str, ...]: return layer_types -#Selects the main tensor output from a model output object. +#Infers cache tensor shapes from the loaded model config and attention modules. X +def cache_spec_from_model(cls: type[CacheSpec], model: torch.nn.Module, batch_size: int, past_sequence_length: int) -> CacheSpec: + config = getattr(model, "config", None) + text_config = get_text_config(config) + num_key_value_heads = int(getattr(text_config, "num_key_value_heads", getattr(text_config, "num_attention_heads", 1))) + num_attention_heads = int(getattr(text_config, "num_attention_heads", num_key_value_heads)) + hidden_size = int(getattr(text_config, "hidden_size", num_attention_heads)) + head_dim = int(getattr(text_config, "head_dim", max(1, hidden_size // max(1, num_attention_heads)))) + sliding_window = getattr(text_config, "sliding_window", None) or getattr(text_config, "attention_chunk_size", None) + decoder_layers = model_decoder_layers(model) + layers: list[CacheLayerSpec] = [] + + for index, layer_type in enumerate(cache_layer_types(text_config)): + cache_sequence_length = past_sequence_length + layer_sliding_window = None + layer_num_key_value_heads = num_key_value_heads + layer_head_dim = head_dim + + if index < len(decoder_layers): + attention = attention_module(decoder_layers[index]) + layer_head_dim = int(getattr(attention, "head_dim", layer_head_dim)) + key_out_features = linear_out_features(getattr(attention, "k_proj", None)) + + if key_out_features is not None: + layer_num_key_value_heads = max(1, key_out_features // max(1, layer_head_dim)) + + if layer_type in ("sliding_attention", "chunked_attention") and sliding_window is not None: + layer_sliding_window = int(sliding_window) + cache_sequence_length = min(past_sequence_length, max(1, layer_sliding_window - 1)) + + cache_shape = (batch_size, layer_num_key_value_heads, cache_sequence_length, layer_head_dim) + layers.append( + CacheLayerSpec( + index=index, + layer_type=str(layer_type), + key_shape=cache_shape, + value_shape=cache_shape, + sliding_window=layer_sliding_window, + ) + ) + + return cls( + layers=tuple(layers), + config=config, + past_sequence_length=past_sequence_length, + ) + + +#Creates zero-filled flat KV tensors that match this cache spec. X +def cache_spec_empty_tensors( + cache_spec: CacheSpec, + dtype: torch.dtype = torch.float32, + device: torch.device | None = None, +) -> tuple[torch.Tensor, ...]: + tensors: list[torch.Tensor] = [] + device = device or torch.device("cpu") + + for layer in cache_spec.layers: + tensors.append(torch.zeros(layer.key_shape, dtype=dtype, device=device)) + tensors.append(torch.zeros(layer.value_shape, dtype=dtype, device=device)) + + return tuple(tensors) + + +#Rebuilds an HF DynamicCache from flat exported tensor inputs. X +def cache_spec_to_dynamic_cache(cache_spec: CacheSpec, flat_tensors: tuple[torch.Tensor, ...]): + from transformers.cache_utils import DynamicCache + + if len(flat_tensors) != len(cache_spec.layers) * 2: + raise ValueError(f"Expected {len(cache_spec.layers) * 2} cache tensors, got {len(flat_tensors)}") + + cache = DynamicCache(config=cache_spec.config) + + for index, layer_spec in enumerate(cache_spec.layers): + flat_index = index * 2 + layer = cache.layers[layer_spec.index] + key = flat_tensors[flat_index] + value = flat_tensors[flat_index + 1] + layer.keys = key + layer.values = value + layer.dtype = key.dtype + layer.device = key.device + layer.is_initialized = True + + if hasattr(layer, "cumulative_length"): + layer.cumulative_length = cache_spec.past_sequence_length + + return cache + + +#Selects the main tensor output from a model output object. X def primary_model_output(outputs: Any) -> torch.Tensor: if isinstance(outputs, dict): for key in ("logits", "last_hidden_state"): @@ -399,7 +375,7 @@ def primary_model_output(outputs: Any) -> torch.Tensor: raise ValueError("Could not find a tensor output to export") -#Flattens an HF DynamicCache into key/value tensor outputs. +#Flattens an HF DynamicCache into key/value tensor outputs. X def flatten_dynamic_cache(cache: Any) -> tuple[torch.Tensor, ...]: if cache is None: return () @@ -431,7 +407,7 @@ def flatten_dynamic_cache(cache: Any) -> tuple[torch.Tensor, ...]: return tuple(flat) -#Drops kwargs unsupported by a module's forward signature. +#Drops kwargs unsupported by a module's forward signature. X def filter_forward_kwargs(model: torch.nn.Module, kwargs: dict[str, Any]) -> dict[str, Any]: try: signature = inspect.signature(model.forward) @@ -444,7 +420,7 @@ def filter_forward_kwargs(model: torch.nn.Module, kwargs: dict[str, Any]) -> dic return {key: value for key, value in kwargs.items() if key in signature.parameters} -#Extracts tensor metadata from an exported FX node. +#Extracts tensor metadata from an exported FX node. X def extract_tensor_meta(node: torch.fx.Node) -> Any | None: if "val" not in node.meta: return None @@ -452,7 +428,7 @@ def extract_tensor_meta(node: torch.fx.Node) -> Any | None: return jsonable(node.meta["val"]) -#Extracts module stack metadata from an exported FX node. +#Extracts module stack metadata from an exported FX node. X def extract_module_stack(node: torch.fx.Node) -> Any | None: stack = node.meta.get("nn_module_stack", None) if stack is None: @@ -477,101 +453,7 @@ def extract_module_stack(node: torch.fx.Node) -> Any | None: return out -#Builds representative model inputs for the requested modalities and inference mode. -def build_input( - mp: MP_Models.ModelProfile, - input_modalities: tuple[str, ...], - model_id: str | None = None, - inference_mode: str = "prefill_no_cache", -) -> Input | None: - if not all(modality in mp.supported_modalties for modality in input_modalities): - print("Requesting unsupported modalities") - return None - - configs = IU.load_configs(mp, model_id) - modalities = tuple(input_modalities) - uses_dynamic_cache = DYNAMIC_CACHE_POLICY in mp.cache_policy - - if model_id is not None and mp.input_strategy != IU.SYNTHETIC_INPUT_STRATEGY: - try: - return Input( - args=(), - kwargs=IU.build_processor_kwargs( - model_id=model_id, - input_modalities=modalities, - configs=configs, - input_strategy=mp.input_strategy, - ), - modalities=modalities, - inference_mode=inference_mode, - ) - except Exception as e: - print(f"{mp.input_strategy} input build failed for {model_id}, falling back to synthetic inputs: {e}") - - return Input( - args=(), - kwargs=IU.build_synthetic_kwargs( - modalities=modalities, - configs=configs, - inference_mode=inference_mode, - uses_dynamic_cache=uses_dynamic_cache, - ), - modalities=modalities, - inference_mode=inference_mode, - ) - - -#Builds one-token decode inputs and flat KV tensors for cache-mode exports. -def build_decode_with_cache_input( - model: torch.nn.Module, - input: Input, - drop_multimodal: bool, -) -> Input: - kwargs = dict(input.kwargs) - token_key = "input_ids" if "input_ids" in kwargs else "decoder_input_ids" - - if token_key not in kwargs: - raise ValueError(f"{IU.DECODE_WITH_CACHE_MODE} requires input_ids or decoder_input_ids") - - token_ids = kwargs[token_key] - batch_size = int(token_ids.shape[0]) - past_sequence_length = int(token_ids.shape[1]) - cache_spec = CacheSpec.from_model( - model=model, - batch_size=batch_size, - past_sequence_length=past_sequence_length, - ) - - kwargs[token_key] = token_ids[:, -1:].clone() - kwargs["attention_mask"] = torch.ones( - (batch_size, past_sequence_length + 1), - dtype=torch.long, - device=token_ids.device, - ) - kwargs["cache_position"] = torch.arange( - past_sequence_length, - past_sequence_length + 1, - dtype=torch.long, - device=token_ids.device, - ) - kwargs["past_key_values"] = cache_spec.empty_tensors( - dtype=model_dtype(model), - device=token_ids.device, - ) - - if drop_multimodal: - for multimodal_key in IU.MULTIMODAL_KEYS: - kwargs.pop(multimodal_key, None) - - return Input( - args=input.args, - kwargs=kwargs, - modalities=input.modalities, - inference_mode=input.inference_mode, - ) - - -#Loads the best HF model class for the requested model profile. +#Loads the best HF model class for the requested model profile. X def load_model(model_id: str, mp: MP_Models.ModelProfile | None = None) -> torch.nn.Module: token = os.environ.get("HF_TOKEN") load_kwargs: dict[str, Any] = {"trust_remote_code": True} @@ -582,7 +464,7 @@ def load_model(model_id: str, mp: MP_Models.ModelProfile | None = None) -> torch candidate_classes = (AutoModel,) export_patches = () else: - candidate_classes = LOAD_STRATEGIES.get(mp.load_strategy, (AutoModel,)) + candidate_classes = constants.LOAD_STRATEGIES.get(mp.load_strategy, (AutoModel,)) export_patches = mp.export_patches last_error: Exception | None = None @@ -611,7 +493,7 @@ def load_model(model_id: str, mp: MP_Models.ModelProfile | None = None) -> torch raise RuntimeError(f"Unable to load model {model_id}") from last_error -#Sets an attribute on a config object and common nested config objects. +#Sets an attribute on a config object and common nested config objects. X def _set_config_attr_recursive(config: Any, attr: str, value: Any, seen: set[int] | None = None) -> None: if config is None: return @@ -637,18 +519,18 @@ def _set_config_attr_recursive(config: Any, attr: str, value: Any, seen: set[int _set_config_attr_recursive(getattr(config, child_name, None), attr, value, seen) -#Configures cache-related model settings before torch.export. +#Configures cache-related model settings before torch.export. X def configure_model_for_export(model: torch.nn.Module, should_use_cache: bool) -> None: for config_name in ("config", "generation_config"): config = getattr(model, config_name, None) _set_config_attr_recursive(config, "use_cache", should_use_cache) -#Gemma-specific: replaces Gemma4's audio bidirectional mask helper with an exportable tensor implementation. +#Gemma-specific: replaces Gemma4's audio bidirectional mask helper with an exportable tensor implementation. X def patch_gemma4_audio_mask_for_export() -> None: import transformers.models.gemma4.modeling_gemma4 as gemma4_modeling - #Gemma-specific: computes Gemma4 audio tower local bidirectional attention masks without vmap/proxy issues. + #Gemma-specific: computes Gemma4 audio tower local bidirectional attention masks without vmap/proxy issues. X def exportable_bidirectional_mask(config, inputs_embeds, attention_mask=None, and_mask_function=None, **kwargs): batch_size, seq_len = inputs_embeds.shape[:2] device = inputs_embeds.device @@ -675,68 +557,51 @@ def exportable_bidirectional_mask(config, inputs_embeds, attention_mask=None, an } -#Builds a loaded model bundle from profile, modalities, model id, and inference mode. +#Builds a loaded model bundle from profile, modalities, model id, and inference mode. X def create_model( mp: MP_Models.ModelProfile, input_modalities: tuple[str, ...], model_id: str, inference_mode: str = "prefill_no_cache", ) -> Model: - input_ = build_input(mp, input_modalities, model_id=model_id, inference_mode=inference_mode) + input_ = IU.build_input( + mp=mp, + input_modalities=input_modalities, + input_cls=Input, + model_id=model_id, + inference_mode=inference_mode, + ) if input_ is None: raise ValueError(f"Could not build input for modalities {input_modalities}") loaded_model = load_model(model_id, mp) - if input_.inference_mode == IU.DECODE_WITH_CACHE_MODE and DYNAMIC_CACHE_POLICY in mp.cache_policy: - input_ = build_decode_with_cache_input( + if input_.inference_mode == IU.DECODE_WITH_CACHE_MODE and constants.DYNAMIC_CACHE_POLICY in mp.cache_policy: + input_ = IU.build_decode_with_cache_input( model=loaded_model, - input=input_, - drop_multimodal=DROP_MULTIMODAL_ON_DECODE_POLICY in mp.cache_policy, + input_=input_, + input_cls=Input, + cache_spec_cls=CacheSpec, + model_dtype_fn=model_dtype, + drop_multimodal=constants.DROP_MULTIMODAL_ON_DECODE_POLICY in mp.cache_policy, ) return Model(name=model_id, model_profile=mp, input=input_, model=loaded_model) -#Infers batch size from the first tensor-shaped input. -def infer_batch_size(kwargs: dict[str, Any]) -> int: - for value in kwargs.values(): - if isinstance(value, torch.Tensor) and value.ndim > 0: - return int(value.shape[0]) - - return 1 - - -#Infers the past/prompt sequence length represented by the export input. -def infer_past_sequence_length(input: Input) -> int: - if input.inference_mode == IU.DECODE_WITH_CACHE_MODE and "cache_position" in input.kwargs: - return int(input.kwargs["cache_position"][0].item()) - - if "input_ids" in input.kwargs: - return int(input.kwargs["input_ids"].shape[1]) - - if "decoder_input_ids" in input.kwargs: - return int(input.kwargs["decoder_input_ids"].shape[1]) - - if "attention_mask" in input.kwargs: - return int(input.kwargs["attention_mask"].shape[1]) - - return 0 - - -#Exports the prepared model and serializes its FX graph into a LayerMap. +#Exports the prepared model and serializes its FX graph into a LayerMap. X def export_(model: Model, input: Input) -> LayerMap: - should_use_cache = input.inference_mode in CACHE_INFERENCE_MODES + should_use_cache = input.inference_mode in constants.CACHE_INFERENCE_MODES configure_model_for_export(model.model, should_use_cache=should_use_cache) export_model = model.model - if should_use_cache and DYNAMIC_CACHE_POLICY in model.model_profile.cache_policy: + if should_use_cache and constants.DYNAMIC_CACHE_POLICY in model.model_profile.cache_policy: export_model = CacheExportWrapper( model=model.model, cache_spec=CacheSpec.from_model( model=model.model, - batch_size=infer_batch_size(input.kwargs), - past_sequence_length=infer_past_sequence_length(input), + batch_size=IU.infer_batch_size(input.kwargs), + past_sequence_length=IU.infer_past_sequence_length(input), ), ) From af0797fa86959440a16e4271b4614735973aa22d Mon Sep 17 00:00:00 2001 From: Parkirat Date: Wed, 22 Jul 2026 16:44:02 -0400 Subject: [PATCH 37/84] Updated required files for Gemma --- python/cactus/transpiler/ModelProfiles/profiles.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/python/cactus/transpiler/ModelProfiles/profiles.py b/python/cactus/transpiler/ModelProfiles/profiles.py index b03c1760a..73ee9e677 100644 --- a/python/cactus/transpiler/ModelProfiles/profiles.py +++ b/python/cactus/transpiler/ModelProfiles/profiles.py @@ -21,7 +21,7 @@ inference_type=GEMMA4_INFERENCE_PATTERNS, cache_type=("attention_kv", "sliding_window_attention_kv"), cache_policy=("dynamic_cache", "drop_multimodal_on_decode", "shared_kv_layers"), - files=("config.json", "generation_config.json", "processor_config.json", "tokenizer_config.json"), + files=("config.json", "generation_config.json", "processor_config.json", "tokenizer_config.json", "merges.txt"), fusion_fields=( "generic", "embedding", From 052a2ba969e7e5adc0d455d2d20374dd9ea7f48b Mon Sep 17 00:00:00 2001 From: Parkirat Date: Thu, 23 Jul 2026 11:38:00 -0400 Subject: [PATCH 38/84] Converter reorganization --- .../transpiler/Converter/cache_utils.py | 362 +++++++++++++ .../transpiler/Converter/export_patches.py | 25 + python/cactus/transpiler/Converter/models.py | 488 ++---------------- 3 files changed, 422 insertions(+), 453 deletions(-) create mode 100644 python/cactus/transpiler/Converter/cache_utils.py create mode 100644 python/cactus/transpiler/Converter/export_patches.py diff --git a/python/cactus/transpiler/Converter/cache_utils.py b/python/cactus/transpiler/Converter/cache_utils.py new file mode 100644 index 000000000..a3ddfa03c --- /dev/null +++ b/python/cactus/transpiler/Converter/cache_utils.py @@ -0,0 +1,362 @@ +import inspect +from dataclasses import dataclass +from typing import Any + +import torch + + +@dataclass(slots=True) +class CacheLayerSpec: + index: int + layer_type: str + key_shape: tuple[int, ...] + value_shape: tuple[int, ...] + sliding_window: int | None = None + + +@dataclass(slots=True) +class CacheSpec: + layers: tuple[CacheLayerSpec, ...] + config: Any + past_sequence_length: int + + #Delegates cache tensor shape inference to the utility helper below. X + @classmethod + def from_model(cls, model: torch.nn.Module, batch_size: int, past_sequence_length: int) -> "CacheSpec": + return cache_spec_from_model(cls, model, batch_size, past_sequence_length) + + #Delegates creation of zero-filled flat KV tensors to the utility helper below. X + def empty_tensors(self, dtype: torch.dtype = torch.float32, device: torch.device | None = None) -> tuple[torch.Tensor, ...]: + return cache_spec_empty_tensors(self, dtype=dtype, device=device) + + #Delegates rebuilding an HF DynamicCache from flat exported tensor inputs. X + def to_dynamic_cache(self, flat_tensors: tuple[torch.Tensor, ...]): + return cache_spec_to_dynamic_cache(self, flat_tensors) + + +#Wraps HF cache-mode forwards so torch.export only sees tensor inputs/outputs. X +class CacheExportWrapper(torch.nn.Module): + #Stores the wrapped model and cache spec used to tensorize cache state. X + def __init__(self, model: torch.nn.Module, cache_spec: CacheSpec): + super().__init__() + self.model = model + self.cache_spec = cache_spec + + #Runs the HF model with DynamicCache internally and returns logits plus flat cache tensors. X + def forward( + self, + input_ids=None, + pixel_values=None, + pixel_values_videos=None, + input_features=None, + attention_mask=None, + input_features_mask=None, + position_ids=None, + image_position_ids=None, + video_position_ids=None, + cache_position=None, + past_key_values=None, + mm_token_type_ids=None, + inputs_embeds=None, + ): + cache = self.cache_spec.to_dynamic_cache(past_key_values) if past_key_values is not None else None + model_kwargs = { + "input_ids": input_ids, + "pixel_values": pixel_values, + "pixel_values_videos": pixel_values_videos, + "input_features": input_features, + "attention_mask": attention_mask, + "input_features_mask": input_features_mask, + "position_ids": position_ids, + "image_position_ids": image_position_ids, + "video_position_ids": video_position_ids, + "cache_position": cache_position, + "past_key_values": cache, + "mm_token_type_ids": mm_token_type_ids, + "inputs_embeds": inputs_embeds, + "use_cache": True, + } + model_kwargs = {key: value for key, value in model_kwargs.items() if value is not None} + outputs = self.model(**filter_forward_kwargs(self.model, model_kwargs)) + + return (primary_model_output(outputs), *flatten_dynamic_cache(getattr(outputs, "past_key_values", None))) + + +#Returns the text/decoder config for models that wrap language configs inside multimodal configs. X +def get_text_config(config: Any) -> Any: + if config is not None and hasattr(config, "get_text_config"): + try: + return config.get_text_config(decoder=True) + except TypeError: + return config.get_text_config() + + return config + + +#Finds the dtype of the loaded model parameters. X +def model_dtype(model: torch.nn.Module) -> torch.dtype: + try: + param = next(model.parameters()) + return param.dtype + except StopIteration: + return torch.float32 + + +#Finds decoder layer modules across common HF model wrapper layouts. X +def model_decoder_layers(model: torch.nn.Module) -> tuple[Any, ...]: + candidates = ( + ("model", "language_model", "layers"), + ("language_model", "layers"), + ("model", "decoder", "layers"), + ("decoder", "layers"), + ("model", "layers"), + ("layers",), + ) + + for path in candidates: + value: Any = model + + for attr in path: + value = getattr(value, attr, None) + if value is None: + break + + if value is not None: + return tuple(value) + + return () + + +#Finds the attention module inside one decoder layer across common naming conventions. X +def attention_module(layer: Any) -> Any: + for attr in ("self_attn", "attention", "attn"): + value = getattr(layer, attr, None) + if value is not None: + return value + + return layer + + +#Returns a linear-like module's output feature count when available. X +def linear_out_features(module: Any) -> int | None: + if module is None: + return None + + if hasattr(module, "out_features"): + return int(module.out_features) + + weight = getattr(module, "weight", None) + if weight is None and hasattr(module, "linear"): + weight = getattr(module.linear, "weight", None) + + if weight is None: + return None + + return int(weight.shape[0]) + + +#Infers which decoder layers actually own KV cache entries. X +def cache_layer_types(text_config: Any) -> tuple[str, ...]: + num_hidden_layers = int(getattr(text_config, "num_hidden_layers", 1)) + layer_types = getattr(text_config, "layer_types", None) + + if layer_types is None: + sliding_window = getattr(text_config, "sliding_window", None) or getattr(text_config, "attention_chunk_size", None) + default_layer_type = "sliding_attention" if sliding_window is not None else "full_attention" + layer_types = tuple(default_layer_type for _ in range(num_hidden_layers)) + else: + layer_types = tuple(str(layer_type) for layer_type in layer_types) + + num_kv_shared_layers = int(getattr(text_config, "num_kv_shared_layers", 0) or 0) + if num_kv_shared_layers > 0: + layer_types = layer_types[: -num_kv_shared_layers] + + return layer_types + + +#Infers cache tensor shapes from the loaded model config and attention modules. X +def cache_spec_from_model(cls: type[CacheSpec], model: torch.nn.Module, batch_size: int, past_sequence_length: int) -> CacheSpec: + config = getattr(model, "config", None) + text_config = get_text_config(config) + num_key_value_heads = int(getattr(text_config, "num_key_value_heads", getattr(text_config, "num_attention_heads", 1))) + num_attention_heads = int(getattr(text_config, "num_attention_heads", num_key_value_heads)) + hidden_size = int(getattr(text_config, "hidden_size", num_attention_heads)) + head_dim = int(getattr(text_config, "head_dim", max(1, hidden_size // max(1, num_attention_heads)))) + sliding_window = getattr(text_config, "sliding_window", None) or getattr(text_config, "attention_chunk_size", None) + decoder_layers = model_decoder_layers(model) + layers: list[CacheLayerSpec] = [] + + for index, layer_type in enumerate(cache_layer_types(text_config)): + cache_sequence_length = past_sequence_length + layer_sliding_window = None + layer_num_key_value_heads = num_key_value_heads + layer_head_dim = head_dim + + if index < len(decoder_layers): + attention = attention_module(decoder_layers[index]) + layer_head_dim = int(getattr(attention, "head_dim", layer_head_dim)) + key_out_features = linear_out_features(getattr(attention, "k_proj", None)) + + if key_out_features is not None: + layer_num_key_value_heads = max(1, key_out_features // max(1, layer_head_dim)) + + if layer_type in ("sliding_attention", "chunked_attention") and sliding_window is not None: + layer_sliding_window = int(sliding_window) + cache_sequence_length = min(past_sequence_length, max(1, layer_sliding_window - 1)) + + cache_shape = (batch_size, layer_num_key_value_heads, cache_sequence_length, layer_head_dim) + layers.append( + CacheLayerSpec( + index=index, + layer_type=str(layer_type), + key_shape=cache_shape, + value_shape=cache_shape, + sliding_window=layer_sliding_window, + ) + ) + + return cls( + layers=tuple(layers), + config=config, + past_sequence_length=past_sequence_length, + ) + + +#Creates zero-filled flat KV tensors that match this cache spec. X +def cache_spec_empty_tensors( + cache_spec: CacheSpec, + dtype: torch.dtype = torch.float32, + device: torch.device | None = None, +) -> tuple[torch.Tensor, ...]: + tensors: list[torch.Tensor] = [] + device = device or torch.device("cpu") + + for layer in cache_spec.layers: + tensors.append(torch.zeros(layer.key_shape, dtype=dtype, device=device)) + tensors.append(torch.zeros(layer.value_shape, dtype=dtype, device=device)) + + return tuple(tensors) + + +#Rebuilds an HF DynamicCache from flat exported tensor inputs. X +def cache_spec_to_dynamic_cache(cache_spec: CacheSpec, flat_tensors: tuple[torch.Tensor, ...]): + from transformers.cache_utils import DynamicCache + + if len(flat_tensors) != len(cache_spec.layers) * 2: + raise ValueError(f"Expected {len(cache_spec.layers) * 2} cache tensors, got {len(flat_tensors)}") + + cache = DynamicCache(config=cache_spec.config) + + for index, layer_spec in enumerate(cache_spec.layers): + flat_index = index * 2 + layer = cache.layers[layer_spec.index] + key = flat_tensors[flat_index] + value = flat_tensors[flat_index + 1] + layer.keys = key + layer.values = value + layer.dtype = key.dtype + layer.device = key.device + layer.is_initialized = True + + if hasattr(layer, "cumulative_length"): + layer.cumulative_length = cache_spec.past_sequence_length + + return cache + + +#Selects the main tensor output from a model output object. X +def primary_model_output(outputs: Any) -> torch.Tensor: + if isinstance(outputs, dict): + for key in ("logits", "last_hidden_state"): + output = outputs.get(key) + if output is not None: + return output + + if hasattr(outputs, "logits") and outputs.logits is not None: + return outputs.logits + + if hasattr(outputs, "last_hidden_state") and outputs.last_hidden_state is not None: + return outputs.last_hidden_state + + if isinstance(outputs, (tuple, list)) and len(outputs) > 0: + return outputs[0] + + raise ValueError("Could not find a tensor output to export") + + +#Flattens an HF DynamicCache into key/value tensor outputs. X +def flatten_dynamic_cache(cache: Any) -> tuple[torch.Tensor, ...]: + if cache is None: + return () + + flat: list[torch.Tensor] = [] + layers = getattr(cache, "layers", None) + + if layers is None and isinstance(cache, (tuple, list)): + for entry in cache: + if isinstance(entry, (tuple, list)) and len(entry) >= 2: + key, value = entry[:2] + if isinstance(key, torch.Tensor) and isinstance(value, torch.Tensor): + flat.extend((key, value)) + return tuple(flat) + + if layers is None: + return () + + for layer in layers: + key = getattr(layer, "keys", None) + value = getattr(layer, "values", None) + + if key is None or value is None: + continue + + flat.append(key) + flat.append(value) + + return tuple(flat) + + +#Drops kwargs unsupported by a module's forward signature. X +def filter_forward_kwargs(model: torch.nn.Module, kwargs: dict[str, Any]) -> dict[str, Any]: + try: + signature = inspect.signature(model.forward) + except (TypeError, ValueError): + return kwargs + + if any(param.kind == inspect.Parameter.VAR_KEYWORD for param in signature.parameters.values()): + return kwargs + + return {key: value for key, value in kwargs.items() if key in signature.parameters} + + +#Sets an attribute on a config object and common nested config objects. X +def _set_config_attr_recursive(config: Any, attr: str, value: Any, seen: set[int] | None = None) -> None: + if config is None: + return + + seen = seen or set() + config_id = id(config) + if config_id in seen: + return + + seen.add(config_id) + + if hasattr(config, attr): + setattr(config, attr, value) + + for child_name in ( + "text_config", + "vision_config", + "audio_config", + "encoder", + "decoder", + "model_config", + ): + _set_config_attr_recursive(getattr(config, child_name, None), attr, value, seen) + + +#Configures cache-related model settings before torch.export. X +def configure_model_for_export(model: torch.nn.Module, should_use_cache: bool) -> None: + for config_name in ("config", "generation_config"): + config = getattr(model, config_name, None) + _set_config_attr_recursive(config, "use_cache", should_use_cache) diff --git a/python/cactus/transpiler/Converter/export_patches.py b/python/cactus/transpiler/Converter/export_patches.py new file mode 100644 index 000000000..32338820c --- /dev/null +++ b/python/cactus/transpiler/Converter/export_patches.py @@ -0,0 +1,25 @@ +import torch + +def patch_gemma4_audio_mask_for_export() -> None: + import transformers.models.gemma4.modeling_gemma4 as gemma4_modeling + + #Gemma-specific: computes Gemma4 audio tower local bidirectional attention masks without vmap/proxy issues. X + def exportable_bidirectional_mask(config, inputs_embeds, attention_mask=None, and_mask_function=None, **kwargs): + batch_size, seq_len = inputs_embeds.shape[:2] + device = inputs_embeds.device + q_idx = torch.arange(seq_len, device=device).view(1, 1, seq_len, 1) + kv_idx = torch.arange(seq_len, device=device).view(1, 1, 1, seq_len) + + left_window_size = getattr(config, "attention_context_left", seq_len) - 1 + right_window_size = getattr(config, "attention_context_right", 0) + distance = q_idx - kv_idx + left_mask = (distance >= 0) & (distance < left_window_size) + right_mask = (distance < 0) & (-distance < right_window_size) + mask = left_mask | right_mask + + if attention_mask is not None: + mask = mask & attention_mask[:, None, None, :].bool() + + return mask.expand(batch_size, 1, seq_len, seq_len) + + gemma4_modeling.create_bidirectional_mask = exportable_bidirectional_mask \ No newline at end of file diff --git a/python/cactus/transpiler/Converter/models.py b/python/cactus/transpiler/Converter/models.py index d66f98f14..048c9b1d8 100644 --- a/python/cactus/transpiler/Converter/models.py +++ b/python/cactus/transpiler/Converter/models.py @@ -1,5 +1,4 @@ import os -import inspect from dataclasses import dataclass from typing import Any from pydantic import BaseModel @@ -7,9 +6,13 @@ from transformers import AutoModel from . import input_utils as IU +from . import cache_utils as CU +from . import export_patches as EP from ..ModelProfiles import models as MP_Models from . import constants +token = os.environ.get("HF_TOKEN") + @dataclass(slots=True) class Input: @@ -29,83 +32,6 @@ def export(self, input: Input) -> "LayerMap": return export_(model=self, input=input) -@dataclass(slots=True) -class CacheLayerSpec: - index: int - layer_type: str - key_shape: tuple[int, ...] - value_shape: tuple[int, ...] - sliding_window: int | None = None - - -@dataclass(slots=True) -class CacheSpec: - layers: tuple[CacheLayerSpec, ...] - config: Any - past_sequence_length: int - - #Delegates cache tensor shape inference to the utility helper below. X - @classmethod - def from_model(cls, model: torch.nn.Module, batch_size: int, past_sequence_length: int) -> "CacheSpec": - return cache_spec_from_model(cls, model, batch_size, past_sequence_length) - - #Delegates creation of zero-filled flat KV tensors to the utility helper below. X - def empty_tensors(self, dtype: torch.dtype = torch.float32, device: torch.device | None = None) -> tuple[torch.Tensor, ...]: - return cache_spec_empty_tensors(self, dtype=dtype, device=device) - - #Delegates rebuilding an HF DynamicCache from flat exported tensor inputs. X - def to_dynamic_cache(self, flat_tensors: tuple[torch.Tensor, ...]): - return cache_spec_to_dynamic_cache(self, flat_tensors) - - -#Wraps HF cache-mode forwards so torch.export only sees tensor inputs/outputs. X -class CacheExportWrapper(torch.nn.Module): - #Stores the wrapped model and cache spec used to tensorize cache state. X - def __init__(self, model: torch.nn.Module, cache_spec: CacheSpec): - super().__init__() - self.model = model - self.cache_spec = cache_spec - - #Runs the HF model with DynamicCache internally and returns logits plus flat cache tensors. X - def forward( - self, - input_ids=None, - pixel_values=None, - pixel_values_videos=None, - input_features=None, - attention_mask=None, - input_features_mask=None, - position_ids=None, - image_position_ids=None, - video_position_ids=None, - cache_position=None, - past_key_values=None, - mm_token_type_ids=None, - inputs_embeds=None, - ): - cache = self.cache_spec.to_dynamic_cache(past_key_values) if past_key_values is not None else None - model_kwargs = { - "input_ids": input_ids, - "pixel_values": pixel_values, - "pixel_values_videos": pixel_values_videos, - "input_features": input_features, - "attention_mask": attention_mask, - "input_features_mask": input_features_mask, - "position_ids": position_ids, - "image_position_ids": image_position_ids, - "video_position_ids": video_position_ids, - "cache_position": cache_position, - "past_key_values": cache, - "mm_token_type_ids": mm_token_type_ids, - "inputs_embeds": inputs_embeds, - "use_cache": True, - } - model_kwargs = {key: value for key, value in model_kwargs.items() if value is not None} - outputs = self.model(**filter_forward_kwargs(self.model, model_kwargs)) - - return (primary_model_output(outputs), *flatten_dynamic_cache(getattr(outputs, "past_key_values", None))) - - #Serializable record for one exported FX graph node. class LayerRecord(BaseModel): index: int @@ -173,293 +99,31 @@ def jsonable(x: Any) -> Any: return repr(x) -#Returns the text/decoder config for models that wrap language configs inside multimodal configs. X -def get_text_config(config: Any) -> Any: - if config is not None and hasattr(config, "get_text_config"): - try: - return config.get_text_config(decoder=True) - except TypeError: - return config.get_text_config() - - return config - - -#Finds the dtype of the loaded model parameters. X -def model_dtype(model: torch.nn.Module) -> torch.dtype: - try: - param = next(model.parameters()) - return param.dtype - except StopIteration: - return torch.float32 - - -#Finds decoder layer modules across common HF model wrapper layouts. X -def model_decoder_layers(model: torch.nn.Module) -> tuple[Any, ...]: - candidates = ( - ("model", "language_model", "layers"), - ("language_model", "layers"), - ("model", "decoder", "layers"), - ("decoder", "layers"), - ("model", "layers"), - ("layers",), - ) - - for path in candidates: - value: Any = model - - for attr in path: - value = getattr(value, attr, None) - if value is None: - break - - if value is not None: - return tuple(value) - - return () - - -#Finds the attention module inside one decoder layer across common naming conventions. X -def attention_module(layer: Any) -> Any: - for attr in ("self_attn", "attention", "attn"): - value = getattr(layer, attr, None) - if value is not None: - return value - - return layer - - -#Returns a linear-like module's output feature count when available. X -def linear_out_features(module: Any) -> int | None: - if module is None: - return None - - if hasattr(module, "out_features"): - return int(module.out_features) - - weight = getattr(module, "weight", None) - if weight is None and hasattr(module, "linear"): - weight = getattr(module.linear, "weight", None) - - if weight is None: - return None - - return int(weight.shape[0]) - - -#Infers which decoder layers actually own KV cache entries. X -def cache_layer_types(text_config: Any) -> tuple[str, ...]: - num_hidden_layers = int(getattr(text_config, "num_hidden_layers", 1)) - layer_types = getattr(text_config, "layer_types", None) - - if layer_types is None: - sliding_window = getattr(text_config, "sliding_window", None) or getattr(text_config, "attention_chunk_size", None) - default_layer_type = "sliding_attention" if sliding_window is not None else "full_attention" - layer_types = tuple(default_layer_type for _ in range(num_hidden_layers)) - else: - layer_types = tuple(str(layer_type) for layer_type in layer_types) - - num_kv_shared_layers = int(getattr(text_config, "num_kv_shared_layers", 0) or 0) - if num_kv_shared_layers > 0: - layer_types = layer_types[: -num_kv_shared_layers] - - return layer_types - - -#Infers cache tensor shapes from the loaded model config and attention modules. X -def cache_spec_from_model(cls: type[CacheSpec], model: torch.nn.Module, batch_size: int, past_sequence_length: int) -> CacheSpec: - config = getattr(model, "config", None) - text_config = get_text_config(config) - num_key_value_heads = int(getattr(text_config, "num_key_value_heads", getattr(text_config, "num_attention_heads", 1))) - num_attention_heads = int(getattr(text_config, "num_attention_heads", num_key_value_heads)) - hidden_size = int(getattr(text_config, "hidden_size", num_attention_heads)) - head_dim = int(getattr(text_config, "head_dim", max(1, hidden_size // max(1, num_attention_heads)))) - sliding_window = getattr(text_config, "sliding_window", None) or getattr(text_config, "attention_chunk_size", None) - decoder_layers = model_decoder_layers(model) - layers: list[CacheLayerSpec] = [] - - for index, layer_type in enumerate(cache_layer_types(text_config)): - cache_sequence_length = past_sequence_length - layer_sliding_window = None - layer_num_key_value_heads = num_key_value_heads - layer_head_dim = head_dim - - if index < len(decoder_layers): - attention = attention_module(decoder_layers[index]) - layer_head_dim = int(getattr(attention, "head_dim", layer_head_dim)) - key_out_features = linear_out_features(getattr(attention, "k_proj", None)) - - if key_out_features is not None: - layer_num_key_value_heads = max(1, key_out_features // max(1, layer_head_dim)) - - if layer_type in ("sliding_attention", "chunked_attention") and sliding_window is not None: - layer_sliding_window = int(sliding_window) - cache_sequence_length = min(past_sequence_length, max(1, layer_sliding_window - 1)) - - cache_shape = (batch_size, layer_num_key_value_heads, cache_sequence_length, layer_head_dim) - layers.append( - CacheLayerSpec( - index=index, - layer_type=str(layer_type), - key_shape=cache_shape, - value_shape=cache_shape, - sliding_window=layer_sliding_window, - ) - ) - - return cls( - layers=tuple(layers), - config=config, - past_sequence_length=past_sequence_length, - ) - - -#Creates zero-filled flat KV tensors that match this cache spec. X -def cache_spec_empty_tensors( - cache_spec: CacheSpec, - dtype: torch.dtype = torch.float32, - device: torch.device | None = None, -) -> tuple[torch.Tensor, ...]: - tensors: list[torch.Tensor] = [] - device = device or torch.device("cpu") - - for layer in cache_spec.layers: - tensors.append(torch.zeros(layer.key_shape, dtype=dtype, device=device)) - tensors.append(torch.zeros(layer.value_shape, dtype=dtype, device=device)) - - return tuple(tensors) - - -#Rebuilds an HF DynamicCache from flat exported tensor inputs. X -def cache_spec_to_dynamic_cache(cache_spec: CacheSpec, flat_tensors: tuple[torch.Tensor, ...]): - from transformers.cache_utils import DynamicCache - - if len(flat_tensors) != len(cache_spec.layers) * 2: - raise ValueError(f"Expected {len(cache_spec.layers) * 2} cache tensors, got {len(flat_tensors)}") - - cache = DynamicCache(config=cache_spec.config) - - for index, layer_spec in enumerate(cache_spec.layers): - flat_index = index * 2 - layer = cache.layers[layer_spec.index] - key = flat_tensors[flat_index] - value = flat_tensors[flat_index + 1] - layer.keys = key - layer.values = value - layer.dtype = key.dtype - layer.device = key.device - layer.is_initialized = True - - if hasattr(layer, "cumulative_length"): - layer.cumulative_length = cache_spec.past_sequence_length - - return cache - - -#Selects the main tensor output from a model output object. X -def primary_model_output(outputs: Any) -> torch.Tensor: - if isinstance(outputs, dict): - for key in ("logits", "last_hidden_state"): - output = outputs.get(key) - if output is not None: - return output - - if hasattr(outputs, "logits") and outputs.logits is not None: - return outputs.logits - - if hasattr(outputs, "last_hidden_state") and outputs.last_hidden_state is not None: - return outputs.last_hidden_state - - if isinstance(outputs, (tuple, list)) and len(outputs) > 0: - return outputs[0] - - raise ValueError("Could not find a tensor output to export") - - -#Flattens an HF DynamicCache into key/value tensor outputs. X -def flatten_dynamic_cache(cache: Any) -> tuple[torch.Tensor, ...]: - if cache is None: - return () - - flat: list[torch.Tensor] = [] - layers = getattr(cache, "layers", None) - - if layers is None and isinstance(cache, (tuple, list)): - for entry in cache: - if isinstance(entry, (tuple, list)) and len(entry) >= 2: - key, value = entry[:2] - if isinstance(key, torch.Tensor) and isinstance(value, torch.Tensor): - flat.extend((key, value)) - return tuple(flat) - - if layers is None: - return () - - for layer in layers: - key = getattr(layer, "keys", None) - value = getattr(layer, "values", None) - - if key is None or value is None: - continue - - flat.append(key) - flat.append(value) - - return tuple(flat) - - -#Drops kwargs unsupported by a module's forward signature. X -def filter_forward_kwargs(model: torch.nn.Module, kwargs: dict[str, Any]) -> dict[str, Any]: - try: - signature = inspect.signature(model.forward) - except (TypeError, ValueError): - return kwargs - - if any(param.kind == inspect.Parameter.VAR_KEYWORD for param in signature.parameters.values()): - return kwargs - - return {key: value for key, value in kwargs.items() if key in signature.parameters} - - -#Extracts tensor metadata from an exported FX node. X +#Extracts tensor metadata from an exported FX node. def extract_tensor_meta(node: torch.fx.Node) -> Any | None: - if "val" not in node.meta: - return None + return jsonable(node.meta["val"]) if "val" not in node.meta else None - return jsonable(node.meta["val"]) - -#Extracts module stack metadata from an exported FX node. X +#Extracts module stack metadata from an exported FX node. def extract_module_stack(node: torch.fx.Node) -> Any | None: stack = node.meta.get("nn_module_stack", None) + out = [] + if stack is None: return None - out = [] - for key, value in stack.items(): if isinstance(value, tuple) and len(value) >= 2: - module_path = value[0] - module_type = value[1] - out.append( - { - "key": str(key), - "module_path": str(module_path), - "module_type": getattr(module_type, "__name__", str(module_type)), - } - ) + module_path, module_type = value[0], value[1] + out.append({"key": str(key), "module_path": str(module_path), "module_type": getattr(module_type, "__name__", str(module_type))}) else: out.append({"key": str(key), "value": repr(value)}) return out -#Loads the best HF model class for the requested model profile. X +#Loads the best HF model class for the requested model profile. def load_model(model_id: str, mp: MP_Models.ModelProfile | None = None) -> torch.nn.Module: - token = os.environ.get("HF_TOKEN") - load_kwargs: dict[str, Any] = {"trust_remote_code": True} - if token: - load_kwargs["token"] = token - if mp is None: candidate_classes = (AutoModel,) export_patches = () @@ -467,122 +131,50 @@ def load_model(model_id: str, mp: MP_Models.ModelProfile | None = None) -> torch candidate_classes = constants.LOAD_STRATEGIES.get(mp.load_strategy, (AutoModel,)) export_patches = mp.export_patches - last_error: Exception | None = None seen_classes: set[Any] = set() for model_class in candidate_classes: if model_class in seen_classes: continue - seen_classes.add(model_class) - - for attempt_kwargs in IU.hf_load_kwargs(load_kwargs): - try: - model = model_class.from_pretrained(model_id, **attempt_kwargs) - model.eval() - - for patch in export_patches: - patch_fn = EXPORT_PATCHES.get(patch) - if patch_fn is None: - raise ValueError(f"Unknown export patch: {patch}") - patch_fn() - - return model - except Exception as e: - last_error = e - - raise RuntimeError(f"Unable to load model {model_id}") from last_error - -#Sets an attribute on a config object and common nested config objects. X -def _set_config_attr_recursive(config: Any, attr: str, value: Any, seen: set[int] | None = None) -> None: - if config is None: - return - - seen = seen or set() - config_id = id(config) - if config_id in seen: - return - - seen.add(config_id) - - if hasattr(config, attr): - setattr(config, attr, value) - - for child_name in ( - "text_config", - "vision_config", - "audio_config", - "encoder", - "decoder", - "model_config", - ): - _set_config_attr_recursive(getattr(config, child_name, None), attr, value, seen) - - -#Configures cache-related model settings before torch.export. X -def configure_model_for_export(model: torch.nn.Module, should_use_cache: bool) -> None: - for config_name in ("config", "generation_config"): - config = getattr(model, config_name, None) - _set_config_attr_recursive(config, "use_cache", should_use_cache) - - -#Gemma-specific: replaces Gemma4's audio bidirectional mask helper with an exportable tensor implementation. X -def patch_gemma4_audio_mask_for_export() -> None: - import transformers.models.gemma4.modeling_gemma4 as gemma4_modeling - - #Gemma-specific: computes Gemma4 audio tower local bidirectional attention masks without vmap/proxy issues. X - def exportable_bidirectional_mask(config, inputs_embeds, attention_mask=None, and_mask_function=None, **kwargs): - batch_size, seq_len = inputs_embeds.shape[:2] - device = inputs_embeds.device - q_idx = torch.arange(seq_len, device=device).view(1, 1, seq_len, 1) - kv_idx = torch.arange(seq_len, device=device).view(1, 1, 1, seq_len) + seen_classes.add(model_class) + try: + model = model_class.from_pretrained(model_id) + model.eval() + except: + pass - left_window_size = getattr(config, "attention_context_left", seq_len) - 1 - right_window_size = getattr(config, "attention_context_right", 0) - distance = q_idx - kv_idx - left_mask = (distance >= 0) & (distance < left_window_size) - right_mask = (distance < 0) & (-distance < right_window_size) - mask = left_mask | right_mask + for patch in export_patches: + patch_fn = EXPORT_PATCHES[patch] + patch_fn() - if attention_mask is not None: - mask = mask & attention_mask[:, None, None, :].bool() + return model - return mask.expand(batch_size, 1, seq_len, seq_len) + raise RuntimeError(f"Unable to load model {model_id}") - gemma4_modeling.create_bidirectional_mask = exportable_bidirectional_mask EXPORT_PATCHES = { - "gemma4_audio_mask": patch_gemma4_audio_mask_for_export, + "gemma4_audio_mask": EP.patch_gemma4_audio_mask_for_export, } #Builds a loaded model bundle from profile, modalities, model id, and inference mode. X -def create_model( - mp: MP_Models.ModelProfile, - input_modalities: tuple[str, ...], - model_id: str, - inference_mode: str = "prefill_no_cache", -) -> Model: - input_ = IU.build_input( - mp=mp, - input_modalities=input_modalities, - input_cls=Input, - model_id=model_id, - inference_mode=inference_mode, - ) +def create_model(mp: MP_Models.ModelProfile, input_modalities: tuple[str, ...], model_id: str, inference_mode: str = "prefill_no_cache") -> Model: + input_ = IU.build_input(mp, input_modalities, Input, model_id, inference_mode) if input_ is None: raise ValueError(f"Could not build input for modalities {input_modalities}") loaded_model = load_model(model_id, mp) + #Build a decode with cache input only if decoder/ctc head supports cache if input_.inference_mode == IU.DECODE_WITH_CACHE_MODE and constants.DYNAMIC_CACHE_POLICY in mp.cache_policy: input_ = IU.build_decode_with_cache_input( model=loaded_model, input_=input_, input_cls=Input, - cache_spec_cls=CacheSpec, - model_dtype_fn=model_dtype, + cache_spec_cls=CU.CacheSpec, + model_dtype_fn=CU.model_dtype, drop_multimodal=constants.DROP_MULTIMODAL_ON_DECODE_POLICY in mp.cache_policy, ) @@ -592,13 +184,13 @@ def create_model( #Exports the prepared model and serializes its FX graph into a LayerMap. X def export_(model: Model, input: Input) -> LayerMap: should_use_cache = input.inference_mode in constants.CACHE_INFERENCE_MODES - configure_model_for_export(model.model, should_use_cache=should_use_cache) + CU.configure_model_for_export(model.model, should_use_cache=should_use_cache) export_model = model.model if should_use_cache and constants.DYNAMIC_CACHE_POLICY in model.model_profile.cache_policy: - export_model = CacheExportWrapper( + export_model = CU.CacheExportWrapper( model=model.model, - cache_spec=CacheSpec.from_model( + cache_spec=CU.CacheSpec.from_model( model=model.model, batch_size=IU.infer_batch_size(input.kwargs), past_sequence_length=IU.infer_past_sequence_length(input), @@ -607,21 +199,11 @@ def export_(model: Model, input: Input) -> LayerMap: export_kwargs = dict(input.kwargs) export_kwargs["use_cache"] = should_use_cache - export_kwargs = filter_forward_kwargs(export_model, export_kwargs) + export_kwargs = CU.filter_forward_kwargs(export_model, export_kwargs) with torch.no_grad(): - exported = torch.export.export( - export_model, - args=input.args, - kwargs=export_kwargs, - strict=False, - ).run_decompositions() + exported = torch.export.export(export_model, args=input.args, kwargs=export_kwargs, strict=False).run_decompositions() graph = exported.graph_module.graph if hasattr(exported, "graph_module") else exported.graph records = [LayerRecord.from_node(i, node) for i, node in enumerate(graph.nodes)] - return LayerMap.from_data( - x=exported, - name=model.name, - model_task=input.inference_mode, - nodes_list=records, - ) + return LayerMap.from_data(x=exported, name=model.name, model_task=input.inference_mode, nodes_list=records) From 198f46c0965d6968a1bc87b16738169cc2745438 Mon Sep 17 00:00:00 2001 From: Parkirat Date: Thu, 23 Jul 2026 16:03:42 -0400 Subject: [PATCH 39/84] Code refactoring --- .../transpiler/Converter/cache_utils.py | 60 ++- .../cactus/transpiler/Converter/constants.py | 55 ++- .../transpiler/Converter/input_processor.py | 65 +++ .../transpiler/Converter/input_utils.py | 413 ++---------------- python/cactus/transpiler/Converter/models.py | 29 +- .../{export_patches.py => overrides.py} | 0 6 files changed, 199 insertions(+), 423 deletions(-) create mode 100644 python/cactus/transpiler/Converter/input_processor.py rename python/cactus/transpiler/Converter/{export_patches.py => overrides.py} (100%) diff --git a/python/cactus/transpiler/Converter/cache_utils.py b/python/cactus/transpiler/Converter/cache_utils.py index a3ddfa03c..22e35666e 100644 --- a/python/cactus/transpiler/Converter/cache_utils.py +++ b/python/cactus/transpiler/Converter/cache_utils.py @@ -1,7 +1,6 @@ import inspect from dataclasses import dataclass from typing import Any - import torch @@ -20,29 +19,34 @@ class CacheSpec: config: Any past_sequence_length: int - #Delegates cache tensor shape inference to the utility helper below. X + #How: delegates model/cache inspection to cache_spec_from_model. + #Why: used by models.export_ and input_utils.build_decode_with_cache_input to create cache shapes for export. @classmethod def from_model(cls, model: torch.nn.Module, batch_size: int, past_sequence_length: int) -> "CacheSpec": return cache_spec_from_model(cls, model, batch_size, past_sequence_length) - #Delegates creation of zero-filled flat KV tensors to the utility helper below. X + #How: delegates zero tensor creation to cache_spec_empty_tensors. + #Why: used when building decode inputs so torch.export receives real tensor placeholders for past_key_values. def empty_tensors(self, dtype: torch.dtype = torch.float32, device: torch.device | None = None) -> tuple[torch.Tensor, ...]: return cache_spec_empty_tensors(self, dtype=dtype, device=device) - #Delegates rebuilding an HF DynamicCache from flat exported tensor inputs. X + #How: delegates flat tensor to DynamicCache reconstruction to cache_spec_to_dynamic_cache. + #Why: used inside CacheExportWrapper.forward because HF models expect DynamicCache, not a flat tuple. def to_dynamic_cache(self, flat_tensors: tuple[torch.Tensor, ...]): return cache_spec_to_dynamic_cache(self, flat_tensors) -#Wraps HF cache-mode forwards so torch.export only sees tensor inputs/outputs. X +#Wraps HF cache-mode forwards so torch.export only sees tensor inputs/outputs. class CacheExportWrapper(torch.nn.Module): - #Stores the wrapped model and cache spec used to tensorize cache state. X + #How: stores the original HF model plus the CacheSpec that describes flattened cache tensors. + #Why: constructed by models.export_ whenever a cache-mode export needs tensor-only inputs and outputs. def __init__(self, model: torch.nn.Module, cache_spec: CacheSpec): super().__init__() self.model = model self.cache_spec = cache_spec - #Runs the HF model with DynamicCache internally and returns logits plus flat cache tensors. X + #How: rebuilds DynamicCache from flat tensors, calls the model, then flattens the returned cache. + #Why: torch.export cannot directly expose HF cache objects, so this is the bridge for decode/prefill cache graphs. def forward( self, input_ids=None, @@ -82,7 +86,8 @@ def forward( return (primary_model_output(outputs), *flatten_dynamic_cache(getattr(outputs, "past_key_values", None))) -#Returns the text/decoder config for models that wrap language configs inside multimodal configs. X +#How: asks HF configs for their decoder/text sub-config when available, otherwise returns the original config. +#Why: cache_spec_from_model needs decoder fields even when the top-level config is multimodal, like Gemma. def get_text_config(config: Any) -> Any: if config is not None and hasattr(config, "get_text_config"): try: @@ -93,7 +98,8 @@ def get_text_config(config: Any) -> Any: return config -#Finds the dtype of the loaded model parameters. X +#How: reads the dtype from the first model parameter and falls back to float32 for parameterless modules. +#Why: input_utils.build_decode_with_cache_input uses this so synthetic cache tensors match the model dtype. def model_dtype(model: torch.nn.Module) -> torch.dtype: try: param = next(model.parameters()) @@ -102,7 +108,8 @@ def model_dtype(model: torch.nn.Module) -> torch.dtype: return torch.float32 -#Finds decoder layer modules across common HF model wrapper layouts. X +#How: tries common attribute paths until it finds a layers collection. +#Why: cache_spec_from_model uses real layers to validate or refine cache head dimensions beyond config defaults. def model_decoder_layers(model: torch.nn.Module) -> tuple[Any, ...]: candidates = ( ("model", "language_model", "layers"), @@ -127,7 +134,8 @@ def model_decoder_layers(model: torch.nn.Module) -> tuple[Any, ...]: return () -#Finds the attention module inside one decoder layer across common naming conventions. X +#How: checks common attention attribute names and falls back to the layer itself. +#Why: cache_spec_from_model needs the attention module to inspect fields like head_dim and k_proj. def attention_module(layer: Any) -> Any: for attr in ("self_attn", "attention", "attn"): value = getattr(layer, attr, None) @@ -137,7 +145,8 @@ def attention_module(layer: Any) -> Any: return layer -#Returns a linear-like module's output feature count when available. X +#How: reads out_features or the first dimension of a weight tensor from linear-like modules. +#Why: cache_spec_from_model uses k_proj output size to infer per-layer KV head count. def linear_out_features(module: Any) -> int | None: if module is None: return None @@ -155,7 +164,8 @@ def linear_out_features(module: Any) -> int | None: return int(weight.shape[0]) -#Infers which decoder layers actually own KV cache entries. X +#How: reads layer_types from config, otherwise derives full/sliding attention from window settings. +#Why: cache_spec_from_model uses this to know how many cache tensor pairs to create and which are sliding-window. def cache_layer_types(text_config: Any) -> tuple[str, ...]: num_hidden_layers = int(getattr(text_config, "num_hidden_layers", 1)) layer_types = getattr(text_config, "layer_types", None) @@ -174,7 +184,8 @@ def cache_layer_types(text_config: Any) -> tuple[str, ...]: return layer_types -#Infers cache tensor shapes from the loaded model config and attention modules. X +#How: combines config defaults with actual attention module inspection to create one CacheLayerSpec per cache layer. +#Why: CacheSpec.from_model calls this before decode/cache export so past_key_values tensors have correct shapes. def cache_spec_from_model(cls: type[CacheSpec], model: torch.nn.Module, batch_size: int, past_sequence_length: int) -> CacheSpec: config = getattr(model, "config", None) text_config = get_text_config(config) @@ -222,7 +233,8 @@ def cache_spec_from_model(cls: type[CacheSpec], model: torch.nn.Module, batch_si ) -#Creates zero-filled flat KV tensors that match this cache spec. X +#How: allocates key and value zero tensors for every CacheLayerSpec, in flat key0/value0/key1/value1 order. +#Why: CacheSpec.empty_tensors uses this to create placeholder cache inputs for torch.export decode traces. def cache_spec_empty_tensors( cache_spec: CacheSpec, dtype: torch.dtype = torch.float32, @@ -238,7 +250,8 @@ def cache_spec_empty_tensors( return tuple(tensors) -#Rebuilds an HF DynamicCache from flat exported tensor inputs. X +#How: creates DynamicCache and assigns each flat key/value tensor into the matching cache layer. +#Why: CacheSpec.to_dynamic_cache uses this so CacheExportWrapper.forward can call the original HF model normally. def cache_spec_to_dynamic_cache(cache_spec: CacheSpec, flat_tensors: tuple[torch.Tensor, ...]): from transformers.cache_utils import DynamicCache @@ -264,7 +277,8 @@ def cache_spec_to_dynamic_cache(cache_spec: CacheSpec, flat_tensors: tuple[torch return cache -#Selects the main tensor output from a model output object. X +#How: checks common HF output containers for logits, last_hidden_state, or first tuple item. +#Why: CacheExportWrapper.forward needs a stable first exported output before appending flat cache tensors. def primary_model_output(outputs: Any) -> torch.Tensor: if isinstance(outputs, dict): for key in ("logits", "last_hidden_state"): @@ -284,7 +298,8 @@ def primary_model_output(outputs: Any) -> torch.Tensor: raise ValueError("Could not find a tensor output to export") -#Flattens an HF DynamicCache into key/value tensor outputs. X +#How: walks either legacy tuple caches or DynamicCache layers and emits key/value tensors in flat order. +#Why: CacheExportWrapper.forward returns this flat tuple so the layermap contains explicit cache output tensors. def flatten_dynamic_cache(cache: Any) -> tuple[torch.Tensor, ...]: if cache is None: return () @@ -316,7 +331,8 @@ def flatten_dynamic_cache(cache: Any) -> tuple[torch.Tensor, ...]: return tuple(flat) -#Drops kwargs unsupported by a module's forward signature. X +#How: inspects forward(...) and drops kwargs unless the module accepts **kwargs. +#Why: used by CacheExportWrapper.forward and models.export_ to avoid passing unsupported modality/cache args. def filter_forward_kwargs(model: torch.nn.Module, kwargs: dict[str, Any]) -> dict[str, Any]: try: signature = inspect.signature(model.forward) @@ -329,7 +345,8 @@ def filter_forward_kwargs(model: torch.nn.Module, kwargs: dict[str, Any]) -> dic return {key: value for key, value in kwargs.items() if key in signature.parameters} -#Sets an attribute on a config object and common nested config objects. X +#How: recursively sets one config attr on top-level and nested text/vision/audio/encoder/decoder configs. +#Why: configure_model_for_export uses it to consistently enable or disable use_cache before tracing. def _set_config_attr_recursive(config: Any, attr: str, value: Any, seen: set[int] | None = None) -> None: if config is None: return @@ -355,7 +372,8 @@ def _set_config_attr_recursive(config: Any, attr: str, value: Any, seen: set[int _set_config_attr_recursive(getattr(config, child_name, None), attr, value, seen) -#Configures cache-related model settings before torch.export. X +#How: sets use_cache on model.config and model.generation_config, including nested configs. +#Why: models.export_ calls this so prefill/decode cache exports include cache paths and no-cache exports do not. def configure_model_for_export(model: torch.nn.Module, should_use_cache: bool) -> None: for config_name in ("config", "generation_config"): config = getattr(model, config_name, None) diff --git a/python/cactus/transpiler/Converter/constants.py b/python/cactus/transpiler/Converter/constants.py index 8b1d7ebcf..ce78e5223 100644 --- a/python/cactus/transpiler/Converter/constants.py +++ b/python/cactus/transpiler/Converter/constants.py @@ -1,5 +1,8 @@ from transformers import AutoModel, AutoModelForCausalLM, AutoModelForCTC, AutoModelForImageTextToText, AutoModelForSeq2SeqLM, AutoModelForSpeechSeq2Seq -from . import input_utils as IU +import os +from pathlib import Path + +token = os.environ.get("HF_TOKEN") #Names of component configs in config.txt (put in constants to allow future expansion to tuple to support different names) TEXT_CONFIG_NAME = "text_config" @@ -14,6 +17,52 @@ "causal_lm": (AutoModelForCausalLM, AutoModel), } -CACHE_INFERENCE_MODES = {IU.PREFILL_WITH_CACHE_MODE, IU.DECODE_WITH_CACHE_MODE} -DYNAMIC_CACHE_POLICY = IU.DYNAMIC_CACHE_POLICY +CACHE_INFERENCE_MODES = {"prefill_with_cache", "decode_with_cache"} +DYNAMIC_CACHE_POLICY = "dynamic_cache" DROP_MULTIMODAL_ON_DECODE_POLICY = "drop_multimodal_on_decode" + +CONFIG_SEARCH_SECTIONS = ( + (), + ("text_config",), + ("vision_config",), + ("audio_config",), + ("encoder",), + ("decoder",), + ("model_config",), +) + +DEFAULT_BATCH_SIZE = 1 +DEFAULT_SEQUENCE_LENGTH = 8 +DEFAULT_PAST_SEQUENCE_LENGTH = 8 +DEFAULT_IMAGE_HEIGHT = 224 +DEFAULT_IMAGE_WIDTH = 224 +DEFAULT_BIN_NUM = 80 +DEFAULT_SOURCE_POSITION = 1500 + +SYNTHETIC_INPUT_STRATEGY = "synthetic" +PROCESSOR_INPUT_STRATEGY = "processor" + +PREFILL_WITH_CACHE_MODE = "prefill_with_cache" +DECODE_WITH_CACHE_MODE = "decode_with_cache" +DYNAMIC_CACHE_POLICY = "dynamic_cache" + +CONVERTER_JSON_DIR = Path(__file__).resolve().parent / "jsons" +ASSET_DIR = Path(__file__).resolve().parents[2] / "assets" + + +MULTIMODAL_KEYS = ( + "pixel_values", + "pixel_values_videos", + "input_features", + "input_features_mask", + "image_position_ids", + "video_position_ids", + "mm_token_type_ids", +) + +#Locations to local sample inputs used by processor-backed multimodal exports. +MODALITY_INPUT_PATH = { + "vision": ASSET_DIR / "test_monkey.png", + "audio": ASSET_DIR / "test.wav", +} + diff --git a/python/cactus/transpiler/Converter/input_processor.py b/python/cactus/transpiler/Converter/input_processor.py new file mode 100644 index 000000000..a816d7e98 --- /dev/null +++ b/python/cactus/transpiler/Converter/input_processor.py @@ -0,0 +1,65 @@ +from typing import Any +from huggingface_hub import hf_hub_download +from . import constants +from pathlib import Path + + +#How: returns a config copy without the HF *_type key that component constructors do not expect. +#Why: _build_gemma4_processor uses this when manually constructing Gemma processor pieces from processor_config.json. +def _without_type_key(config: dict[str, Any], type_key: str) -> dict[str, Any]: + return {key: value for key, value in config.items() if key != type_key} + +def gemma4_processor(model_id: str, configs: dict[str, dict[str, Any]]): + + from transformers import PreTrainedTokenizerFast + from transformers.models.gemma4 import Gemma4AudioFeatureExtractor, Gemma4ImageProcessor, Gemma4Processor, Gemma4VideoProcessor + + processor_config = configs.get("processor_config.json") + if processor_config is None or processor_config.get("processor_class") != "Gemma4Processor": + return None + + #TODO: Path to where we downloaded the config files + tokenizer_path = constants.CONVERTER_JSON_DIR+"/"+str + + tokenizer_config = configs.get("tokenizer_config.json", {}) + tokenizer = PreTrainedTokenizerFast(tokenizer_file=tokenizer_path, bos_token=tokenizer_config.get("bos_token", ""), eos_token=tokenizer_config.get("eos_token", ""), unk_token=tokenizer_config.get("unk_token", ""), pad_token=tokenizer_config.get("pad_token", "")) + + for token_name in ("image_token", "audio_token", "boi_token", "eoi_token", "boa_token", "eoa_token"): + token_value = tokenizer_config.get(token_name) + if token_value is not None: + setattr(tokenizer, token_name, token_value) + + tokenizer.image_token_id = tokenizer.convert_tokens_to_ids(tokenizer.image_token) + tokenizer.audio_token_id = tokenizer.convert_tokens_to_ids(tokenizer.audio_token) + image_processor = Gemma4ImageProcessor(**_without_type_key(processor_config.get("image_processor", {}), "image_processor_type")) + feature_extractor = Gemma4AudioFeatureExtractor(**_without_type_key(processor_config.get("feature_extractor", {}), "feature_extractor_type")) + video_processor = Gemma4VideoProcessor(**_without_type_key(processor_config.get("video_processor", {}), "video_processor_type")) + + return Gemma4Processor( + feature_extractor=feature_extractor, + image_processor=image_processor, + tokenizer=tokenizer, + video_processor=video_processor, + image_seq_length=processor_config.get("image_seq_length", 280), + audio_seq_length=processor_config.get("audio_seq_length", 750), + audio_ms_per_token=processor_config.get("audio_ms_per_token", 40), + ) + + +#How: loops over hf_load_kwargs attempts and returns the first AutoProcessor that loads. +#Why: build_processor_kwargs needs an HF processor to convert real sample assets into model-ready tensors. +def default_processor(model_id: str, processor_kwargs: dict[str, Any]): + from transformers import AutoProcessor + try: + return AutoProcessor.from_pretrained(model_id, processor_kwargs) + except Exception as e: + raise RuntimeError(f"Unable to load processor for {model_id}") + + +PROCESSOR_MAP = { + "google/gemma-4-E2B": gemma4_processor, + "openai/whisper-tiny": default_processor, + "nvidia/parakeet-tdt-0.6b-v3": default_processor, + "LiquidAI/LFM2-VL-3B": default_processor, + "Qwen/Qwen2.5-0.5B": default_processor, +} \ No newline at end of file diff --git a/python/cactus/transpiler/Converter/input_utils.py b/python/cactus/transpiler/Converter/input_utils.py index b83219ea7..1c554ed14 100644 --- a/python/cactus/transpiler/Converter/input_utils.py +++ b/python/cactus/transpiler/Converter/input_utils.py @@ -1,61 +1,18 @@ import json -import os from pathlib import Path from typing import Any - from huggingface_hub import hf_hub_download import numpy as np import torch +from . import constants +from . import input_processor as IP from ..ModelProfiles import models as MP_Models -DEFAULT_BATCH_SIZE = 1 -DEFAULT_SEQUENCE_LENGTH = 8 -DEFAULT_PAST_SEQUENCE_LENGTH = 8 -SYNTHETIC_INPUT_STRATEGY = "synthetic" -PROCESSOR_INPUT_STRATEGY = "processor" -PREFILL_WITH_CACHE_MODE = "prefill_with_cache" -DECODE_WITH_CACHE_MODE = "decode_with_cache" -DYNAMIC_CACHE_POLICY = "dynamic_cache" -CONVERTER_JSON_DIR = Path(__file__).resolve().parent / "jsons" -ASSET_DIR = Path(__file__).resolve().parents[2] / "assets" - -CONFIG_SEARCH_SECTIONS = ( - (), - ("text_config",), - ("vision_config",), - ("audio_config",), - ("encoder",), - ("decoder",), - ("model_config",), -) - -MULTIMODAL_KEYS = ( - "pixel_values", - "pixel_values_videos", - "input_features", - "input_features_mask", - "image_position_ids", - "video_position_ids", - "mm_token_type_ids", -) - -#Locations to local sample inputs used by processor-backed multimodal exports. -MODALITY_INPUT_PATH = { - "vision": ASSET_DIR / "test_monkey.png", - "audio": ASSET_DIR / "test.wav", -} - - -#Downloads profile-declared JSON files from Hugging Face and loads them. def load_configs(mp: MP_Models.ModelProfile, model_id: str | None) -> dict[str, dict[str, Any]]: - if model_id is None: - return {} - - output_dir = CONVERTER_JSON_DIR / mp.model_profiles + output_dir = constants.CONVERTER_JSON_DIR / mp.model_profiles output_dir.mkdir(parents=True, exist_ok=True) - token = os.environ.get("HF_TOKEN") configs: dict[str, dict[str, Any]] = {} for filename in mp.files: @@ -66,7 +23,7 @@ def load_configs(mp: MP_Models.ModelProfile, model_id: str | None) -> dict[str, if not local_path.exists(): try: - local_path = Path(hf_hub_download(repo_id=model_id, filename=filename, local_dir=output_dir, token=token)) + local_path = Path(hf_hub_download(repo_id=model_id, filename=filename, local_dir=output_dir, token=constants.token)) except Exception as e: print(f"Error downloading {filename} from {model_id}: {e}") continue @@ -77,223 +34,37 @@ def load_configs(mp: MP_Models.ModelProfile, model_id: str | None) -> dict[str, return configs -#Searches through configs to find desired keys, and their values -def _find_config_value(configs: dict[str, dict[str, Any]], keys: tuple[str, ...]) -> Any | None: - for config in configs.values(): - for section_path in CONFIG_SEARCH_SECTIONS: - section = config - - for section_name in section_path: - if not isinstance(section, dict): - section = None - break - section = section.get(section_name) - - if not isinstance(section, dict): - continue - - for key in keys: - if key in section and section[key] is not None: - return section[key] - - return None - - -#Returns the first config value that can safely be interpreted as an int. -def _first_int(configs: dict[str, dict[str, Any]], keys: tuple[str, ...], default: int) -> int: - value = _find_config_value(configs, keys) - if isinstance(value, bool): - return default - if isinstance(value, int): - return value - if isinstance(value, float): - return int(value) - if isinstance(value, str) and value.isdigit(): - return int(value) - return default - - -#Builds simple token ids for generic text-only synthetic exports. -def _build_input_ids(batch_size: int, sequence_length: int) -> torch.Tensor: - tokens = torch.arange(sequence_length, dtype=torch.long) + 1 - return tokens.unsqueeze(0).repeat(batch_size, 1) - - -#Builds an all-ones attention mask for synthetic exports. -def _build_attention_mask(batch_size: int, sequence_length: int) -> torch.Tensor: - return torch.ones((batch_size, sequence_length), dtype=torch.long) - - -#Builds synthetic image pixel tensors from config-derived image dimensions. -def _build_pixel_values(batch_size: int, configs: dict[str, dict[str, Any]]) -> torch.Tensor: - size = _find_config_value(configs, ("size", "image_size")) - channels = _first_int(configs, ("num_channels", "image_num_channels"), 3) - height = 224 - width = 224 - - if isinstance(size, int): - height = size - width = size - elif isinstance(size, dict): - height = int(size.get("height") or size.get("shortest_edge") or height) - width = int(size.get("width") or size.get("shortest_edge") or width) - - return torch.zeros((batch_size, channels, height, width), dtype=torch.float32) - - -#Builds synthetic audio feature tensors from config-derived audio dimensions. -def _build_input_features(batch_size: int, configs: dict[str, dict[str, Any]]) -> torch.Tensor: - feature_bins = _first_int(configs, ("feature_size", "num_mel_bins", "n_mels"), 80) - max_source_positions = _first_int(configs, ("max_source_positions",), 1500) - feature_frames = _first_int(configs, ("nb_max_frames", "max_length", "feature_frames"), max_source_positions * 2) - return torch.zeros((batch_size, feature_bins, feature_frames), dtype=torch.float32) - - -#Builds generic tuple-style KV cache tensors for models that still accept legacy cache inputs. -def _build_past_key_values(configs: dict[str, dict[str, Any]], batch_size: int, past_sequence_length: int) -> tuple[tuple[torch.Tensor, torch.Tensor], ...]: - num_layers = _first_int(configs, ("num_hidden_layers", "n_layer", "num_layers", "decoder_layers"), 1) - num_attention_heads = _first_int(configs, ("num_attention_heads", "n_head", "decoder_attention_heads"), 1) - num_key_value_heads = _first_int(configs, ("num_key_value_heads", "num_kv_heads"), num_attention_heads) - hidden_size = _first_int(configs, ("hidden_size", "n_embd", "d_model"), 64) - head_dim = _first_int(configs, ("head_dim", "kv_channels"), max(1, hidden_size // max(1, num_attention_heads))) - cache_shape = (batch_size, num_key_value_heads, past_sequence_length, head_dim) - - return tuple( - ( - torch.zeros(cache_shape, dtype=torch.float32), - torch.zeros(cache_shape, dtype=torch.float32), - ) - for _ in range(num_layers) - ) - - -#Finds a decoder start token id for seq2seq-style synthetic inputs. -def _decoder_start_token_id(configs: dict[str, dict[str, Any]]) -> int: - return _first_int(configs, ("decoder_start_token_id", "bos_token_id", "pad_token_id"), 1) - - -#Builds config-shaped synthetic kwargs when a profile does not use a processor. -def build_synthetic_kwargs( - modalities: tuple[str, ...], - configs: dict[str, dict[str, Any]], - inference_mode: str, - uses_dynamic_cache: bool, -) -> dict[str, Any]: - batch_size = DEFAULT_BATCH_SIZE - sequence_length = DEFAULT_SEQUENCE_LENGTH - past_sequence_length = DEFAULT_PAST_SEQUENCE_LENGTH - kwargs: dict[str, Any] = {} - - if "text" in modalities: - if inference_mode == DECODE_WITH_CACHE_MODE and not uses_dynamic_cache: - sequence_length = 1 - - kwargs["input_ids"] = _build_input_ids(batch_size, sequence_length) - kwargs["attention_mask"] = _build_attention_mask(batch_size, sequence_length) - - if "vision" in modalities: - kwargs["pixel_values"] = _build_pixel_values(batch_size, configs) - - if "audio" in modalities: - kwargs["input_features"] = _build_input_features(batch_size, configs) - - if "audio" in modalities and "text" in modalities and "vision" not in modalities and "speech" in inference_mode: - kwargs.pop("input_ids", None) - kwargs.pop("attention_mask", None) - kwargs["decoder_input_ids"] = torch.full( - (batch_size, 1), - _decoder_start_token_id(configs), - dtype=torch.long, - ) - - if "audio" in modalities and "text" not in modalities: - input_features = kwargs["input_features"] - kwargs["attention_mask"] = _build_attention_mask(batch_size, input_features.shape[-1]) - - if inference_mode == DECODE_WITH_CACHE_MODE and not uses_dynamic_cache: - total_sequence_length = past_sequence_length + sequence_length - kwargs["attention_mask"] = _build_attention_mask(batch_size, total_sequence_length) - kwargs["past_key_values"] = _build_past_key_values(configs, batch_size, past_sequence_length) - kwargs["cache_position"] = torch.arange( - past_sequence_length, - past_sequence_length + sequence_length, - dtype=torch.long, - ) - kwargs["use_cache"] = True - elif inference_mode == PREFILL_WITH_CACHE_MODE: - kwargs["use_cache"] = True - - return kwargs - - -#Builds representative model inputs for the requested modalities and inference mode. X -def build_input( - mp: MP_Models.ModelProfile, - input_modalities: tuple[str, ...], - input_cls: Any, - model_id: str | None = None, - inference_mode: str = "prefill_no_cache", -) -> Any | None: +def build_input(mp: MP_Models.ModelProfile, input_modalities: tuple[str, ...], input_cls: Any, model_id: str | None = None, inference_mode: str = "prefill_no_cache") -> Any | None: if not all(modality in mp.supported_modalties for modality in input_modalities): print("Requesting unsupported modalities") return None configs = load_configs(mp, model_id) - modalities = tuple(input_modalities) - uses_dynamic_cache = DYNAMIC_CACHE_POLICY in mp.cache_policy - - if model_id is not None and mp.input_strategy != SYNTHETIC_INPUT_STRATEGY: - try: - return input_cls( - args=(), - kwargs=build_processor_kwargs( - model_id=model_id, - input_modalities=modalities, - configs=configs, - input_strategy=mp.input_strategy, - ), - modalities=modalities, - inference_mode=inference_mode, - ) - except Exception as e: - print(f"{mp.input_strategy} input build failed for {model_id}, falling back to synthetic inputs: {e}") + + if mp.input_strategy == constants.SYNTHETIC_INPUT_STRATEGY: + raise ValueError("Synthetic input creation is currently disabled") return input_cls( args=(), - kwargs=build_synthetic_kwargs( - modalities=modalities, - configs=configs, - inference_mode=inference_mode, - uses_dynamic_cache=uses_dynamic_cache, - ), - modalities=modalities, + kwargs=build_processor_kwargs(model_id, input_modalities), + modalities=input_modalities, inference_mode=inference_mode, ) -#Builds one-token decode inputs and flat KV tensors for cache-mode exports. X -def build_decode_with_cache_input( - model: torch.nn.Module, - input_: Any, - input_cls: Any, - cache_spec_cls: Any, - model_dtype_fn: Any, - drop_multimodal: bool, -) -> Any: +#How: slices full prompt tokens to one decode token, adds cache_position, attention_mask, and flat cache tensors. +#Why: models.create_model uses this for decode_with_cache so the exported graph represents one-token generation. +def build_decode_with_cache_input(model: torch.nn.Module, input_: Any, input_cls: Any, cache_spec_cls: Any, model_dtype_fn: Any, drop_multimodal: bool) -> Any: kwargs = dict(input_.kwargs) token_key = "input_ids" if "input_ids" in kwargs else "decoder_input_ids" if token_key not in kwargs: - raise ValueError(f"{DECODE_WITH_CACHE_MODE} requires input_ids or decoder_input_ids") + raise ValueError(f"{constants.DECODE_WITH_CACHE_MODE} requires input_ids or decoder_input_ids") token_ids = kwargs[token_key] batch_size = int(token_ids.shape[0]) past_sequence_length = int(token_ids.shape[1]) - cache_spec = cache_spec_cls.from_model( - model=model, - batch_size=batch_size, - past_sequence_length=past_sequence_length, - ) + cache_spec = cache_spec_cls.from_model(model=model, batch_size=batch_size, past_sequence_length=past_sequence_length) kwargs[token_key] = token_ids[:, -1:].clone() kwargs["attention_mask"] = torch.ones( @@ -313,7 +84,7 @@ def build_decode_with_cache_input( ) if drop_multimodal: - for multimodal_key in MULTIMODAL_KEYS: + for multimodal_key in constants.MULTIMODAL_KEYS: kwargs.pop(multimodal_key, None) return input_cls( @@ -324,7 +95,8 @@ def build_decode_with_cache_input( ) -#Infers batch size from the first tensor-shaped input. X +#How: scans kwargs for the first tensor with a batch dimension and returns shape[0]. +#Why: models.export_ uses this to size CacheSpec when wrapping cache-mode exports. def infer_batch_size(kwargs: dict[str, Any]) -> int: for value in kwargs.values(): if isinstance(value, torch.Tensor) and value.ndim > 0: @@ -333,9 +105,10 @@ def infer_batch_size(kwargs: dict[str, Any]) -> int: return 1 -#Infers the past/prompt sequence length represented by the export input. X +#How: prefers cache_position for decode, otherwise checks token or attention-mask sequence length. +#Why: models.export_ uses this to tell CacheSpec how much past context the exported cache represents. def infer_past_sequence_length(input_: Any) -> int: - if input_.inference_mode == DECODE_WITH_CACHE_MODE and "cache_position" in input_.kwargs: + if input_.inference_mode == constants.DECODE_WITH_CACHE_MODE and "cache_position" in input_.kwargs: return int(input_.kwargs["cache_position"][0].item()) if "input_ids" in input_.kwargs: @@ -350,18 +123,19 @@ def infer_past_sequence_length(input_: Any) -> int: return 0 -#Loads the local image asset used for processor-backed multimodal exports. +#How: opens the configured image path and converts it to RGB. +#Why: build_processor_kwargs needs a real image object when using HF processors for vision models. def _load_image_asset(): from PIL import Image + return Image.open(constants.MODALITY_INPUT_PATH["vision"]).convert("RGB") - return Image.open(MODALITY_INPUT_PATH["vision"]).convert("RGB") - -#Loads and normalizes the local audio asset used for processor-backed multimodal exports. +#How: reads the configured WAV file, mixes stereo to mono, and normalizes integer samples to float32. +#Why: build_processor_kwargs needs real waveform data when using HF processors for audio models. def _load_audio_asset() -> np.ndarray: from scipy.io import wavfile - _sample_rate, audio = wavfile.read(MODALITY_INPUT_PATH["audio"]) + _sample_rate, audio = wavfile.read(constants.MODALITY_INPUT_PATH["audio"]) if audio.ndim > 1: audio = audio.mean(axis=1) @@ -374,137 +148,12 @@ def _load_audio_asset() -> np.ndarray: return audio -#Removes Hugging Face processor type metadata before constructing processor components directly. -def _without_type_key(config: dict[str, Any], type_key: str) -> dict[str, Any]: - return {key: value for key, value in config.items() if key != type_key} - - -#Finds a previously cached Hugging Face file without requiring network access. -def _find_cached_hf_file(model_id: str, filename: str) -> str | None: - cache_model_id = model_id.replace("/", "--") - cache_root = Path.home() / ".cache" / "huggingface" / "hub" / f"models--{cache_model_id}" - - if not cache_root.exists(): - return None - - for candidate in cache_root.glob(f"snapshots/*/{filename}"): - if candidate.is_file(): - return str(candidate) - - return None -#Gemma-specific: manually builds Gemma4Processor when AutoProcessor cannot be loaded cleanly offline. -def _build_gemma4_processor(model_id: str, configs: dict[str, dict[str, Any]]): - from transformers import PreTrainedTokenizerFast - from transformers.models.gemma4 import ( - Gemma4AudioFeatureExtractor, - Gemma4ImageProcessor, - Gemma4Processor, - Gemma4VideoProcessor, - ) - - processor_config = configs.get("processor_config.json") - if processor_config is None or processor_config.get("processor_class") != "Gemma4Processor": - return None - - token = os.environ.get("HF_TOKEN") - tokenizer_path = _find_cached_hf_file(model_id, "tokenizer.json") - if tokenizer_path is None: - try: - tokenizer_path = hf_hub_download(repo_id=model_id, filename="tokenizer.json", token=token, local_files_only=True) - except Exception: - tokenizer_path = hf_hub_download(repo_id=model_id, filename="tokenizer.json", token=token) - - tokenizer_config = configs.get("tokenizer_config.json", {}) - tokenizer = PreTrainedTokenizerFast( - tokenizer_file=tokenizer_path, - bos_token=tokenizer_config.get("bos_token", ""), - eos_token=tokenizer_config.get("eos_token", ""), - unk_token=tokenizer_config.get("unk_token", ""), - pad_token=tokenizer_config.get("pad_token", ""), - ) - - for token_name in ( - "image_token", - "audio_token", - "boi_token", - "eoi_token", - "boa_token", - "eoa_token", - ): - token_value = tokenizer_config.get(token_name) - if token_value is not None: - setattr(tokenizer, token_name, token_value) - - tokenizer.image_token_id = tokenizer.convert_tokens_to_ids(tokenizer.image_token) - tokenizer.audio_token_id = tokenizer.convert_tokens_to_ids(tokenizer.audio_token) - image_processor = Gemma4ImageProcessor( - **_without_type_key(processor_config.get("image_processor", {}), "image_processor_type") - ) - feature_extractor = Gemma4AudioFeatureExtractor( - **_without_type_key(processor_config.get("feature_extractor", {}), "feature_extractor_type") - ) - video_processor = Gemma4VideoProcessor( - **_without_type_key(processor_config.get("video_processor", {}), "video_processor_type") - ) - - return Gemma4Processor( - feature_extractor=feature_extractor, - image_processor=image_processor, - tokenizer=tokenizer, - video_processor=video_processor, - image_seq_length=processor_config.get("image_seq_length", 280), - audio_seq_length=processor_config.get("audio_seq_length", 750), - audio_ms_per_token=processor_config.get("audio_ms_per_token", 40), - ) - - -#Loads AutoProcessor from local HF cache first, then falls back to normal HF loading. -def _load_auto_processor(model_id: str, processor_kwargs: dict[str, Any]): - from transformers import AutoProcessor - - last_error: Exception | None = None - - for attempt_kwargs in hf_load_kwargs(processor_kwargs): - try: - return AutoProcessor.from_pretrained(model_id, **attempt_kwargs) - except Exception as e: - last_error = e - - raise RuntimeError(f"Unable to load processor for {model_id}") from last_error - - -#Returns local-cache-first HF loading kwargs plus a normal network-capable fallback. -def hf_load_kwargs(load_kwargs: dict[str, Any]) -> tuple[dict[str, Any], ...]: - if load_kwargs.get("local_files_only") is True: - return (load_kwargs,) - - return ( - {**load_kwargs, "local_files_only": True}, - load_kwargs, - ) - - -#Builds real processor-backed tensor kwargs for multimodal models. -def build_processor_kwargs( - model_id: str, - input_modalities: tuple[str, ...], - configs: dict[str, dict[str, Any]] | None = None, - input_strategy: str = PROCESSOR_INPUT_STRATEGY, -) -> dict[str, Any]: - configs = configs or {} - token = os.environ.get("HF_TOKEN") - processor_kwargs: dict[str, Any] = {"trust_remote_code": True} - if token: - processor_kwargs["token"] = token - - if input_strategy == PROCESSOR_INPUT_STRATEGY: - processor = _load_auto_processor(model_id, processor_kwargs) - elif input_strategy == "manual_gemma4_processor": - processor = _build_gemma4_processor(model_id, configs) or _load_auto_processor(model_id, processor_kwargs) - else: - raise ValueError(f"Unknown input strategy: {input_strategy}") +#How: loads the requested processor, builds a prompt with modality tokens/assets, then keeps only tensor outputs. +#Why: build_input uses this for profile-driven processor exports so the graph sees realistic multimodal inputs. +def build_processor_kwargs(model_id: str, input_modalities: tuple[str, ...]) -> dict[str, Any]: + processor = IP.PROCESSOR_MAP[model_id] prompt_parts = [] call_kwargs: dict[str, Any] = {"return_tensors": "pt"} diff --git a/python/cactus/transpiler/Converter/models.py b/python/cactus/transpiler/Converter/models.py index 048c9b1d8..cbddc0e5e 100644 --- a/python/cactus/transpiler/Converter/models.py +++ b/python/cactus/transpiler/Converter/models.py @@ -26,7 +26,7 @@ class Model: name: str model_profile: MP_Models.ModelProfile input: Input - model: Any + model: torch.nn.Module def export(self, input: Input) -> "LayerMap": return export_(model=self, input=input) @@ -159,24 +159,23 @@ def load_model(model_id: str, mp: MP_Models.ModelProfile | None = None) -> torch } -#Builds a loaded model bundle from profile, modalities, model id, and inference mode. X +#Builds a loaded model bundle from profile, modalities, model id, and inference mode. def create_model(mp: MP_Models.ModelProfile, input_modalities: tuple[str, ...], model_id: str, inference_mode: str = "prefill_no_cache") -> Model: + input_ = IU.build_input(mp, input_modalities, Input, model_id, inference_mode) - if input_ is None: raise ValueError(f"Could not build input for modalities {input_modalities}") loaded_model = load_model(model_id, mp) - #Build a decode with cache input only if decoder/ctc head supports cache + if input_.inference_mode == IU.DECODE_WITH_CACHE_MODE and constants.DYNAMIC_CACHE_POLICY in mp.cache_policy: - input_ = IU.build_decode_with_cache_input( - model=loaded_model, - input_=input_, - input_cls=Input, - cache_spec_cls=CU.CacheSpec, - model_dtype_fn=CU.model_dtype, - drop_multimodal=constants.DROP_MULTIMODAL_ON_DECODE_POLICY in mp.cache_policy, - ) + input_ = IU.build_decode_with_cache_input(model=loaded_model, input_=input_, input_cls=Input, cache_spec_cls=CU.CacheSpec, model_dtype_fn=CU.model_dtype, drop_multimodal=constants.DROP_MULTIMODAL_ON_DECODE_POLICY in mp.cache_policy) + + #TODO: Add for models that do support multimodal decoders + # elif input_.inference_mode == IU.DECODE_WITH_CACHE_MODE: + # input_ = IU.build_decode_with_cache_input() + + return Model(name=model_id, model_profile=mp, input=input_, model=loaded_model) @@ -190,11 +189,7 @@ def export_(model: Model, input: Input) -> LayerMap: if should_use_cache and constants.DYNAMIC_CACHE_POLICY in model.model_profile.cache_policy: export_model = CU.CacheExportWrapper( model=model.model, - cache_spec=CU.CacheSpec.from_model( - model=model.model, - batch_size=IU.infer_batch_size(input.kwargs), - past_sequence_length=IU.infer_past_sequence_length(input), - ), + cache_spec=CU.CacheSpec.from_model(model=model.model, batch_size=IU.infer_batch_size(input.kwargs), past_sequence_length=IU.infer_past_sequence_length(input)), ) export_kwargs = dict(input.kwargs) diff --git a/python/cactus/transpiler/Converter/export_patches.py b/python/cactus/transpiler/Converter/overrides.py similarity index 100% rename from python/cactus/transpiler/Converter/export_patches.py rename to python/cactus/transpiler/Converter/overrides.py From 1c4a1f092d8a7db2fa7e8cd6369ca994cbf990dc Mon Sep 17 00:00:00 2001 From: Parkirat Date: Thu, 23 Jul 2026 16:27:13 -0400 Subject: [PATCH 40/84] Converter input_utils removal + input loading code simplification --- .../transpiler/Converter/cache_utils.py | 4 +- .../transpiler/Converter/input_processor.py | 33 ++- .../transpiler/Converter/input_utils.py | 178 ---------------- python/cactus/transpiler/Converter/models.py | 194 ++++++++++++++++-- 4 files changed, 204 insertions(+), 205 deletions(-) delete mode 100644 python/cactus/transpiler/Converter/input_utils.py diff --git a/python/cactus/transpiler/Converter/cache_utils.py b/python/cactus/transpiler/Converter/cache_utils.py index 22e35666e..af6134d55 100644 --- a/python/cactus/transpiler/Converter/cache_utils.py +++ b/python/cactus/transpiler/Converter/cache_utils.py @@ -20,7 +20,7 @@ class CacheSpec: past_sequence_length: int #How: delegates model/cache inspection to cache_spec_from_model. - #Why: used by models.export_ and input_utils.build_decode_with_cache_input to create cache shapes for export. + #Why: used by models.export_ and models.build_decode_with_cache_input to create cache shapes for export. @classmethod def from_model(cls, model: torch.nn.Module, batch_size: int, past_sequence_length: int) -> "CacheSpec": return cache_spec_from_model(cls, model, batch_size, past_sequence_length) @@ -99,7 +99,7 @@ def get_text_config(config: Any) -> Any: #How: reads the dtype from the first model parameter and falls back to float32 for parameterless modules. -#Why: input_utils.build_decode_with_cache_input uses this so synthetic cache tensors match the model dtype. +#Why: models.build_decode_with_cache_input uses this so synthetic cache tensors match the model dtype. def model_dtype(model: torch.nn.Module) -> torch.dtype: try: param = next(model.parameters()) diff --git a/python/cactus/transpiler/Converter/input_processor.py b/python/cactus/transpiler/Converter/input_processor.py index a816d7e98..35330262b 100644 --- a/python/cactus/transpiler/Converter/input_processor.py +++ b/python/cactus/transpiler/Converter/input_processor.py @@ -4,22 +4,38 @@ from pathlib import Path +def _model_json_dir(model_profile: str) -> Path: + output_dir = constants.CONVERTER_JSON_DIR / model_profile + output_dir.mkdir(parents=True, exist_ok=True) + return output_dir + + +def _local_tokenizer_path(model_id: str, model_profile: str) -> str: + output_dir = _model_json_dir(model_profile) + tokenizer_path = output_dir / "tokenizer.json" + + if tokenizer_path.exists(): + return str(tokenizer_path) + + return str(Path(hf_hub_download(repo_id=model_id, filename="tokenizer.json", local_dir=output_dir, token=constants.token))) + + #How: returns a config copy without the HF *_type key that component constructors do not expect. #Why: _build_gemma4_processor uses this when manually constructing Gemma processor pieces from processor_config.json. def _without_type_key(config: dict[str, Any], type_key: str) -> dict[str, Any]: return {key: value for key, value in config.items() if key != type_key} -def gemma4_processor(model_id: str, configs: dict[str, dict[str, Any]]): + +def gemma4_processor(model_id: str, configs: dict[str, dict[str, Any]], model_profile: str): from transformers import PreTrainedTokenizerFast from transformers.models.gemma4 import Gemma4AudioFeatureExtractor, Gemma4ImageProcessor, Gemma4Processor, Gemma4VideoProcessor processor_config = configs.get("processor_config.json") if processor_config is None or processor_config.get("processor_class") != "Gemma4Processor": - return None + return default_processor(model_id=model_id, configs=configs, model_profile=model_profile) - #TODO: Path to where we downloaded the config files - tokenizer_path = constants.CONVERTER_JSON_DIR+"/"+str + tokenizer_path = _local_tokenizer_path(model_id=model_id, model_profile=model_profile) tokenizer_config = configs.get("tokenizer_config.json", {}) tokenizer = PreTrainedTokenizerFast(tokenizer_file=tokenizer_path, bos_token=tokenizer_config.get("bos_token", ""), eos_token=tokenizer_config.get("eos_token", ""), unk_token=tokenizer_config.get("unk_token", ""), pad_token=tokenizer_config.get("pad_token", "")) @@ -48,10 +64,13 @@ def gemma4_processor(model_id: str, configs: dict[str, dict[str, Any]]): #How: loops over hf_load_kwargs attempts and returns the first AutoProcessor that loads. #Why: build_processor_kwargs needs an HF processor to convert real sample assets into model-ready tensors. -def default_processor(model_id: str, processor_kwargs: dict[str, Any]): +def default_processor(model_id: str, configs: dict[str, dict[str, Any]], model_profile: str): from transformers import AutoProcessor + + processor_kwargs: dict[str, Any] = {"trust_remote_code": True, "token": constants.token} + try: - return AutoProcessor.from_pretrained(model_id, processor_kwargs) + return AutoProcessor.from_pretrained(model_id, **processor_kwargs) except Exception as e: raise RuntimeError(f"Unable to load processor for {model_id}") @@ -62,4 +81,4 @@ def default_processor(model_id: str, processor_kwargs: dict[str, Any]): "nvidia/parakeet-tdt-0.6b-v3": default_processor, "LiquidAI/LFM2-VL-3B": default_processor, "Qwen/Qwen2.5-0.5B": default_processor, -} \ No newline at end of file +} diff --git a/python/cactus/transpiler/Converter/input_utils.py b/python/cactus/transpiler/Converter/input_utils.py deleted file mode 100644 index 1c554ed14..000000000 --- a/python/cactus/transpiler/Converter/input_utils.py +++ /dev/null @@ -1,178 +0,0 @@ -import json -from pathlib import Path -from typing import Any -from huggingface_hub import hf_hub_download -import numpy as np -import torch - -from . import constants -from . import input_processor as IP -from ..ModelProfiles import models as MP_Models - - -def load_configs(mp: MP_Models.ModelProfile, model_id: str | None) -> dict[str, dict[str, Any]]: - output_dir = constants.CONVERTER_JSON_DIR / mp.model_profiles - output_dir.mkdir(parents=True, exist_ok=True) - configs: dict[str, dict[str, Any]] = {} - - for filename in mp.files: - if not filename.endswith(".json"): - continue - - local_path = output_dir / filename - - if not local_path.exists(): - try: - local_path = Path(hf_hub_download(repo_id=model_id, filename=filename, local_dir=output_dir, token=constants.token)) - except Exception as e: - print(f"Error downloading {filename} from {model_id}: {e}") - continue - - with local_path.open("r", encoding="utf-8") as f: - configs[filename] = json.load(f) - - return configs - - -def build_input(mp: MP_Models.ModelProfile, input_modalities: tuple[str, ...], input_cls: Any, model_id: str | None = None, inference_mode: str = "prefill_no_cache") -> Any | None: - if not all(modality in mp.supported_modalties for modality in input_modalities): - print("Requesting unsupported modalities") - return None - - configs = load_configs(mp, model_id) - - if mp.input_strategy == constants.SYNTHETIC_INPUT_STRATEGY: - raise ValueError("Synthetic input creation is currently disabled") - - return input_cls( - args=(), - kwargs=build_processor_kwargs(model_id, input_modalities), - modalities=input_modalities, - inference_mode=inference_mode, - ) - - -#How: slices full prompt tokens to one decode token, adds cache_position, attention_mask, and flat cache tensors. -#Why: models.create_model uses this for decode_with_cache so the exported graph represents one-token generation. -def build_decode_with_cache_input(model: torch.nn.Module, input_: Any, input_cls: Any, cache_spec_cls: Any, model_dtype_fn: Any, drop_multimodal: bool) -> Any: - kwargs = dict(input_.kwargs) - token_key = "input_ids" if "input_ids" in kwargs else "decoder_input_ids" - - if token_key not in kwargs: - raise ValueError(f"{constants.DECODE_WITH_CACHE_MODE} requires input_ids or decoder_input_ids") - - token_ids = kwargs[token_key] - batch_size = int(token_ids.shape[0]) - past_sequence_length = int(token_ids.shape[1]) - cache_spec = cache_spec_cls.from_model(model=model, batch_size=batch_size, past_sequence_length=past_sequence_length) - - kwargs[token_key] = token_ids[:, -1:].clone() - kwargs["attention_mask"] = torch.ones( - (batch_size, past_sequence_length + 1), - dtype=torch.long, - device=token_ids.device, - ) - kwargs["cache_position"] = torch.arange( - past_sequence_length, - past_sequence_length + 1, - dtype=torch.long, - device=token_ids.device, - ) - kwargs["past_key_values"] = cache_spec.empty_tensors( - dtype=model_dtype_fn(model), - device=token_ids.device, - ) - - if drop_multimodal: - for multimodal_key in constants.MULTIMODAL_KEYS: - kwargs.pop(multimodal_key, None) - - return input_cls( - args=input_.args, - kwargs=kwargs, - modalities=input_.modalities, - inference_mode=input_.inference_mode, - ) - - -#How: scans kwargs for the first tensor with a batch dimension and returns shape[0]. -#Why: models.export_ uses this to size CacheSpec when wrapping cache-mode exports. -def infer_batch_size(kwargs: dict[str, Any]) -> int: - for value in kwargs.values(): - if isinstance(value, torch.Tensor) and value.ndim > 0: - return int(value.shape[0]) - - return 1 - - -#How: prefers cache_position for decode, otherwise checks token or attention-mask sequence length. -#Why: models.export_ uses this to tell CacheSpec how much past context the exported cache represents. -def infer_past_sequence_length(input_: Any) -> int: - if input_.inference_mode == constants.DECODE_WITH_CACHE_MODE and "cache_position" in input_.kwargs: - return int(input_.kwargs["cache_position"][0].item()) - - if "input_ids" in input_.kwargs: - return int(input_.kwargs["input_ids"].shape[1]) - - if "decoder_input_ids" in input_.kwargs: - return int(input_.kwargs["decoder_input_ids"].shape[1]) - - if "attention_mask" in input_.kwargs: - return int(input_.kwargs["attention_mask"].shape[1]) - - return 0 - - -#How: opens the configured image path and converts it to RGB. -#Why: build_processor_kwargs needs a real image object when using HF processors for vision models. -def _load_image_asset(): - from PIL import Image - return Image.open(constants.MODALITY_INPUT_PATH["vision"]).convert("RGB") - - -#How: reads the configured WAV file, mixes stereo to mono, and normalizes integer samples to float32. -#Why: build_processor_kwargs needs real waveform data when using HF processors for audio models. -def _load_audio_asset() -> np.ndarray: - from scipy.io import wavfile - - _sample_rate, audio = wavfile.read(constants.MODALITY_INPUT_PATH["audio"]) - - if audio.ndim > 1: - audio = audio.mean(axis=1) - - if np.issubdtype(audio.dtype, np.integer): - audio = audio.astype(np.float32) / np.iinfo(audio.dtype).max - else: - audio = audio.astype(np.float32) - - return audio - - - - -#How: loads the requested processor, builds a prompt with modality tokens/assets, then keeps only tensor outputs. -#Why: build_input uses this for profile-driven processor exports so the graph sees realistic multimodal inputs. -def build_processor_kwargs(model_id: str, input_modalities: tuple[str, ...]) -> dict[str, Any]: - processor = IP.PROCESSOR_MAP[model_id] - - prompt_parts = [] - call_kwargs: dict[str, Any] = {"return_tensors": "pt"} - - if "vision" in input_modalities: - image_token = getattr(processor, "image_token", None) - if image_token is not None: - prompt_parts.append(image_token) - call_kwargs["images"] = _load_image_asset() - - if "audio" in input_modalities: - audio_token = getattr(processor, "audio_token", None) - if audio_token is not None: - prompt_parts.append(audio_token) - call_kwargs["audio"] = _load_audio_asset() - - if "text" in input_modalities: - prompt_parts.append("Describe this input.") - - call_kwargs["text"] = " ".join(prompt_parts) if prompt_parts else "Hello" - processed = processor(**call_kwargs) - return {key: value for key, value in dict(processed).items() if isinstance(value, torch.Tensor)} diff --git a/python/cactus/transpiler/Converter/models.py b/python/cactus/transpiler/Converter/models.py index cbddc0e5e..a28d776cd 100644 --- a/python/cactus/transpiler/Converter/models.py +++ b/python/cactus/transpiler/Converter/models.py @@ -1,17 +1,24 @@ +import json import os from dataclasses import dataclass +from pathlib import Path from typing import Any +from huggingface_hub import hf_hub_download +import numpy as np from pydantic import BaseModel import torch from transformers import AutoModel -from . import input_utils as IU from . import cache_utils as CU -from . import export_patches as EP +from . import input_processor as IP +from . import overrides as OV from ..ModelProfiles import models as MP_Models from . import constants -token = os.environ.get("HF_TOKEN") +token = constants.token +EXPORT_PATCHES = { + "gemma4_audio_mask": OV.patch_gemma4_audio_mask_for_export, +} @dataclass(slots=True) @@ -99,12 +106,10 @@ def jsonable(x: Any) -> Any: return repr(x) -#Extracts tensor metadata from an exported FX node. def extract_tensor_meta(node: torch.fx.Node) -> Any | None: return jsonable(node.meta["val"]) if "val" not in node.meta else None -#Extracts module stack metadata from an exported FX node. def extract_module_stack(node: torch.fx.Node) -> Any | None: stack = node.meta.get("nn_module_stack", None) out = [] @@ -122,6 +127,166 @@ def extract_module_stack(node: torch.fx.Node) -> Any | None: return out +def load_configs(mp: MP_Models.ModelProfile, model_id: str | None) -> dict[str, dict[str, Any]]: + output_dir = constants.CONVERTER_JSON_DIR / mp.model_profiles + output_dir.mkdir(parents=True, exist_ok=True) + configs: dict[str, dict[str, Any]] = {} + + for filename in mp.files: + if not filename.endswith(".json"): + continue + + local_path = output_dir / filename + + if not local_path.exists(): + try: + local_path = Path(hf_hub_download(repo_id=model_id, filename=filename, local_dir=output_dir, token=constants.token)) + except Exception as e: + print(f"Error downloading {filename} from {model_id}: {e}") + continue + + with local_path.open("r", encoding="utf-8") as f: + configs[filename] = json.load(f) + + return configs + + +def build_input(mp: MP_Models.ModelProfile, input_modalities: tuple[str, ...], input_cls: Any, model_id: str | None = None, inference_mode: str = "prefill_no_cache") -> Any | None: + if not all(modality in mp.supported_modalties for modality in input_modalities): + print("Requesting unsupported modalities") + return None + + configs = load_configs(mp, model_id) + + if mp.input_strategy == constants.SYNTHETIC_INPUT_STRATEGY: + raise ValueError("Synthetic input creation is currently disabled") + + return input_cls( + args=(), + kwargs=build_processor_kwargs(model_id, input_modalities, configs, mp.model_profiles), + modalities=input_modalities, + inference_mode=inference_mode, + ) + + +#How: slices full prompt tokens to one decode token, adds cache_position, attention_mask, and flat cache tensors. +#Why: create_model uses this for decode_with_cache so the exported graph represents one-token generation. +def build_decode_with_cache_input(model: torch.nn.Module, input_: Any, input_cls: Any, cache_spec_cls: Any, model_dtype_fn: Any, drop_multimodal: bool) -> Any: + kwargs = dict(input_.kwargs) + token_key = "input_ids" if "input_ids" in kwargs else "decoder_input_ids" + + if token_key not in kwargs: + raise ValueError(f"{constants.DECODE_WITH_CACHE_MODE} requires input_ids or decoder_input_ids") + + token_ids = kwargs[token_key] + batch_size = int(token_ids.shape[0]) + past_sequence_length = int(token_ids.shape[1]) + cache_spec = cache_spec_cls.from_model(model=model, batch_size=batch_size, past_sequence_length=past_sequence_length) + + kwargs[token_key] = token_ids[:, -1:].clone() + kwargs["attention_mask"] = torch.ones( + (batch_size, past_sequence_length + 1), + dtype=torch.long, + device=token_ids.device, + ) + kwargs["cache_position"] = torch.arange( + past_sequence_length, + past_sequence_length + 1, + dtype=torch.long, + device=token_ids.device, + ) + kwargs["past_key_values"] = cache_spec.empty_tensors( + dtype=model_dtype_fn(model), + device=token_ids.device, + ) + + if drop_multimodal: + for multimodal_key in constants.MULTIMODAL_KEYS: + kwargs.pop(multimodal_key, None) + + return input_cls( + args=input_.args, + kwargs=kwargs, + modalities=input_.modalities, + inference_mode=input_.inference_mode, + ) + + + +def infer_batch_size(kwargs: dict[str, Any]) -> int: + for value in kwargs.values(): + if isinstance(value, torch.Tensor) and value.ndim > 0: + return int(value.shape[0]) + + return 1 + + +#How: prefers cache_position for decode, otherwise checks token or attention-mask sequence length. +#Why: export_ uses this to tell CacheSpec how much past context the exported cache represents. +def infer_past_sequence_length(input_: Any) -> int: + if input_.inference_mode == constants.DECODE_WITH_CACHE_MODE and "cache_position" in input_.kwargs: + return int(input_.kwargs["cache_position"][0].item()) + + if "input_ids" in input_.kwargs: + return int(input_.kwargs["input_ids"].shape[1]) + + if "decoder_input_ids" in input_.kwargs: + return int(input_.kwargs["decoder_input_ids"].shape[1]) + + if "attention_mask" in input_.kwargs: + return int(input_.kwargs["attention_mask"].shape[1]) + + return 0 + + +def _load_image_asset(): + from PIL import Image + return Image.open(constants.MODALITY_INPUT_PATH["vision"]).convert("RGB") + + +def _load_audio_asset() -> np.ndarray: + from scipy.io import wavfile + + _sample_rate, audio = wavfile.read(constants.MODALITY_INPUT_PATH["audio"]) + + if audio.ndim > 1: + audio = audio.mean(axis=1) + + if np.issubdtype(audio.dtype, np.integer): + audio = audio.astype(np.float32) / np.iinfo(audio.dtype).max + else: + audio = audio.astype(np.float32) + + return audio + + +def build_processor_kwargs(model_id: str, input_modalities: tuple[str, ...], configs: dict[str, dict[str, Any]], model_profile: str) -> dict[str, Any]: + processor_builder = IP.PROCESSOR_MAP.get(model_id, IP.default_processor) + processor = processor_builder(model_id=model_id, configs=configs, model_profile=model_profile) + + prompt = [] + call_kwargs: dict[str, Any] = {"return_tensors": "pt"} + + if "vision" in input_modalities: + image_token = getattr(processor, "image_token", None) + if image_token is not None: + prompt.append(image_token) + call_kwargs["images"] = _load_image_asset() + + if "audio" in input_modalities: + audio_token = getattr(processor, "audio_token", None) + if audio_token is not None: + prompt.append(audio_token) + call_kwargs["audio"] = _load_audio_asset() + + if "text" in input_modalities: + prompt.append("Describe this input.") + + call_kwargs["text"] = " ".join(prompt) if prompt else "Hello" + processed = processor(**call_kwargs) + return {key: value for key, value in dict(processed).items() if isinstance(value, torch.Tensor)} + + #Loads the best HF model class for the requested model profile. def load_model(model_id: str, mp: MP_Models.ModelProfile | None = None) -> torch.nn.Module: if mp is None: @@ -153,34 +318,27 @@ def load_model(model_id: str, mp: MP_Models.ModelProfile | None = None) -> torch raise RuntimeError(f"Unable to load model {model_id}") - -EXPORT_PATCHES = { - "gemma4_audio_mask": EP.patch_gemma4_audio_mask_for_export, -} - - #Builds a loaded model bundle from profile, modalities, model id, and inference mode. def create_model(mp: MP_Models.ModelProfile, input_modalities: tuple[str, ...], model_id: str, inference_mode: str = "prefill_no_cache") -> Model: - input_ = IU.build_input(mp, input_modalities, Input, model_id, inference_mode) + input_ = build_input(mp, input_modalities, Input, model_id, inference_mode) if input_ is None: raise ValueError(f"Could not build input for modalities {input_modalities}") loaded_model = load_model(model_id, mp) - if input_.inference_mode == IU.DECODE_WITH_CACHE_MODE and constants.DYNAMIC_CACHE_POLICY in mp.cache_policy: - input_ = IU.build_decode_with_cache_input(model=loaded_model, input_=input_, input_cls=Input, cache_spec_cls=CU.CacheSpec, model_dtype_fn=CU.model_dtype, drop_multimodal=constants.DROP_MULTIMODAL_ON_DECODE_POLICY in mp.cache_policy) + if input_.inference_mode == constants.DECODE_WITH_CACHE_MODE and constants.DYNAMIC_CACHE_POLICY in mp.cache_policy: + input_ = build_decode_with_cache_input(model=loaded_model, input_=input_, input_cls=Input, cache_spec_cls=CU.CacheSpec, model_dtype_fn=CU.model_dtype, drop_multimodal=constants.DROP_MULTIMODAL_ON_DECODE_POLICY in mp.cache_policy) #TODO: Add for models that do support multimodal decoders - # elif input_.inference_mode == IU.DECODE_WITH_CACHE_MODE: - # input_ = IU.build_decode_with_cache_input() + # elif input_.inference_mode == constants.DECODE_WITH_CACHE_MODE: + # input_ = build_decode_with_cache_input() return Model(name=model_id, model_profile=mp, input=input_, model=loaded_model) -#Exports the prepared model and serializes its FX graph into a LayerMap. X def export_(model: Model, input: Input) -> LayerMap: should_use_cache = input.inference_mode in constants.CACHE_INFERENCE_MODES CU.configure_model_for_export(model.model, should_use_cache=should_use_cache) @@ -189,7 +347,7 @@ def export_(model: Model, input: Input) -> LayerMap: if should_use_cache and constants.DYNAMIC_CACHE_POLICY in model.model_profile.cache_policy: export_model = CU.CacheExportWrapper( model=model.model, - cache_spec=CU.CacheSpec.from_model(model=model.model, batch_size=IU.infer_batch_size(input.kwargs), past_sequence_length=IU.infer_past_sequence_length(input)), + cache_spec=CU.CacheSpec.from_model(model=model.model, batch_size=infer_batch_size(input.kwargs), past_sequence_length=infer_past_sequence_length(input)), ) export_kwargs = dict(input.kwargs) From 72e6874c9dbcda51ea9fecf84ba939de963e9a16 Mon Sep 17 00:00:00 2001 From: Parkirat Date: Thu, 23 Jul 2026 16:36:28 -0400 Subject: [PATCH 41/84] Removed unecessary code --- python/cactus/transpiler/ModelProfiles/profiles.py | 10 ---------- 1 file changed, 10 deletions(-) diff --git a/python/cactus/transpiler/ModelProfiles/profiles.py b/python/cactus/transpiler/ModelProfiles/profiles.py index 73ee9e677..3809af2fb 100644 --- a/python/cactus/transpiler/ModelProfiles/profiles.py +++ b/python/cactus/transpiler/ModelProfiles/profiles.py @@ -133,13 +133,3 @@ "LiquidAI/LFM2-VL-3B": LFM_VLM_PROFILE, "Qwen/Qwen2.5-0.5B": QWEN2_5_0_5B_PROFILE, } - -MODEL_ID_HINTS = ( - (("gemma",), GEMMA4_E2B_PROFILE), - (("whisper",), WHISPER_PROFILE), - (("parakeet",), PARAKEET_PROFILE), - (("lfm", "vl"), LFM_VLM_PROFILE), - (("lfm", "vision"), LFM_VLM_PROFILE), - (("qwen2.5",), QWEN2_5_0_5B_PROFILE), - (("qwen2_5",), QWEN2_5_0_5B_PROFILE), -) From de8514e6e31a1797be70aeaed24f8b896628203d Mon Sep 17 00:00:00 2001 From: Parkirat Date: Fri, 24 Jul 2026 11:13:45 -0400 Subject: [PATCH 42/84] Added fusion patterns --- python/cactus/transpiler/Fusions/edges.py | 202 +++++ python/cactus/transpiler/Fusions/fusions.py | 817 ++++++++++++++++++++ python/cactus/transpiler/Fusions/models.py | 360 ++++++++- python/cactus/transpiler/Fusions/nodes.py | 359 +++++++++ 4 files changed, 1730 insertions(+), 8 deletions(-) create mode 100644 python/cactus/transpiler/Fusions/edges.py create mode 100644 python/cactus/transpiler/Fusions/fusions.py create mode 100644 python/cactus/transpiler/Fusions/nodes.py diff --git a/python/cactus/transpiler/Fusions/edges.py b/python/cactus/transpiler/Fusions/edges.py new file mode 100644 index 000000000..1a313fc49 --- /dev/null +++ b/python/cactus/transpiler/Fusions/edges.py @@ -0,0 +1,202 @@ +from . import models as M + +EXPERT_GROUP = "experts" + + +def _edge( + source: str, + dest: str, + dest_input_index: int | None = 0, + *, + source_output_index: int | None = None, + optional: bool = False, + transparent: bool = False, + repeated: bool = False, + repeated_group: str | None = None, + metadata: dict | None = None, +) -> M.FusionEdge: + return M.FusionEdge( + source=source, + dest=dest, + dest_input_index=dest_input_index, + source_output_index=source_output_index, + optional=optional, + transparent=transparent, + repeated=repeated, + repeated_group=repeated_group, + metadata=metadata or {}, + ) + + +def _expert_edge( + source: str, + dest: str, + dest_input_index: int | None = 0, +) -> M.FusionEdge: + return _edge( + source, + dest, + dest_input_index, + repeated=True, + repeated_group=EXPERT_GROUP, + ) + + +def _index_edges_by_source(edges: dict[str, M.FusionEdge]) -> dict[str, tuple[str, ...]]: + index: dict[str, tuple[str, ...]] = {} + + for edge_name, edge in edges.items(): + index[edge.source] = (*index.get(edge.source, ()), edge_name) + + return index + + +def _index_edges_by_dest(edges: dict[str, M.FusionEdge]) -> dict[str, tuple[str, ...]]: + index: dict[str, tuple[str, ...]] = {} + + for edge_name, edge in edges.items(): + index[edge.dest] = (*index.get(edge.dest, ()), edge_name) + + return index + + +EDGES: dict[str, M.FusionEdge] = { + "linear_mm_to_bias_add": _edge("linear_mm", "linear_bias_add", 0), + "rms_square_to_mean": _edge("rms_square", "rms_mean", 0), + "rms_mean_to_eps_add": _edge("rms_mean", "rms_eps_add", 0), + "rms_eps_add_to_inv": _edge("rms_eps_add", "rms_inv", 0), + "rms_inv_to_scale": _edge("rms_inv", "rms_scale", 1), + "rms_scale_to_weight_mul": _edge("rms_scale", "rms_weight_mul", 0), + "ln_mean_to_center": _edge("ln_mean", "ln_center", 1), + "ln_center_to_square": _edge("ln_center", "ln_square", 0), + "ln_square_to_var": _edge("ln_square", "ln_var", 0), + "ln_var_to_eps_add": _edge("ln_var", "ln_eps_add", 0), + "ln_eps_add_to_inv": _edge("ln_eps_add", "ln_inv", 0), + "ln_center_to_norm": _edge("ln_center", "ln_norm", 0), + "ln_inv_to_norm": _edge("ln_inv", "ln_norm", 1), + "ln_norm_to_weight_mul": _edge("ln_norm", "ln_weight_mul", 0), + "ln_weight_mul_to_bias_add": _edge("ln_weight_mul", "ln_bias_add", 0, optional=True), + "gate_proj_to_activation": _edge("gate_proj", "mlp_activation", 0), + "mlp_activation_to_product": _edge("mlp_activation", "mlp_product", 0), + "up_proj_to_product": _edge("up_proj", "mlp_product", 1), + "mlp_product_to_down_proj": _edge("mlp_product", "down_proj", 0), + "up_proj_to_gelu": _edge("up_proj", "gelu", 0), + "gelu_to_down_proj": _edge("gelu", "down_proj", 0), + "q_proj_to_q_view": _edge("q_proj", "q_view", 0, optional=True, transparent=True), + "k_proj_to_k_view": _edge("k_proj", "k_view", 0, optional=True, transparent=True), + "v_proj_to_v_view": _edge("v_proj", "v_view", 0, optional=True, transparent=True), + "q_view_to_q_rope": _edge("q_view", "q_rope", 0, optional=True, transparent=True), + "k_view_to_k_rope": _edge("k_view", "k_rope", 0, optional=True, transparent=True), + "q_rope_to_qk": _edge("q_rope", "attn_qk", 0), + "k_rope_to_qk": _edge("k_rope", "attn_qk", 1), + "attn_qk_to_scale": _edge("attn_qk", "attn_scale", 0), + "attn_scale_to_mask_add": _edge("attn_scale", "attn_mask_add", 0, optional=True), + "attn_scale_to_softmax": _edge("attn_scale", "attn_softmax", 0), + "attn_mask_add_to_softmax": _edge("attn_mask_add", "attn_softmax", 0, optional=True), + "attn_softmax_to_value": _edge("attn_softmax", "attn_value", 0), + "v_view_to_value": _edge("v_view", "attn_value", 1, optional=True, transparent=True), + "rope_slice_odd_to_neg": _edge("rope_slice_odd", "rope_neg", 0), + "rope_neg_to_rotate_cat": _edge("rope_neg", "rope_rotate_cat", 0), + "rope_slice_even_to_rotate_cat": _edge("rope_slice_even", "rope_rotate_cat", 1), + "rope_cos_mul_to_add": _edge("rope_cos_mul", "rope_add", 0), + "rope_sin_mul_to_add": _edge("rope_sin_mul", "rope_add", 1), + "rope_rotate_cat_to_sin_mul": _edge("rope_rotate_cat", "rope_sin_mul", 0), + "conv_to_bias_add": _edge("conv", "conv_bias_add", 0, optional=True), + "lstm_gate_mm_to_add": _edge("lstm_gate_mm", "lstm_gate_add", 0), + "lstm_recurrent_mm_to_add": _edge("lstm_recurrent_mm", "lstm_gate_add", 1), + "lstm_add_to_sigmoid": _edge("lstm_gate_add", "lstm_sigmoid", 0), + "lstm_add_to_tanh": _edge("lstm_gate_add", "lstm_tanh", 0), + "delta_q_to_decode": _edge("delta_q", "delta_gate", 0, optional=True), + "delta_k_to_decode": _edge("delta_k", "delta_gate", 1, optional=True), + "delta_v_to_decode": _edge("delta_v", "delta_gate", 2, optional=True), + "router_logits_to_probs": _edge("router_logits", "routing_probs", 0), + "routing_probs_to_topk": _edge("routing_probs", "topk", 0), + "expert_gate_to_activation": _expert_edge("expert_gate", "expert_activation", 0), + "expert_activation_to_product": _expert_edge("expert_activation", "expert_product", 0), + "expert_up_to_product": _expert_edge("expert_up", "expert_product", 1), + "expert_product_to_down": _expert_edge("expert_product", "expert_down", 0), + "expert_down_to_weighted": _expert_edge("expert_down", "expert_weighted", 0), + "expert_weighted_to_combine": _expert_edge("expert_weighted", "moe_combine", 0), + "sample_topk_to_softmax": _edge("sample_topk", "sample_softmax", 0), +} + + +EDGE_GROUPS: dict[str, tuple[str, ...]] = { + "linear_bias": ("linear_mm_to_bias_add",), + "rms_norm": ( + "rms_square_to_mean", + "rms_mean_to_eps_add", + "rms_eps_add_to_inv", + "rms_inv_to_scale", + "rms_scale_to_weight_mul", + ), + "layernorm": ( + "ln_mean_to_center", + "ln_center_to_square", + "ln_square_to_var", + "ln_var_to_eps_add", + "ln_eps_add_to_inv", + "ln_center_to_norm", + "ln_inv_to_norm", + "ln_norm_to_weight_mul", + "ln_weight_mul_to_bias_add", + ), + "swiglu_mlp": ( + "gate_proj_to_activation", + "mlp_activation_to_product", + "up_proj_to_product", + "mlp_product_to_down_proj", + ), + "gelu_mlp": ( + "up_proj_to_gelu", + "gelu_to_down_proj", + ), + "attention_core": ( + "attn_qk_to_scale", + "attn_scale_to_mask_add", + "attn_scale_to_softmax", + "attn_mask_add_to_softmax", + "attn_softmax_to_value", + ), + "projected_attention": ( + "q_proj_to_q_view", + "k_proj_to_k_view", + "v_proj_to_v_view", + "q_view_to_q_rope", + "k_view_to_k_rope", + "q_rope_to_qk", + "k_rope_to_qk", + "attn_qk_to_scale", + "attn_scale_to_mask_add", + "attn_mask_add_to_softmax", + "attn_softmax_to_value", + "v_view_to_value", + ), + "rope": ( + "rope_slice_odd_to_neg", + "rope_neg_to_rotate_cat", + "rope_slice_even_to_rotate_cat", + "rope_rotate_cat_to_sin_mul", + "rope_cos_mul_to_add", + "rope_sin_mul_to_add", + ), + "conv_bias": ("conv_to_bias_add",), + "lstm_cell": ( + "lstm_gate_mm_to_add", + "lstm_recurrent_mm_to_add", + "lstm_add_to_sigmoid", + "lstm_add_to_tanh", + ), + "moe_expert": ( + "expert_gate_to_activation", + "expert_activation_to_product", + "expert_up_to_product", + "expert_product_to_down", + "expert_down_to_weighted", + "expert_weighted_to_combine", + ), +} + + +EDGE_NAMES_BY_SOURCE: dict[str, tuple[str, ...]] = _index_edges_by_source(EDGES) +EDGE_NAMES_BY_DEST: dict[str, tuple[str, ...]] = _index_edges_by_dest(EDGES) diff --git a/python/cactus/transpiler/Fusions/fusions.py b/python/cactus/transpiler/Fusions/fusions.py new file mode 100644 index 000000000..f513eb962 --- /dev/null +++ b/python/cactus/transpiler/Fusions/fusions.py @@ -0,0 +1,817 @@ +from . import edges as E +from . import models as M +from . import nodes as N + + +def _nodes(*names: str) -> dict[str, M.FusionNode]: + return {name: N.NODES[name] for name in names} + + +def _edges(*names: str) -> tuple[M.FusionEdge, ...]: + return tuple(E.EDGES[name] for name in names) + + +def _input(role: str, node: str, parent_index: int | None = None, **kwargs) -> M.FusionInput: + return M.FusionInput(role=role, source=M.NodeRef(node=node, parent_index=parent_index), **kwargs) + + +def _variadic_input( + role: str, + node: str, + start_parent_index: int, + *, + min_count: int = 1, + max_count: int | None = None, + end_parent_index: int | None = None, +) -> M.FusionInput: + return M.FusionInput( + role=role, + source=M.NodeRef(node=node, parent_index=start_parent_index), + variadic=True, + min_count=min_count, + max_count=max_count, + end_parent_index=end_parent_index, + ) + + +def _shared_input( + left_node: str, + left_parent_index: int, + right_node: str, + right_parent_index: int, +) -> tuple[M.NodeRef, M.NodeRef]: + return ( + M.NodeRef(left_node, left_parent_index), + M.NodeRef(right_node, right_parent_index), + ) + + +def _output(node: str, role: str = "out", output_index: int | None = None) -> M.FusionOutput: + return M.FusionOutput(role=role, node=node, output_index=output_index) + + +def _cache_input( + role: str, + node: str, + parent_index: int, + *, + cache_kind: str = M.CacheKind.KV, + tensor_role: str | None = None, + optional: bool = False, +) -> M.CacheInput: + return M.CacheInput( + role, + M.NodeRef(node, parent_index), + cache_kind=cache_kind, + tensor_role=tensor_role, + optional=optional, + ) + + +def _cache_output( + role: str, + node: str, + *, + cache_kind: str = M.CacheKind.KV, + tensor_role: str | None = None, +) -> M.CacheOutput: + return M.CacheOutput(role, node, cache_kind=cache_kind, tensor_role=tensor_role) + + +def _cache_mutation( + name: str, + *, + cache_kind: str = M.CacheKind.KV, + read_roles: tuple[str, ...] = (), + write_roles: tuple[str, ...] = (), +) -> M.CacheMutation: + return M.CacheMutation( + name, + cache_kind=cache_kind, + read_roles=read_roles, + write_roles=write_roles, + ) + + +def _required_attrs(**attrs) -> dict: + return {"required_attrs": attrs} + + +def _graph( + name: str, + root: str, + node_names: tuple[str, ...], + *, + edge_names: tuple[str, ...] = (), + inputs: tuple[M.FusionInput, ...] = (), + shared_inputs: tuple[tuple[M.NodeRef, M.NodeRef], ...] = (), + outputs: tuple[M.FusionOutput, ...] | None = None, + attr_captures: tuple[M.AttrCapture, ...] = (), + optional_nodes: tuple[str, ...] = (), + repeated_subgraphs: tuple[M.RepeatedSubgraph, ...] = (), + cache_inputs: tuple[M.CacheInput, ...] = (), + cache_outputs: tuple[M.CacheOutput, ...] = (), + cache_mutations: tuple[M.CacheMutation, ...] = (), + constraints: tuple[str, ...] = (), + variants: tuple[str, ...] = (), + metadata: dict | None = None, + allow_root_external_children: bool = True, +) -> M.FusionGraph: + return M.FusionGraph( + name=name, + root=root, + nodes=_nodes(*node_names), + edges=_edges(*edge_names), + inputs=inputs, + shared_inputs=shared_inputs, + outputs=outputs or (_output(root),), + attr_captures=attr_captures, + optional_nodes=optional_nodes, + repeated_subgraphs=repeated_subgraphs, + cache_inputs=cache_inputs, + cache_outputs=cache_outputs, + cache_mutations=cache_mutations, + constraints=constraints, + variants=variants, + metadata=metadata or {}, + allow_root_external_children=allow_root_external_children, + ) + + +def _single_node_graph( + name: str, + root: str, + input_roles: tuple[str, ...] = (), + *, + attr_captures: tuple[M.AttrCapture, ...] = (), + metadata: dict | None = None, +) -> M.FusionGraph: + return _graph( + name, + root, + (root,), + inputs=tuple(_input(role, root, index) for index, role in enumerate(input_roles)), + attr_captures=attr_captures, + metadata=metadata, + ) + + +def _definition( + name: str, + cactus_op: str, + graph: M.FusionGraph, + *, + fusion_fields: tuple[str, ...] = ("generic",), + supported_inference_modes: tuple[str, ...] = (), + supported_modalities: tuple[str, ...] = (), + metadata: dict | None = None, +) -> M.FusionDefinition: + return M.FusionDefinition( + name=name, + target=f"cactus.{cactus_op}", + graph=graph, + fusion_fields=fusion_fields, + supported_inference_modes=supported_inference_modes, + supported_modalities=supported_modalities, + cactus_op=cactus_op, + metadata=metadata or {}, + ) + + +def _index_by_target(fusions: dict[str, M.FusionDefinition]) -> dict[str, M.FusionDefinition]: + return {fusion.target: fusion for fusion in fusions.values()} + + +def _index_by_cactus_op( + fusions: dict[str, M.FusionDefinition], +) -> dict[str, tuple[M.FusionDefinition, ...]]: + index: dict[str, tuple[M.FusionDefinition, ...]] = {} + + for fusion in fusions.values(): + if fusion.cactus_op is None: + continue + + index[fusion.cactus_op] = (*index.get(fusion.cactus_op, ()), fusion) + + return index + + +def _index_by_field( + fusions: dict[str, M.FusionDefinition], +) -> dict[str, tuple[M.FusionDefinition, ...]]: + index: dict[str, tuple[M.FusionDefinition, ...]] = {} + + for fusion in fusions.values(): + for field in fusion.fusion_fields: + index[field] = (*index.get(field, ()), fusion) + + return index + + +def _index_by_root_op( + fusions: dict[str, M.FusionDefinition], +) -> dict[str, tuple[M.FusionDefinition, ...]]: + index: dict[str, tuple[M.FusionDefinition, ...]] = {} + + for fusion in fusions.values(): + root_node = fusion.graph.nodes.get(fusion.graph.root) + if root_node is None: + continue + + for op in root_node.ops: + index[op] = (*index.get(op, ()), fusion) + + return index + + +DIRECT_GRAPHS: dict[str, M.FusionGraph] = { + "add": _single_node_graph("add", "add", ("a", "b")), + "subtract": _single_node_graph("subtract", "subtract", ("a", "b")), + "multiply": _single_node_graph("multiply", "multiply", ("a", "b")), + "divide": _single_node_graph("divide", "divide", ("a", "b")), + "abs": _single_node_graph("abs", "abs", ("x",)), + "pow": _single_node_graph( + "pow", + "pow", + ("x",), + attr_captures=(M.AttrCapture("exponent", "pow", "exponent", required=False),), + ), + "sqrt": _single_node_graph("sqrt", "sqrt", ("x",)), + "scalar_add": _single_node_graph( + "scalar_add", + "scalar_add", + ("x",), + attr_captures=(M.AttrCapture("value", "scalar_add", "other", required=False),), + ), + "scalar_subtract": _single_node_graph( + "scalar_subtract", + "scalar_subtract", + ("x",), + attr_captures=(M.AttrCapture("value", "scalar_subtract", "other", required=False),), + ), + "scalar_multiply": _single_node_graph( + "scalar_multiply", + "scalar_multiply", + ("x",), + attr_captures=(M.AttrCapture("value", "scalar_multiply", "other", required=False),), + ), + "scalar_divide": _single_node_graph( + "scalar_divide", + "scalar_divide", + ("x",), + attr_captures=(M.AttrCapture("value", "scalar_divide", "other", required=False),), + ), + "scalar_floor_divide": _single_node_graph( + "scalar_floor_divide", + "scalar_floor_divide", + ("x",), + attr_captures=(M.AttrCapture("value", "scalar_floor_divide", "other", required=False),), + ), + "scalar_exp": _single_node_graph("scalar_exp", "scalar_exp", ("x",)), + "scalar_sqrt": _single_node_graph("scalar_sqrt", "scalar_sqrt", ("x",)), + "scalar_cos": _single_node_graph("scalar_cos", "scalar_cos", ("x",)), + "scalar_sin": _single_node_graph("scalar_sin", "scalar_sin", ("x",)), + "scalar_log": _single_node_graph("scalar_log", "scalar_log", ("x",)), + "mean": _single_node_graph( + "mean", + "mean", + ("x",), + attr_captures=(M.AttrCapture("axis", "mean", "dim", default=-1, required=False),), + ), + "sum": _single_node_graph( + "sum", + "sum", + ("x",), + attr_captures=(M.AttrCapture("axis", "sum", "dim", default=-1, required=False),), + ), + "variance": _single_node_graph( + "variance", + "variance", + ("x",), + attr_captures=(M.AttrCapture("axis", "variance", "dim", default=-1, required=False),), + ), + "min": _single_node_graph( + "min", + "min", + ("x",), + attr_captures=(M.AttrCapture("axis", "min", "dim", default=-1, required=False),), + ), + "max": _single_node_graph( + "max", + "max", + ("x",), + attr_captures=(M.AttrCapture("axis", "max", "dim", default=-1, required=False),), + ), + "softmax": _single_node_graph( + "softmax", + "softmax", + ("x",), + attr_captures=(M.AttrCapture("axis", "softmax", "dim", default=-1, required=False),), + ), + "topk": _single_node_graph( + "topk", + "topk_direct", + ("x",), + attr_captures=(M.AttrCapture("k", "topk_direct", "k", required=True),), + ), + "view": _single_node_graph( + "view", + "view", + ("x",), + attr_captures=(M.AttrCapture("shape", "view", "shape", required=False),), + ), + "reshape": _single_node_graph( + "reshape", + "reshape", + ("x",), + attr_captures=(M.AttrCapture("shape", "reshape", "shape", required=False),), + ), + "flatten": _single_node_graph( + "flatten", + "flatten", + ("x",), + attr_captures=( + M.AttrCapture("start_dim", "flatten", "start_dim", default=0, required=False), + M.AttrCapture("end_dim", "flatten", "end_dim", default=-1, required=False), + ), + ), + "transpose": _single_node_graph( + "transpose", + "transpose", + ("x",), + attr_captures=(M.AttrCapture("permutation", "transpose", "permutation", required=False),), + ), + "slice": _single_node_graph( + "slice", + "slice", + ("x",), + attr_captures=( + M.AttrCapture("axis", "slice", "dim", required=False), + M.AttrCapture("start", "slice", "start", required=False), + M.AttrCapture("end", "slice", "end", required=False), + M.AttrCapture("step", "slice", "step", required=False), + ), + ), + "index": _single_node_graph( + "index", + "index", + ("x",), + attr_captures=( + M.AttrCapture("index_value", "index", "index", required=False), + M.AttrCapture("axis", "index", "dim", default=0, required=False), + ), + ), + "cat": _graph("cat", "cat", ("cat",), inputs=(_variadic_input("values", "cat", 0, min_count=2),), attr_captures=(M.AttrCapture("axis", "cat", "dim", default=0, required=False),)), + "gather": _single_node_graph("gather", "gather", ("tensor", "indices")), + "relu": _single_node_graph("relu", "relu", ("x",)), + "silu": _single_node_graph("silu", "silu", ("x",)), + "gelu": _single_node_graph("gelu", "gelu", ("x",)), + "gelu_erf": _single_node_graph("gelu_erf", "gelu_erf", ("x",)), + "sigmoid": _single_node_graph("sigmoid", "sigmoid", ("x",)), + "tanh": _single_node_graph("tanh", "tanh", ("x",)), + "clamp": _single_node_graph( + "clamp", + "clamp", + ("x",), + attr_captures=( + M.AttrCapture("lo", "clamp", "min", required=False), + M.AttrCapture("hi", "clamp", "max", required=False), + ), + ), + "glu": _single_node_graph( + "glu", + "glu", + ("x",), + attr_captures=(M.AttrCapture("axis", "glu", "dim", default=-1, required=False),), + ), + "groupnorm": _single_node_graph( + "groupnorm", + "groupnorm", + ("x", "weight", "bias"), + attr_captures=( + M.AttrCapture("num_groups", "groupnorm", "num_groups", required=False), + M.AttrCapture("epsilon", "groupnorm", "eps", default=1e-5, required=False), + ), + ), + "batchnorm": _single_node_graph( + "batchnorm", + "batchnorm", + ("x", "weight", "bias", "running_mean", "running_var"), + attr_captures=( + M.AttrCapture("axis", "batchnorm", "axis", default=1, required=False), + M.AttrCapture("epsilon", "batchnorm", "eps", default=1e-5, required=False), + ), + ), + "embedding_from_tensor": _single_node_graph( + "embedding_from_tensor", + "embedding", + ("embedding_tensor", "indices"), + ), +} + + +LINEAR_GRAPH = _graph( + "linear", + "linear_mm", + ("linear_mm",), + inputs=(_input("x", "linear_mm", 0), _input("weight", "linear_mm", 1)), + attr_captures=(M.AttrCapture("pretransposed_rhs", default=False, required=False),), + constraints=("Weight layout/pretranspose must be verified against exported attrs or weight metadata.",), +) + +LINEAR_BIAS_GRAPH = _graph( + "linear_bias", + "linear_bias_add", + ("linear_mm", "linear_bias_add"), + edge_names=("linear_mm_to_bias_add",), + inputs=(_input("x", "linear_mm", 0), _input("weight", "linear_mm", 1), _input("bias", "linear_bias_add", 1)), + attr_captures=(M.AttrCapture("pretransposed_rhs", default=False, required=False),), + constraints=("Bias add must consume only the linear output plus a parameter/buffer bias.",), +) + +RMS_NORM_GRAPH = _graph( + "rms_norm", + "rms_weight_mul", + ("rms_square", "rms_mean", "rms_eps_add", "rms_inv", "rms_scale", "rms_weight_mul"), + edge_names=E.EDGE_GROUPS["rms_norm"], + inputs=(_input("x", "rms_square", 0), _input("weight", "rms_weight_mul", 1)), + shared_inputs=(_shared_input("rms_square", 0, "rms_scale", 0),), + attr_captures=(M.AttrCapture("epsilon", "rms_eps_add", "other", default=1e-6, required=False),), + constraints=("Mean dimension must be the hidden dimension.", "Final multiply must consume the RMS-normalized activation and norm weight."), +) + +LAYERNORM_GRAPH = _graph( + "layernorm", + "ln_bias_add", + ("ln_mean", "ln_center", "ln_square", "ln_var", "ln_eps_add", "ln_inv", "ln_norm", "ln_weight_mul", "ln_bias_add"), + edge_names=E.EDGE_GROUPS["layernorm"], + inputs=(_input("x", "ln_mean", 0), _input("weight", "ln_weight_mul", 1), _input("bias", "ln_bias_add", 1, optional=True)), + shared_inputs=(_shared_input("ln_mean", 0, "ln_center", 0),), + attr_captures=(M.AttrCapture("epsilon", "ln_eps_add", "other", default=1e-5, required=False),), + optional_nodes=("ln_bias_add",), + constraints=("Mean/variance dimensions must match the normalized feature axis.",), +) + +SWIGLU_MLP_GRAPH = _graph( + "swiglu_mlp", + "down_proj", + ("gate_proj", "up_proj", "mlp_activation", "mlp_product", "down_proj"), + edge_names=E.EDGE_GROUPS["swiglu_mlp"], + inputs=(_input("hidden", "gate_proj", 0), _input("gate_weight", "gate_proj", 1), _input("up_weight", "up_proj", 1), _input("down_weight", "down_proj", 1)), + shared_inputs=(_shared_input("gate_proj", 0, "up_proj", 0),), + attr_captures=(M.AttrCapture("product_scale", default=1.0, required=False),), + constraints=("gate_proj and up_proj must consume the same hidden input.", "down_proj must consume silu(gate_proj(hidden)) * up_proj(hidden)."), +) + +GELU_MLP_GRAPH = _graph( + "gelu_mlp", + "down_proj", + ("up_proj", "gelu", "down_proj"), + inputs=(_input("hidden", "up_proj", 0), _input("up_weight", "up_proj", 1), _input("down_weight", "down_proj", 1)), + edge_names=E.EDGE_GROUPS["gelu_mlp"], + constraints=("This pattern uses gelu as the activation before the down projection; IR should bind the activation output to down_proj input 0.",), +) + +ATTENTION_DIRECT_GRAPH = _graph( + "scaled_dot_product_attention", + "attn_sdp", + ("attn_sdp",), + inputs=(_input("query", "attn_sdp", 0), _input("key", "attn_sdp", 1), _input("value", "attn_sdp", 2), _input("mask", "attn_sdp", 3, optional=True)), + attr_captures=(M.AttrCapture("scale", "attn_sdp", "scale", required=False), M.AttrCapture("is_causal", "attn_sdp", "is_causal", default=True, required=False)), +) + +ATTENTION_CORE_GRAPH = _graph( + "attention_core", + "attn_value", + ("attn_qk", "attn_scale", "attn_softmax", "attn_value"), + edge_names=("attn_qk_to_scale", "attn_scale_to_softmax", "attn_softmax_to_value"), + inputs=(_input("query", "attn_qk", 0), _input("key", "attn_qk", 1), _input("value", "attn_value", 1)), + attr_captures=(M.AttrCapture("scale", "attn_scale", "other", required=False), M.AttrCapture("is_causal", default=True, required=False), M.AttrCapture("window_size", default=0, required=False)), + constraints=("QK matmul must be over query/key head dimensions.", "Softmax axis must be the attention score axis.", "Value matmul must consume attention probabilities and value states."), +) + +ATTENTION_MASKED_GRAPH = _graph( + "attention_masked", + "attn_value", + ("attn_qk", "attn_scale", "attn_mask_add", "attn_softmax", "attn_value"), + edge_names=("attn_qk_to_scale", "attn_scale_to_mask_add", "attn_mask_add_to_softmax", "attn_softmax_to_value"), + inputs=(_input("query", "attn_qk", 0), _input("key", "attn_qk", 1), _input("mask", "attn_mask_add", 1), _input("value", "attn_value", 1)), + attr_captures=(M.AttrCapture("scale", "attn_scale", "other", required=False), M.AttrCapture("use_mask", default=True), M.AttrCapture("additive_mask", default=True, required=False)), + constraints=("Mask input must be broadcast-compatible with attention scores.",), +) + +ROPE_GRAPH = _graph( + "rope", + "rope_add", + ("rope_slice_even", "rope_slice_odd", "rope_neg", "rope_rotate_cat", "rope_cos_mul", "rope_sin_mul", "rope_add"), + edge_names=E.EDGE_GROUPS["rope"], + inputs=(_input("x", "rope_slice_even", 0), _input("cos", "rope_cos_mul", 1), _input("sin", "rope_sin_mul", 1)), + shared_inputs=( + _shared_input("rope_slice_even", 0, "rope_slice_odd", 0), + _shared_input("rope_slice_even", 0, "rope_cos_mul", 0), + ), + attr_captures=(M.AttrCapture("theta", default=None, required=False), M.AttrCapture("position_offset", default=0, required=False)), + constraints=("RoPE cos/sin tables or generated constants must correspond to the same theta and position offset.",), +) + +CONV_GRAPH = _graph( + "conv", + "conv", + ("conv",), + inputs=(_input("x", "conv", 0), _input("weight", "conv", 1), _input("bias", "conv", 2, optional=True)), + attr_captures=(M.AttrCapture("stride", "conv", "stride", default=1, required=False), M.AttrCapture("padding", "conv", "padding", default=0, required=False), M.AttrCapture("dilation", "conv", "dilation", default=1, required=False), M.AttrCapture("groups", "conv", "groups", default=1, required=False)), +) + +CONV_BIAS_GRAPH = _graph( + "conv_bias", + "conv_bias_add", + ("conv", "conv_bias_add"), + edge_names=("conv_to_bias_add",), + inputs=(_input("x", "conv", 0), _input("weight", "conv", 1), _input("bias", "conv_bias_add", 1)), + attr_captures=(M.AttrCapture("stride", "conv", "stride", default=1, required=False), M.AttrCapture("padding", "conv", "padding", default=0, required=False), M.AttrCapture("dilation", "conv", "dilation", default=1, required=False), M.AttrCapture("groups", "conv", "groups", default=1, required=False)), +) + +KV_CACHE_APPEND_GRAPH = _graph( + "kv_cache_append", + "kv_cache_append", + ("kv_cache_append",), + inputs=(_input("new_kv", "kv_cache_append", 0), _input("cache_state", "kv_cache_append", 1)), + cache_inputs=(_cache_input("cache_state", "kv_cache_append", 1, tensor_role=M.CacheTensorRole.STATE),), + cache_outputs=(_cache_output("updated_cache_state", "kv_cache_append", tensor_role=M.CacheTensorRole.STATE),), + cache_mutations=(_cache_mutation("append_kv", read_roles=("cache_state",), write_roles=("updated_cache_state",)),), +) + +ATTENTION_CACHED_GRAPH = _graph( + "attention_cached", + "attn_value", + ("attn_qk", "attn_scale", "attn_softmax", "attn_value"), + edge_names=("attn_qk_to_scale", "attn_scale_to_softmax", "attn_softmax_to_value"), + inputs=(_input("query", "attn_qk", 0), _input("key_new", "attn_qk", 1), _input("value_new", "attn_value", 1)), + cache_inputs=( + _cache_input("key_cache_state", "attn_qk", 1, tensor_role=M.CacheTensorRole.KEY, optional=True), + _cache_input("value_cache_state", "attn_value", 1, tensor_role=M.CacheTensorRole.VALUE, optional=True), + ), + attr_captures=(M.AttrCapture("scale", "attn_scale", "other", required=False), M.AttrCapture("position_offset", default=0, required=False), M.AttrCapture("window_size", default=0, required=False)), + constraints=("Cached attention must verify that key/value cache states belong to the same decoder layer.",), +) + +EXPERT_BRANCH_GRAPH = _graph( + "moe_expert_branch", + "expert_weighted", + ("expert_gate", "expert_up", "expert_activation", "expert_product", "expert_down", "expert_weighted"), + edge_names=("expert_gate_to_activation", "expert_activation_to_product", "expert_up_to_product", "expert_product_to_down", "expert_down_to_weighted"), + inputs=(_input("hidden", "expert_gate", 0), _input("gate_weight", "expert_gate", 1), _input("up_weight", "expert_up", 1), _input("down_weight", "expert_down", 1), _input("routing_weight", "expert_weighted", 1)), + shared_inputs=(_shared_input("expert_gate", 0, "expert_up", 0),), + constraints=("Each repeated expert branch must correspond to the same expert index selected by routing.",), + allow_root_external_children=False, +) + +MOE_GATED_GRAPH = _graph( + "moe_layer_gated", + "moe_combine", + ("router_logits", "routing_probs", "topk", "moe_combine"), + edge_names=("router_logits_to_probs", "routing_probs_to_topk"), + inputs=(_input("hidden", "router_logits", 0), _input("router_weight", "router_logits", 1)), + repeated_subgraphs=(M.RepeatedSubgraph("experts", EXPERT_BRANCH_GRAPH, min_count=1, anchor_node="topk"),), + attr_captures=(M.AttrCapture("num_experts", default=None, required=False), M.AttrCapture("num_experts_per_tok", "topk", "k", required=False), M.AttrCapture("normalize_routing", default=True, required=False), M.AttrCapture("epsilon", default=1e-6, required=False), M.AttrCapture("routed_scaling_factor", default=1.0, required=False)), + constraints=("Top-k indices must select the repeated expert branches.", "Routing probabilities must be the weights used for expert output combination."), +) + +LSTM_CELL_GRAPH = _graph( + "lstm_cell", + "lstm_tanh", + ("lstm_gate_mm", "lstm_recurrent_mm", "lstm_gate_add", "lstm_sigmoid", "lstm_tanh"), + edge_names=E.EDGE_GROUPS["lstm_cell"], + inputs=(_input("input", "lstm_gate_mm", 0), _input("h_prev", "lstm_recurrent_mm", 0), _input("weight_ih", "lstm_gate_mm", 1), _input("weight_hh", "lstm_recurrent_mm", 1)), + constraints=("IR must verify the four LSTM gate splits and c_prev update path before using this fusion.",), +) + +DELTANET_DECODE_GRAPH = _graph( + "gated_deltanet_decode", + "delta_gate", + ("delta_q", "delta_k", "delta_v", "delta_gate"), + inputs=(_input("query", "delta_q", 0), _input("key", "delta_k", 0), _input("value", "delta_v", 0), _input("gate_log", "delta_gate", 0), _input("beta", "delta_gate", 1), _input("initial_state", "delta_gate", 2)), + attr_captures=(M.AttrCapture("scale", default=None, required=False),), + cache_inputs=(_cache_input("initial_state", "delta_gate", 2, cache_kind=M.CacheKind.RECURRENT, tensor_role=M.CacheTensorRole.STATE),), +) + +DELTANET_PREFILL_GRAPH = _graph( + "gated_deltanet_prefill", + "delta_prefill", + ("delta_q", "delta_k", "delta_v", "delta_gate", "delta_prefill"), + inputs=(_input("query", "delta_q", 0), _input("key", "delta_k", 0), _input("value", "delta_v", 0), _input("gate_log", "delta_gate", 0), _input("beta", "delta_gate", 1), _input("initial_state", "delta_prefill", 2)), + attr_captures=(M.AttrCapture("chunk_size", "delta_prefill", "chunk_size", required=False), M.AttrCapture("scale", default=None, required=False),), + cache_inputs=(_cache_input("initial_state", "delta_prefill", 2, cache_kind=M.CacheKind.RECURRENT, tensor_role=M.CacheTensorRole.STATE),), +) + +REL_POS_BIAS_GRAPH = _graph( + "rel_pos_bias", + "rel_pos_bias", + ("rel_pos_bias",), + inputs=(_input("query", "rel_pos_bias", 0), _input("relative_key", "rel_pos_bias", 1)), + attr_captures=(M.AttrCapture("scale", "rel_pos_bias", "scale", required=False),), +) + +KV_CACHE_STATE_GRAPH = _graph( + "kv_cache_state", + "kv_cache_state", + ("kv_cache_state",), + outputs=(M.FusionOutput("cache_state", "kv_cache_state"),), + attr_captures=(M.AttrCapture("max_seq_len", default=None, required=False), M.AttrCapture("num_kv_heads", default=None, required=False), M.AttrCapture("head_dim", default=None, required=False), M.AttrCapture("window_size", default=0, required=False), M.AttrCapture("sink_size", default=0, required=False), M.AttrCapture("num_slots", default=1, required=False)), +) + +CONV_CACHE_STATE_GRAPH = _graph( + "conv_cache_state", + "conv_cache_state", + ("conv_cache_state",), + outputs=(M.FusionOutput("cache_state", "conv_cache_state"),), + attr_captures=(M.AttrCapture("window_size", default=None, required=False), M.AttrCapture("hidden_dim", default=None, required=False)), +) + +CONV_CACHE_APPEND_GRAPH = _graph( + "conv_cache_append", + "conv_cache_append", + ("conv_cache_append",), + inputs=(_input("new_data", "conv_cache_append", 0), _input("cache_state", "conv_cache_append", 1)), + cache_inputs=(_cache_input("cache_state", "conv_cache_append", 1, cache_kind=M.CacheKind.CONV, tensor_role=M.CacheTensorRole.STATE),), + cache_outputs=(_cache_output("updated_cache_state", "conv_cache_append", cache_kind=M.CacheKind.CONV, tensor_role=M.CacheTensorRole.STATE),), + cache_mutations=(_cache_mutation("append_conv_state", cache_kind=M.CacheKind.CONV, read_roles=("cache_state",), write_roles=("updated_cache_state",)),), +) + +CONV_CACHE_INITIALIZE_GRAPH = _graph( + "conv_cache_initialize", + "conv_cache_initialize", + ("conv_cache_initialize",), + inputs=(_input("rows", "conv_cache_initialize", 0), _input("cache_state", "conv_cache_initialize", 1)), + cache_inputs=(_cache_input("cache_state", "conv_cache_initialize", 1, cache_kind=M.CacheKind.CONV, tensor_role=M.CacheTensorRole.STATE),), + cache_outputs=(_cache_output("initialized_cache_state", "conv_cache_initialize", cache_kind=M.CacheKind.CONV, tensor_role=M.CacheTensorRole.STATE),), +) + +RECURRENT_CACHE_STATE_GRAPH = _graph( + "recurrent_cache_state", + "recurrent_cache_state", + ("recurrent_cache_state",), + outputs=(M.FusionOutput("cache_state", "recurrent_cache_state"),), + attr_captures=(M.AttrCapture("shape", default=None, required=False), M.AttrCapture("precision", default=None, required=False)), +) + +RECURRENT_CACHE_WRITE_GRAPH = _graph( + "recurrent_cache_write", + "recurrent_cache_write", + ("recurrent_cache_write",), + inputs=(_input("new_value", "recurrent_cache_write", 0), _input("cache_input", "recurrent_cache_write", 1)), + cache_inputs=(_cache_input("cache_input", "recurrent_cache_write", 1, cache_kind=M.CacheKind.RECURRENT, tensor_role=M.CacheTensorRole.STATE),), + cache_outputs=(_cache_output("updated_cache", "recurrent_cache_write", cache_kind=M.CacheKind.RECURRENT, tensor_role=M.CacheTensorRole.STATE),), + cache_mutations=(_cache_mutation("write_recurrent_state", cache_kind=M.CacheKind.RECURRENT, read_roles=("cache_input",), write_roles=("updated_cache",)),), +) + +SAMPLE_GRAPH = _graph( + "sample", + "sample_softmax", + ("sample_topk", "sample_softmax"), + edge_names=("sample_topk_to_softmax",), + inputs=(_input("logits", "sample_topk", 0),), + attr_captures=(M.AttrCapture("temperature", default=0.6, required=False), M.AttrCapture("top_p", default=0.95, required=False), M.AttrCapture("top_k", "sample_topk", "k", default=20, required=False)), +) + +SCATTER_TOPK_GRAPH = _graph( + "scatter_topk", + "scatter", + ("scatter",), + inputs=(_input("indices", "scatter", 0), _input("values", "scatter", 1)), + attr_captures=(M.AttrCapture("num_classes", "scatter", "num_classes", required=False),), +) + +GAUSSIAN_TOPK_GRAPH = _graph( + "gaussian_topk", + "gaussian_topk", + ("gaussian_topk",), + inputs=(_input("x", "gaussian_topk", 0),), + attr_captures=(M.AttrCapture("ppf", "gaussian_topk", "ppf", required=False),), +) + +ALTUP_PREDICT_GRAPH = _graph( + "altup_predict", + "altup_predict", + ("altup_predict",), + inputs=(_input("coefs", "altup_predict", 0), _variadic_input("streams", "altup_predict", 1)), +) + +ALTUP_CORRECT_GRAPH = _graph( + "altup_correct", + "altup_correct", + ("altup_correct",), + inputs=(_input("coefs", "altup_correct", 0), _input("innovation", "altup_correct", 1), _variadic_input("predictions", "altup_correct", 2)), +) + +DSP_GRAPHS: dict[str, M.FusionGraph] = { + "stft": _graph("stft", "stft", ("stft",), inputs=(_input("x", "stft", 0), _input("weight", "stft", 1)), attr_captures=(M.AttrCapture("stride", "stft", "stride", required=False), M.AttrCapture("num_fft_bins", "stft", "num_fft_bins", required=False))), + "rfft": _graph("rfft", "rfft", ("rfft",), inputs=(_input("x", "rfft", 0),)), + "irfft": _graph("irfft", "irfft", ("irfft",), inputs=(_input("x", "irfft", 0),), attr_captures=(M.AttrCapture("output_length", "irfft", "output_length", required=False),)), + "spectrogram": _graph("spectrogram", "spectrogram", ("spectrogram",), inputs=(_input("waveform", "spectrogram", 0), _input("mel_filters", "spectrogram", 1)), metadata={"requires_preprocessing_attrs": True}), + "image_preprocess": _graph("image_preprocess", "image_preprocess", ("image_preprocess",), inputs=(_input("pixel_input", "image_preprocess", 0),), metadata={"requires_preprocessing_attrs": True}), + "bilinear_interpolation": _graph("bilinear_interpolation", "bilinear_interpolation", ("bilinear_interpolation",), inputs=(_input("pos_embeds", "bilinear_interpolation", 0),), attr_captures=(M.AttrCapture("dst_height", "bilinear_interpolation", "dst_height", required=False), M.AttrCapture("dst_width", "bilinear_interpolation", "dst_width", required=False))), +} + + +GRAPH_BY_NAME: dict[str, M.FusionGraph] = { + **DIRECT_GRAPHS, + **DSP_GRAPHS, + "linear": LINEAR_GRAPH, + "linear_bias": LINEAR_BIAS_GRAPH, + "rms_norm": RMS_NORM_GRAPH, + "layernorm": LAYERNORM_GRAPH, + "swiglu_mlp": SWIGLU_MLP_GRAPH, + "gelu_mlp": GELU_MLP_GRAPH, + "scaled_dot_product_attention": ATTENTION_DIRECT_GRAPH, + "attention_core": ATTENTION_CORE_GRAPH, + "attention_masked": ATTENTION_MASKED_GRAPH, + "rope": ROPE_GRAPH, + "conv": CONV_GRAPH, + "conv_bias": CONV_BIAS_GRAPH, + "kv_cache_append": KV_CACHE_APPEND_GRAPH, + "attention_cached": ATTENTION_CACHED_GRAPH, + "moe_layer_gated": MOE_GATED_GRAPH, + "lstm_cell": LSTM_CELL_GRAPH, + "gated_deltanet_decode": DELTANET_DECODE_GRAPH, + "gated_deltanet_prefill": DELTANET_PREFILL_GRAPH, + "rel_pos_bias": REL_POS_BIAS_GRAPH, + "kv_cache_state": KV_CACHE_STATE_GRAPH, + "conv_cache_state": CONV_CACHE_STATE_GRAPH, + "conv_cache_append": CONV_CACHE_APPEND_GRAPH, + "conv_cache_initialize": CONV_CACHE_INITIALIZE_GRAPH, + "recurrent_cache_state": RECURRENT_CACHE_STATE_GRAPH, + "recurrent_cache_write": RECURRENT_CACHE_WRITE_GRAPH, + "sample": SAMPLE_GRAPH, + "scatter_topk": SCATTER_TOPK_GRAPH, + "gaussian_topk": GAUSSIAN_TOPK_GRAPH, + "altup_predict": ALTUP_PREDICT_GRAPH, + "altup_correct": ALTUP_CORRECT_GRAPH, +} + + +DIRECT_CACTUS_OPS: dict[str, str] = { + "sqrt": "scalar_sqrt", +} + +DIRECT_DEFINITIONS: dict[str, M.FusionDefinition] = { + name: _definition(name, cactus_op=DIRECT_CACTUS_OPS.get(name, name), graph=graph, fusion_fields=("generic", "direct")) + for name, graph in DIRECT_GRAPHS.items() +} + + +FUSIONS: dict[str, M.FusionDefinition] = { + **DIRECT_DEFINITIONS, + "linear": _definition("linear", "matmul", LINEAR_GRAPH, fusion_fields=("generic", "linear")), + "linear_bias": _definition("linear_bias", "linear", LINEAR_BIAS_GRAPH, fusion_fields=("generic", "linear")), + "rms_norm": _definition("rms_norm", "rms_norm", RMS_NORM_GRAPH, fusion_fields=("generic", "rmsnorm", "gemma4_rmsnorm", "qwen2_5_rmsnorm")), + "layernorm": _definition("layernorm", "layernorm", LAYERNORM_GRAPH, fusion_fields=("generic", "normalization")), + "swiglu_mlp": _definition("swiglu_mlp", "dense_mlp_tq_fused", SWIGLU_MLP_GRAPH, fusion_fields=("generic", "mlp", "gemma4_mlp", "qwen2_5_mlp", "lfm_mlp")), + "gelu_mlp": _definition("gelu_mlp", "matmul", GELU_MLP_GRAPH, fusion_fields=("generic", "mlp")), + "scaled_dot_product_attention": _definition("scaled_dot_product_attention", "attention", ATTENTION_DIRECT_GRAPH, fusion_fields=("generic", "attention")), + "attention_core": _definition("attention_core", "attention", ATTENTION_CORE_GRAPH, fusion_fields=("generic", "attention", "gemma4_attention", "qwen2_5_attention", "lfm_attention")), + "attention_masked": _definition("attention_masked", "attention", ATTENTION_MASKED_GRAPH, fusion_fields=("generic", "attention", "gemma4_attention")), + "rope": _definition("rope", "rope", ROPE_GRAPH, fusion_fields=("generic", "rope", "gemma4_rope", "qwen2_5_rope")), + "conv": _definition("conv", "conv1d", CONV_GRAPH, fusion_fields=("generic", "conv", "audio", "vision")), + "conv_bias": _definition("conv_bias", "conv1d", CONV_BIAS_GRAPH, fusion_fields=("generic", "conv", "audio", "vision")), + "conv1d_k3": _definition("conv1d_k3", "conv1d_k3", CONV_GRAPH, fusion_fields=("generic", "conv", "audio"), metadata=_required_attrs(kernel_size=3)), + "conv1d_k7s3": _definition("conv1d_k7s3", "conv1d_k7s3", CONV_BIAS_GRAPH, fusion_fields=("generic", "conv", "audio"), metadata=_required_attrs(kernel_size=7, stride=3)), + "conv1d_causal": _definition("conv1d_causal", "conv1d_causal", CONV_GRAPH, fusion_fields=("generic", "conv", "cache"), metadata=_required_attrs(causal=True)), + "conv1d_pointwise": _definition("conv1d_pointwise", "conv1d_pointwise", CONV_GRAPH, fusion_fields=("generic", "conv", "audio"), metadata=_required_attrs(kernel_size=1)), + "conv1d_same_depthwise_k9": _definition("conv1d_same_depthwise_k9", "conv1d_same_depthwise_k9", CONV_GRAPH, fusion_fields=("generic", "conv", "audio"), metadata=_required_attrs(kernel_size=9, depthwise=True)), + "conv2d_k3s2p1": _definition("conv2d_k3s2p1", "conv2d_k3s2p1", CONV_GRAPH, fusion_fields=("generic", "conv", "vision"), metadata=_required_attrs(kernel_size=3, stride=2, padding=1)), + "conv2d_depthwise_k3s2p1": _definition("conv2d_depthwise_k3s2p1", "conv2d_depthwise_k3s2p1", CONV_GRAPH, fusion_fields=("generic", "conv", "vision"), metadata=_required_attrs(kernel_size=3, stride=2, padding=1, depthwise=True)), + "conv2d_pointwise_1x1": _definition("conv2d_pointwise_1x1", "conv2d_pointwise_1x1", CONV_GRAPH, fusion_fields=("generic", "conv", "vision"), metadata=_required_attrs(kernel_size=1)), + "kv_cache_append": _definition("kv_cache_append", "kv_cache_append", KV_CACHE_APPEND_GRAPH, fusion_fields=("generic", "cache"), supported_inference_modes=("prefill_with_cache", "decode_with_cache")), + "attention_cached": _definition("attention_cached", "attention_cached", ATTENTION_CACHED_GRAPH, fusion_fields=("generic", "attention", "cache"), supported_inference_modes=("decode_with_cache",)), + "moe_layer_gated": _definition("moe_layer_gated", "moe_layer_gated", MOE_GATED_GRAPH, fusion_fields=("generic", "moe")), + "lstm_cell": _definition("lstm_cell", "lstm_cell", LSTM_CELL_GRAPH, fusion_fields=("generic", "recurrent", "audio")), + "gated_deltanet_decode": _definition("gated_deltanet_decode", "gated_deltanet_decode", DELTANET_DECODE_GRAPH, fusion_fields=("generic", "recurrent", "cache")), + "gated_deltanet_prefill": _definition("gated_deltanet_prefill", "gated_deltanet_prefill", DELTANET_PREFILL_GRAPH, fusion_fields=("generic", "recurrent", "cache")), + "rel_pos_bias": _definition("rel_pos_bias", "rel_pos_bias", REL_POS_BIAS_GRAPH, fusion_fields=("generic", "attention")), + "kv_cache_state": _definition("kv_cache_state", "kv_cache_state", KV_CACHE_STATE_GRAPH, fusion_fields=("generic", "cache"), supported_inference_modes=("prefill_with_cache", "decode_with_cache")), + "conv_cache_state": _definition("conv_cache_state", "conv_cache_state", CONV_CACHE_STATE_GRAPH, fusion_fields=("generic", "cache", "conv")), + "conv_cache_append": _definition("conv_cache_append", "conv_cache_append", CONV_CACHE_APPEND_GRAPH, fusion_fields=("generic", "cache", "conv")), + "conv_cache_initialize": _definition("conv_cache_initialize", "conv_cache_initialize", CONV_CACHE_INITIALIZE_GRAPH, fusion_fields=("generic", "cache", "conv")), + "recurrent_cache_state": _definition("recurrent_cache_state", "recurrent_cache_state", RECURRENT_CACHE_STATE_GRAPH, fusion_fields=("generic", "cache", "recurrent")), + "recurrent_cache_write": _definition("recurrent_cache_write", "recurrent_cache_write", RECURRENT_CACHE_WRITE_GRAPH, fusion_fields=("generic", "cache", "recurrent")), + "sample": _definition("sample", "sample", SAMPLE_GRAPH, fusion_fields=("generic", "sample")), + "scatter_topk": _definition("scatter_topk", "scatter_topk", SCATTER_TOPK_GRAPH, fusion_fields=("generic", "sample")), + "gaussian_topk": _definition("gaussian_topk", "gaussian_topk", GAUSSIAN_TOPK_GRAPH, fusion_fields=("generic", "sample")), + "altup_predict": _definition("altup_predict", "altup_predict", ALTUP_PREDICT_GRAPH, fusion_fields=("generic", "altup")), + "altup_correct": _definition("altup_correct", "altup_correct", ALTUP_CORRECT_GRAPH, fusion_fields=("generic", "altup")), + "stft": _definition("stft", "stft", DSP_GRAPHS["stft"], fusion_fields=("generic", "audio", "dsp")), + "rfft": _definition("rfft", "rfft", DSP_GRAPHS["rfft"], fusion_fields=("generic", "audio", "dsp")), + "irfft": _definition("irfft", "irfft", DSP_GRAPHS["irfft"], fusion_fields=("generic", "audio", "dsp")), + "spectrogram": _definition("spectrogram", "spectrogram", DSP_GRAPHS["spectrogram"], fusion_fields=("generic", "audio", "dsp")), + "image_preprocess": _definition("image_preprocess", "image_preprocess", DSP_GRAPHS["image_preprocess"], fusion_fields=("generic", "vision", "preprocess")), + "bilinear_interpolation": _definition("bilinear_interpolation", "bilinear_interpolation", DSP_GRAPHS["bilinear_interpolation"], fusion_fields=("generic", "vision")), +} + + +FUSION_CATALOG = M.FusionCatalog(fusions=tuple(FUSIONS.values())) + +FUSIONS_BY_TARGET: dict[str, M.FusionDefinition] = _index_by_target(FUSIONS) +FUSIONS_BY_CACTUS_OP: dict[str, tuple[M.FusionDefinition, ...]] = _index_by_cactus_op(FUSIONS) +FUSIONS_BY_FIELD: dict[str, tuple[M.FusionDefinition, ...]] = _index_by_field(FUSIONS) +FUSIONS_BY_ROOT_OP: dict[str, tuple[M.FusionDefinition, ...]] = _index_by_root_op(FUSIONS) diff --git a/python/cactus/transpiler/Fusions/models.py b/python/cactus/transpiler/Fusions/models.py index 33fb58c8b..bcd3d0fe3 100644 --- a/python/cactus/transpiler/Fusions/models.py +++ b/python/cactus/transpiler/Fusions/models.py @@ -1,20 +1,364 @@ -from dataclasses import dataclass +from dataclasses import dataclass, field from typing import Any +class ValueKind: + """ + Names the semantic role of a real IR value. + + IR matchers use this to tell the difference between parameters, + buffers, runtime inputs, activations, constants, outputs, and cache + tensors without depending on fragile node-name strings. + """ + + UNKNOWN = "unknown" + PARAMETER = "parameter" + BUFFER = "buffer" + USER_INPUT = "user_input" + LIFTED_CONSTANT = "lifted_constant" + ACTIVATION = "activation" + OUTPUT = "output" + CACHE_INPUT = "cache_input" + CACHE_OUTPUT = "cache_output" + CACHE_STATE = "cache_state" + + +class CacheKind: + """ + Names the broad family of cache/state a fusion consumes or produces. + + KV is for transformer key/value cache, CONV is for rolling convolution + state, and RECURRENT is for state-space/RNN-style hidden state. Runtime + planning can use this to route cache tensors differently from weights. + """ + + KV = "kv" + CONV = "conv" + RECURRENT = "recurrent" + + +class CacheTensorRole: + """ + Names the role of one tensor inside a cache family. + + This lets a cache-aware fusion say whether a tensor is a key cache, + value cache, generic recurrent/conv state, or cache position tensor. + """ + + KEY = "key" + VALUE = "value" + STATE = "state" + POSITION = "position" @dataclass(slots=True) -class SynthNode: - attrs: dict[str, Any] +class TensorMeta: + """ + Stores observed tensor metadata from the exported graph. + + The IR graph can attach this to real nodes so fusion matchers can compare + shapes, dtypes, devices, grad flags, and strides against TensorConstraint + without rereading the original torch node metadata each time. + """ + + shape: list[Any] = field(default_factory=list) + dtype: str | None = None + device: str | None = None + requires_grad: bool | None = None + stride: list[Any] | None = None + + +@dataclass(slots=True) +class TensorConstraint: + """ + Describes tensor properties a matched real value must satisfy. + + Fusion patterns use this for rank/dtype/shape checks, fixed dimension + checks, and cross-node dimension relationships such as "query head_dim + equals key head_dim". The actual comparison logic belongs in IR. + """ + + rank: int | None = None + min_rank: int | None = None + max_rank: int | None = None + dtype: str | None = None + shape: tuple[Any, ...] = () + dim_equals: tuple[tuple[int, Any], ...] = () + same_dim_as: tuple[tuple[int, str, int], ...] = () + metadata: dict[str, Any] = field(default_factory=dict) + @dataclass(slots=True) -class Route: - pass +class AttrConstraint: + """ + Describes one required attribute condition for a pattern node. + + Examples are pow exponent equals 2, softmax dim equals -1, or conv stride + equals a target kernel variant. source_node/source_attr allow an IR matcher + to compare this attr to another matched node's attr when needed. + """ + name: str + value: Any = None + source_node: str | None = None + source_attr: str | None = None + comparator: str = "eq" + required: bool = True + metadata: dict[str, Any] = field(default_factory=dict) @dataclass(slots=True) -class SynthGraph: - source_node: SynthNode - route: Route \ No newline at end of file +class AttrCapture: + """ + Describes which raw-node attribute becomes a fused-node attribute. + + A match can use this to copy values like epsilon, axis, stride, padding, + scale, or top-k from the matched graph into the simplified Cactus-facing + fused node. default is used when the raw graph does not expose the attr. + """ + + name: str + source_node: str | None = None + source_attr: str | None = None + default: Any = None + required: bool = True + metadata: dict[str, Any] = field(default_factory=dict) + + +@dataclass(slots=True) +class NodeRef: + """ + Points to a pattern node and optionally one of its input/output positions. + + FusionInput, shared-input declarations, and cache declarations use this + small reference instead of embedding matcher behavior in the Fusion schema. + parent_index refers to the matched node's ordered parents in the IR DAG. + """ + + node: str + parent_index: int | None = None + output_index: int | None = None + + +@dataclass(slots=True) +class FusionNode: + """ + Describes one synthetic node inside a fusion pattern graph. + + It stores the allowed raw op targets, optional node-level attr/tensor + constraints, and structural hints such as optional, transparent, or repeated. + IR matchers bind these descriptions to real exported nodes during fusion. + """ + + name: str + ops: tuple[str, ...] = () + attrs: tuple[AttrConstraint, ...] = () + optional: bool = False + transparent: bool = False + repeated: bool = False + repeated_group: str | None = None + allowed_value_kinds: tuple[str, ...] = () + tensor_constraints: tuple[TensorConstraint, ...] = () + metadata: dict[str, Any] = field(default_factory=dict) + + +@dataclass(slots=True) +class FusionEdge: + """ + Describes one directed producer-to-consumer edge in a fusion pattern. + + source/dest identify synthetic nodes, dest_input_index preserves meaningful + input ordering from FX args, and optional/transparent/repeated flags let IR + match common graph variants without changing the core FusionGraph shape. + """ + + source: str + dest: str + dest_input_index: int | None = None + source_output_index: int | None = None + optional: bool = False + transparent: bool = False + repeated: bool = False + repeated_group: str | None = None + metadata: dict[str, Any] = field(default_factory=dict) + + +@dataclass(slots=True) +class FusionInput: + """ + Describes an external input that must remain visible after fusion. + + The source points to the synthetic node input position where the external + real node appears. variadic/min_count/max_count support ops like cat or + AltUp that accept a variable number of parents. + """ + + role: str + source: NodeRef + optional: bool = False + variadic: bool = False + min_count: int = 1 + max_count: int | None = None + end_parent_index: int | None = None + repeated_group: str | None = None + allowed_value_kinds: tuple[str, ...] = () + tensor_constraints: tuple[TensorConstraint, ...] = () + metadata: dict[str, Any] = field(default_factory=dict) + + +@dataclass(slots=True) +class FusionOutput: + """ + Describes an output produced by the fused op. + + Most fusions expose the root node output, but multi-output/cache-aware + fusions can declare multiple outputs with roles so the simplified graph + and runtime plan know how to reconnect downstream consumers. + """ + + role: str + node: str + output_index: int | None = None + tensor_constraints: tuple[TensorConstraint, ...] = () + metadata: dict[str, Any] = field(default_factory=dict) + + +@dataclass(slots=True) +class CacheInput: + """ + Describes an old cache/state tensor consumed by a fused op. + + This keeps cache semantics separate from ordinary data inputs, allowing IR + and runtime planning to bind KV/conv/recurrent state tensors correctly for + prefill-with-cache and decode-with-cache exports. + """ + + role: str + source: NodeRef + cache_kind: str = CacheKind.KV + tensor_role: str | None = None + layer_index: int | None = None + optional: bool = False + tensor_constraints: tuple[TensorConstraint, ...] = () + metadata: dict[str, Any] = field(default_factory=dict) + + +@dataclass(slots=True) +class CacheOutput: + """ + Describes a new or updated cache/state tensor produced by a fused op. + + IR can use this to reconnect tuple outputs or explicit cache writes after + fusion, while the generator can use cache_kind/tensor_role to emit the + right Cactus runtime cache plumbing. + """ + + role: str + node: str + output_index: int | None = None + cache_kind: str = CacheKind.KV + tensor_role: str | None = None + layer_index: int | None = None + optional: bool = False + tensor_constraints: tuple[TensorConstraint, ...] = () + metadata: dict[str, Any] = field(default_factory=dict) + + +@dataclass(slots=True) +class CacheMutation: + """ + Describes read/write semantics for a cache-aware fusion. + + CacheInput/CacheOutput say which tensors cross the fusion boundary; this + object says how the operation conceptually mutates or advances that state, + such as appending KV rows or rolling a convolution window. + """ + + name: str + cache_kind: str = CacheKind.KV + read_roles: tuple[str, ...] = () + write_roles: tuple[str, ...] = () + sequence_axis: int | None = -2 + window_attr: str | None = None + required: bool = True + metadata: dict[str, Any] = field(default_factory=dict) + + +@dataclass(slots=True) +class RepeatedSubgraph: + """ + Describes a nested pattern that can repeat inside a larger fusion. + + MoE is the main use case: one router/top-k structure fans into repeated + expert branches with the same shape. IR owns the logic for discovering and + binding each repeated branch. + """ + + name: str + graph: "FusionGraph" + min_count: int = 1 + max_count: int | None = None + anchor_node: str | None = None + metadata: dict[str, Any] = field(default_factory=dict) + + +@dataclass(slots=True) +class FusionGraph: + """ + Describes the complete synthetic DAG for one candidate fusion. + + It stores pattern nodes, required edges, exposed inputs/outputs, attrs to + capture, optional nodes, repeated subgraphs, cache boundaries, and extra + textual constraints that IR matchers must enforce before replacing the + matched real subgraph with a fused op. + """ + + name: str + root: str + nodes: dict[str, FusionNode] = field(default_factory=dict) + edges: tuple[FusionEdge, ...] = () + inputs: tuple[FusionInput, ...] = () + shared_inputs: tuple[tuple[NodeRef, NodeRef], ...] = () + outputs: tuple[FusionOutput, ...] = () + attr_captures: tuple[AttrCapture, ...] = () + optional_nodes: tuple[str, ...] = () + repeated_subgraphs: tuple[RepeatedSubgraph, ...] = () + cache_inputs: tuple[CacheInput, ...] = () + cache_outputs: tuple[CacheOutput, ...] = () + cache_mutations: tuple[CacheMutation, ...] = () + constraints: tuple[str, ...] = () + variants: tuple[str, ...] = () + metadata: dict[str, Any] = field(default_factory=dict) + allow_root_external_children: bool = True + + +@dataclass(slots=True) +class FusionDefinition: + """ + Wraps a FusionGraph with Cactus-facing metadata. + + target/cactus_op tell the generator what op this pattern lowers to, while + fusion_fields, inference modes, modalities, and metadata let model profiles + select stronger model-specific patterns before falling back to generic ones. + """ + + name: str + target: str + graph: FusionGraph + fusion_fields: tuple[str, ...] = () + supported_inference_modes: tuple[str, ...] = () + supported_modalities: tuple[str, ...] = () + cactus_op: str | None = None + metadata: dict[str, Any] = field(default_factory=dict) + + +@dataclass(slots=True) +class FusionCatalog: + """ + Stores the registry of all fusion definitions. + + This is the compact object IR can iterate over or index when running fusion + passes. The heavier lookup maps live in fusions.py, not in this data model. + """ + + fusions: tuple[FusionDefinition, ...] = () diff --git a/python/cactus/transpiler/Fusions/nodes.py b/python/cactus/transpiler/Fusions/nodes.py new file mode 100644 index 000000000..a3e8287e0 --- /dev/null +++ b/python/cactus/transpiler/Fusions/nodes.py @@ -0,0 +1,359 @@ +from . import models as M + +EXPERT_GROUP = "experts" + + +ATEN_OPS: dict[str, tuple[str, ...]] = { + "placeholder": ("placeholder",), + "getitem": ("operator.getitem", ""), + "embedding": ("aten.embedding.default", "aten.index_select.default", "aten.gather.default"), + "matmul": ("aten.mm.default", "aten.matmul.default", "aten.bmm.default", "aten.addmm.default", "aten.linear.default"), + "add": ("aten.add.Tensor", "aten.add.Scalar", "operator.add", ""), + "sub": ("aten.sub.Tensor", "aten.sub.Scalar", "operator.sub", ""), + "mul": ("aten.mul.Tensor", "aten.mul.Scalar", "operator.mul", ""), + "div": ("aten.div.Tensor", "aten.div.Scalar", "operator.truediv", ""), + "floor_divide": ("aten.floor_divide.default", "aten.div.Tensor_mode"), + "abs": ("aten.abs.default",), + "pow": ("aten.pow.Tensor_Scalar", "aten.pow.Tensor_Tensor", "aten.pow.Scalar"), + "sqrt": ("aten.sqrt.default", "aten.rsqrt.default"), + "exp": ("aten.exp.default",), + "cos": ("aten.cos.default",), + "sin": ("aten.sin.default",), + "log": ("aten.log.default",), + "mean": ("aten.mean.dim", "aten.mean.default"), + "sum": ("aten.sum.dim_IntList", "aten.sum.default"), + "var": ("aten.var.dim", "aten.var.correction"), + "max": ("aten.max.dim", "aten.amax.default"), + "min": ("aten.min.dim", "aten.amin.default"), + "softmax": ("aten.softmax.int", "aten._softmax.default"), + "topk": ("aten.topk.default",), + "view": ("aten.view.default", "aten.reshape.default", "aten.flatten.using_ints"), + "transpose": ("aten.t.default", "aten.transpose.int", "aten.permute.default"), + "slice": ("aten.slice.Tensor", "aten.select.int", "aten.narrow.default"), + "index": ("aten.select.int", "aten.index.Tensor"), + "cat": ("aten.cat.default",), + "scatter": ("aten.scatter.src", "aten.scatter.value", "aten.scatter_reduce.two"), + "neg": ("aten.neg.default", "operator.neg", ""), + "where": ("aten.where.self",), + "eq": ("aten.eq.Tensor", "aten.eq.Scalar"), + "ne": ("aten.ne.Tensor", "aten.ne.Scalar"), + "lt": ("aten.lt.Tensor", "aten.lt.Scalar"), + "gt": ("aten.gt.Tensor", "aten.gt.Scalar"), + "full": ("aten.full.default", "aten.full_like.default"), + "arange": ("aten.arange.default", "aten.arange.start", "aten.arange.start_step"), + "scalar_tensor": ("aten.scalar_tensor.default",), + "copy": ("aten.clone.default", "aten._to_copy.default", "aten.contiguous.default"), + "relu": ("aten.relu.default",), + "silu": ("aten.silu.default", "aten.sigmoid.default"), + "gelu": ("aten.gelu.default",), + "tanh": ("aten.tanh.default",), + "sigmoid": ("aten.sigmoid.default",), + "clamp": ("aten.clamp.default", "aten.clamp.Tensor"), + "conv": ("aten.convolution.default", "aten.conv1d.default", "aten.conv2d.default"), + "groupnorm": ("aten.native_group_norm.default", "aten.group_norm.default"), + "batchnorm": ("aten.native_batch_norm.default", "aten.batch_norm.default"), + "layernorm": ("aten.native_layer_norm.default", "aten.layer_norm.default"), + "scaled_dot_product_attention": ("aten.scaled_dot_product_attention.default",), + "stft": ("aten.stft.default",), + "fft": ("aten.fft_rfft.default", "aten.fft_irfft.default"), +} + + +def _attr( + name: str, + value: object = None, + *, + required: bool = True, +) -> M.AttrConstraint: + return M.AttrConstraint(name=name, value=value, required=required) + + +def _node( + name: str, + op_key: str, + *, + attrs: tuple[M.AttrConstraint, ...] = (), + optional: bool = False, + transparent: bool = False, + repeated: bool = False, + repeated_group: str | None = None, + allowed_value_kinds: tuple[str, ...] = (), + tensor_constraints: tuple[M.TensorConstraint, ...] = (), + metadata: dict | None = None, +) -> M.FusionNode: + return M.FusionNode( + name=name, + ops=ATEN_OPS[op_key], + attrs=attrs, + optional=optional, + transparent=transparent, + repeated=repeated, + repeated_group=repeated_group, + allowed_value_kinds=allowed_value_kinds, + tensor_constraints=tensor_constraints, + metadata=metadata or {}, + ) + + +def _expert_node(name: str, op_key: str) -> M.FusionNode: + return _node(name, op_key, repeated=True, repeated_group=EXPERT_GROUP) + + +# def _index_nodes_by_op(nodes: dict[str, M.FusionNode]) -> dict[str, tuple[str, ...]]: +# index: dict[str, tuple[str, ...]] = {} + +# for node_name, node in nodes.items(): +# for op in node.ops: +# index[op] = (*index.get(op, ()), node_name) + +# return index + + +NODES: dict[str, M.FusionNode] = { + "add": _node("add", "add"), + "subtract": _node("subtract", "sub"), + "multiply": _node("multiply", "mul"), + "divide": _node("divide", "div"), + "abs": _node("abs", "abs"), + "pow": _node("pow", "pow"), + "sqrt": _node("sqrt", "sqrt"), + "scalar_add": _node("scalar_add", "add"), + "scalar_subtract": _node("scalar_subtract", "sub"), + "scalar_multiply": _node("scalar_multiply", "mul"), + "scalar_divide": _node("scalar_divide", "div"), + "scalar_floor_divide": _node("scalar_floor_divide", "floor_divide"), + "scalar_exp": _node("scalar_exp", "exp"), + "scalar_sqrt": _node("scalar_sqrt", "sqrt"), + "scalar_cos": _node("scalar_cos", "cos"), + "scalar_sin": _node("scalar_sin", "sin"), + "scalar_log": _node("scalar_log", "log"), + "mean": _node("mean", "mean"), + "sum": _node("sum", "sum"), + "variance": _node("variance", "var"), + "min": _node("min", "min"), + "max": _node("max", "max"), + "softmax": _node("softmax", "softmax"), + "topk_direct": _node("topk_direct", "topk"), + "view": _node("view", "view", transparent=True), + "reshape": _node("reshape", "view", transparent=True), + "flatten": _node("flatten", "view", transparent=True), + "transpose": _node("transpose", "transpose", transparent=True), + "slice": _node("slice", "slice"), + "index": _node("index", "index"), + "cat": _node("cat", "cat"), + "gather": _node("gather", "embedding"), + "not_equal": _node("not_equal", "ne"), + "equal": _node("equal", "eq"), + "embedding": _node("embedding", "embedding"), + "linear_mm": _node("linear_mm", "matmul"), + "linear_bias_add": _node("linear_bias_add", "add"), + "rms_square": _node("rms_square", "pow", attrs=(_attr("exponent", 2, required=False),)), + "rms_mean": _node("rms_mean", "mean", attrs=(_attr("dim", -1, required=False),)), + "rms_eps_add": _node("rms_eps_add", "add"), + "rms_inv": _node("rms_inv", "sqrt"), + "rms_scale": _node("rms_scale", "mul"), + "rms_weight_mul": _node("rms_weight_mul", "mul"), + "ln_mean": _node("ln_mean", "mean", attrs=(_attr("dim", -1, required=False),)), + "ln_center": _node("ln_center", "sub"), + "ln_square": _node("ln_square", "pow", attrs=(_attr("exponent", 2, required=False),)), + "ln_var": _node("ln_var", "mean", attrs=(_attr("dim", -1, required=False),)), + "ln_eps_add": _node("ln_eps_add", "add"), + "ln_inv": _node("ln_inv", "sqrt"), + "ln_norm": _node("ln_norm", "mul"), + "ln_weight_mul": _node("ln_weight_mul", "mul"), + "ln_bias_add": _node("ln_bias_add", "add", optional=True), + "groupnorm_mean": _node("groupnorm_mean", "mean"), + "groupnorm_var": _node("groupnorm_var", "var"), + "groupnorm": _node("groupnorm", "groupnorm"), + "batchnorm": _node("batchnorm", "batchnorm"), + "layernorm_direct": _node("layernorm_direct", "layernorm"), + "batchnorm_sub_mean": _node("batchnorm_sub_mean", "sub"), + "batchnorm_div_std": _node("batchnorm_div_std", "div"), + "relu": _node("relu", "relu"), + "silu": _node("silu", "silu"), + "gelu": _node("gelu", "gelu"), + "gelu_erf": _node("gelu_erf", "gelu", metadata={"approximation": "erf"}), + "sigmoid": _node("sigmoid", "sigmoid"), + "tanh": _node("tanh", "tanh"), + "clamp": _node("clamp", "clamp"), + "glu": _node("glu", "mul"), + "gate_proj": _node("gate_proj", "matmul"), + "up_proj": _node("up_proj", "matmul"), + "mlp_activation": _node("mlp_activation", "silu"), + "mlp_product": _node("mlp_product", "mul"), + "down_proj": _node("down_proj", "matmul"), + "q_proj": _node("q_proj", "matmul"), + "k_proj": _node("k_proj", "matmul"), + "v_proj": _node("v_proj", "matmul"), + "q_view": _node("q_view", "view", transparent=True, optional=True), + "k_view": _node("k_view", "view", transparent=True, optional=True), + "v_view": _node("v_view", "view", transparent=True, optional=True), + "q_rope": _node("q_rope", "add"), + "k_rope": _node("k_rope", "add"), + "attn_qk": _node("attn_qk", "matmul"), + "attn_scale": _node("attn_scale", "mul"), + "attn_mask_add": _node("attn_mask_add", "add", optional=True), + "attn_softmax": _node("attn_softmax", "softmax"), + "attn_value": _node("attn_value", "matmul"), + "attn_sdp": _node("attn_sdp", "scaled_dot_product_attention"), + "rel_pos_bias": _node("rel_pos_bias", "matmul"), + "rope_slice_even": _node("rope_slice_even", "slice"), + "rope_slice_odd": _node("rope_slice_odd", "slice"), + "rope_neg": _node("rope_neg", "neg"), + "rope_rotate_cat": _node("rope_rotate_cat", "cat"), + "rope_cos_mul": _node("rope_cos_mul", "mul"), + "rope_sin_mul": _node("rope_sin_mul", "mul"), + "rope_add": _node("rope_add", "add"), + "conv": _node("conv", "conv"), + "conv_bias_add": _node("conv_bias_add", "add", optional=True), + "conv_cache_append": _node("conv_cache_append", "cat"), + "conv_cache_initialize": _node("conv_cache_initialize", "copy"), + "conv_cache_state": _node("conv_cache_state", "placeholder", allowed_value_kinds=(M.ValueKind.CACHE_STATE,)), + "kv_cache_append": _node("kv_cache_append", "cat"), + "kv_cache_state": _node("kv_cache_state", "placeholder", allowed_value_kinds=(M.ValueKind.CACHE_STATE,)), + "recurrent_cache_state": _node("recurrent_cache_state", "placeholder", allowed_value_kinds=(M.ValueKind.CACHE_STATE,)), + "recurrent_cache_write": _node("recurrent_cache_write", "copy"), + "lstm_gate_mm": _node("lstm_gate_mm", "matmul"), + "lstm_recurrent_mm": _node("lstm_recurrent_mm", "matmul"), + "lstm_gate_add": _node("lstm_gate_add", "add"), + "lstm_sigmoid": _node("lstm_sigmoid", "sigmoid"), + "lstm_tanh": _node("lstm_tanh", "tanh"), + "delta_q": _node("delta_q", "matmul"), + "delta_k": _node("delta_k", "matmul"), + "delta_v": _node("delta_v", "matmul"), + "delta_gate": _node("delta_gate", "matmul"), + "delta_prefill": _node("delta_prefill", "matmul"), + "router_logits": _node("router_logits", "matmul"), + "routing_probs": _node("routing_probs", "softmax"), + "topk": _node("topk", "topk"), + "expert_gate": _expert_node("expert_gate", "matmul"), + "expert_up": _expert_node("expert_up", "matmul"), + "expert_activation": _expert_node("expert_activation", "silu"), + "expert_product": _expert_node("expert_product", "mul"), + "expert_down": _expert_node("expert_down", "matmul"), + "expert_weighted": _expert_node("expert_weighted", "mul"), + "moe_combine": _node("moe_combine", "sum"), + "rfft": _node("rfft", "fft"), + "irfft": _node("irfft", "fft"), + "stft": _node("stft", "stft"), + "spectrogram": _node("spectrogram", "stft"), + "image_preprocess": _node("image_preprocess", "view"), + "bilinear_interpolation": _node("bilinear_interpolation", "view"), + "sample_topk": _node("sample_topk", "topk"), + "sample_softmax": _node("sample_softmax", "softmax"), + "scatter_topk": _node("scatter_topk", "topk"), + "scatter": _node("scatter", "scatter"), + "gaussian_topk": _node("gaussian_topk", "topk"), + "altup_predict": _node("altup_predict", "add"), + "altup_correct": _node("altup_correct", "add"), +} + + +# NODE_NAMES_BY_OP: dict[str, tuple[str, ...]] = _index_nodes_by_op(NODES) + + +NODE_GROUPS: dict[str, tuple[str, ...]] = { + "direct_tensor": ( + "add", + "subtract", + "multiply", + "divide", + "abs", + "pow", + "sqrt", + "scalar_add", + "scalar_subtract", + "scalar_multiply", + "scalar_divide", + "scalar_floor_divide", + "scalar_exp", + "scalar_sqrt", + "scalar_cos", + "scalar_sin", + "scalar_log", + "mean", + "sum", + "variance", + "min", + "max", + "softmax", + "topk_direct", + "view", + "reshape", + "flatten", + "transpose", + "slice", + "index", + "cat", + "gather", + "not_equal", + "equal", + ), + "linear": ("linear_mm", "linear_bias_add"), + "normalization": ( + "rms_square", + "rms_mean", + "rms_eps_add", + "rms_inv", + "rms_scale", + "rms_weight_mul", + "ln_mean", + "ln_center", + "ln_square", + "ln_var", + "ln_eps_add", + "ln_inv", + "ln_norm", + "ln_weight_mul", + "ln_bias_add", + ), + "activation": ("relu", "silu", "gelu", "gelu_erf", "sigmoid", "tanh", "clamp", "glu"), + "mlp": ("gate_proj", "up_proj", "mlp_activation", "mlp_product", "down_proj"), + "attention": ( + "q_proj", + "k_proj", + "v_proj", + "q_view", + "k_view", + "v_view", + "q_rope", + "k_rope", + "attn_qk", + "attn_scale", + "attn_mask_add", + "attn_softmax", + "attn_value", + "attn_sdp", + ), + "rope": ( + "rope_slice_even", + "rope_slice_odd", + "rope_neg", + "rope_rotate_cat", + "rope_cos_mul", + "rope_sin_mul", + "rope_add", + ), + "conv": ("conv", "conv_bias_add", "conv_cache_append"), + "cache": ( + "kv_cache_state", + "kv_cache_append", + "conv_cache_state", + "conv_cache_append", + "conv_cache_initialize", + "recurrent_cache_state", + "recurrent_cache_write", + ), + "moe": ( + "router_logits", + "routing_probs", + "topk", + "expert_gate", + "expert_up", + "expert_activation", + "expert_product", + "expert_down", + "expert_weighted", + "moe_combine", + ), +} From 8852af1e390ddca07ac8969bba429d53dbcc3565 Mon Sep 17 00:00:00 2001 From: Parkirat Date: Fri, 24 Jul 2026 15:48:58 -0400 Subject: [PATCH 43/84] Initial graph classes --- python/cactus/transpiler/IR/constants.py | 99 ++++++++++++ python/cactus/transpiler/IR/models.py | 177 +++++++++++++++++++++ python/cactus/transpiler/IR/simplify_ir.py | 0 python/cactus/transpiler/IR/utils.py | 0 4 files changed, 276 insertions(+) create mode 100644 python/cactus/transpiler/IR/constants.py create mode 100644 python/cactus/transpiler/IR/simplify_ir.py create mode 100644 python/cactus/transpiler/IR/utils.py diff --git a/python/cactus/transpiler/IR/constants.py b/python/cactus/transpiler/IR/constants.py new file mode 100644 index 000000000..19b8dcb88 --- /dev/null +++ b/python/cactus/transpiler/IR/constants.py @@ -0,0 +1,99 @@ +POSITIONAL_ATTRS_MAP: dict[str, dict[int, str]] = { + "operator.getitem": {1: "index"}, + "": {1: "index"}, + "aten.add.Tensor": {1: "other", 2: "alpha"}, + "aten.add.Scalar": {1: "other", 2: "alpha"}, + "aten.sub.Tensor": {1: "other", 2: "alpha"}, + "aten.sub.Scalar": {1: "other", 2: "alpha"}, + "aten.mul.Tensor": {1: "other"}, + "aten.mul.Scalar": {1: "other"}, + "aten.div.Tensor": {1: "other"}, + "aten.div.Scalar": {1: "other"}, + "aten.div.Tensor_mode": {1: "other", 2: "rounding_mode"}, + "aten.floor_divide.default": {1: "other"}, + "aten.eq.Tensor": {1: "other"}, + "aten.eq.Scalar": {1: "other"}, + "aten.ne.Tensor": {1: "other"}, + "aten.ne.Scalar": {1: "other"}, + "aten.lt.Tensor": {1: "other"}, + "aten.lt.Scalar": {1: "other"}, + "aten.gt.Tensor": {1: "other"}, + "aten.gt.Scalar": {1: "other"}, + "aten.le.Tensor": {1: "other"}, + "aten.le.Scalar": {1: "other"}, + "aten.ge.Tensor": {1: "other"}, + "aten.ge.Scalar": {1: "other"}, + "aten.bitwise_or.Tensor": {1: "other"}, + "aten.bitwise_and.Tensor": {1: "other"}, + "aten.pow.Tensor_Scalar": {1: "exponent"}, + "aten.pow.Tensor_Tensor": {1: "exponent"}, + "aten.pow.Scalar": {1: "exponent"}, + "aten.mean.dim": {1: "dim", 2: "keepdim"}, + "aten.sum.dim_IntList": {1: "dim", 2: "keepdim", 3: "dtype"}, + "aten.var.dim": {1: "dim", 2: "correction", 3: "keepdim"}, + "aten.var.correction": {1: "dim", 2: "correction", 3: "keepdim"}, + "aten.max.dim": {1: "dim", 2: "keepdim"}, + "aten.amax.default": {1: "dim", 2: "keepdim"}, + "aten.min.dim": {1: "dim", 2: "keepdim"}, + "aten.amin.default": {1: "dim", 2: "keepdim"}, + "aten.softmax.int": {1: "dim", 2: "dtype"}, + "aten._softmax.default": {1: "dim", 2: "half_to_float"}, + "aten.topk.default": {1: "k", 2: "dim", 3: "largest", 4: "sorted"}, + "aten.view.default": {1: "shape"}, + "aten.reshape.default": {1: "shape"}, + "aten.flatten.using_ints": {1: "start_dim", 2: "end_dim"}, + "aten.t.default": {}, + "aten.transpose.int": {1: "dim0", 2: "dim1"}, + "aten.permute.default": {1: "permutation"}, + "aten.slice.Tensor": {1: "dim", 2: "start", 3: "end", 4: "step"}, + "aten.select.int": {1: "dim", 2: "index"}, + "aten.narrow.default": {1: "dim", 2: "start", 3: "length"}, + "aten.index.Tensor": {1: "indices"}, + "aten.cat.default": {1: "dim"}, + "aten.embedding.default": {2: "padding_idx", 3: "scale_grad_by_freq", 4: "sparse"}, + "aten.gather.default": {1: "dim", 2: "index"}, + "aten.scatter.src": {1: "dim", 2: "index"}, + "aten.scatter.value": {1: "dim", 2: "index", 3: "value"}, + "aten.scatter_reduce.two": {1: "dim", 2: "index", 4: "reduce", 5: "include_self"}, + "aten.clamp.default": {1: "min", 2: "max"}, + "aten.clamp.Tensor": {1: "min", 2: "max"}, + "aten.gelu.default": {1: "approximate"}, + "aten.convolution.default": { + 3: "stride", + 4: "padding", + 5: "dilation", + 6: "transposed", + 7: "output_padding", + 8: "groups", + }, + "aten.conv1d.default": {3: "stride", 4: "padding", 5: "dilation", 6: "groups"}, + "aten.conv2d.default": {3: "stride", 4: "padding", 5: "dilation", 6: "groups"}, + "aten.native_group_norm.default": {3: "batch_size", 4: "channels", 5: "height_width", 6: "num_groups", 7: "eps"}, + "aten.group_norm.default": {1: "num_groups", 4: "eps"}, + "aten.native_batch_norm.default": {5: "training", 6: "momentum", 7: "eps"}, + "aten.batch_norm.default": {5: "training", 6: "momentum", 7: "eps", 8: "cudnn_enabled"}, + "aten.native_layer_norm.default": {1: "normalized_shape", 4: "eps"}, + "aten.layer_norm.default": {1: "normalized_shape", 4: "eps"}, + "aten.scaled_dot_product_attention.default": { + 3: "attn_mask", + 4: "dropout_p", + 5: "is_causal", + 6: "scale", + 7: "enable_gqa", + }, + "aten.arange.default": {0: "end"}, + "aten.arange.start": {0: "start", 1: "end"}, + "aten.arange.start_step": {0: "start", 1: "end", 2: "step"}, + "aten.full.default": {0: "size", 1: "fill_value"}, + "aten.full_like.default": {1: "fill_value"}, + "aten.stft.default": { + 1: "n_fft", + 2: "hop_length", + 3: "win_length", + 5: "normalized", + 6: "onesided", + 7: "return_complex", + }, + "aten.fft_rfft.default": {1: "n", 2: "dim", 3: "norm"}, + "aten.fft_irfft.default": {1: "n", 2: "dim", 3: "norm"}, +} \ No newline at end of file diff --git a/python/cactus/transpiler/IR/models.py b/python/cactus/transpiler/IR/models.py index e69de29bb..3bd5aa8b8 100644 --- a/python/cactus/transpiler/IR/models.py +++ b/python/cactus/transpiler/IR/models.py @@ -0,0 +1,177 @@ +from dataclasses import dataclass, field +from typing import Any + +from . import constants + +from ..Converter import models as CModels +from ..Fusions import models as FModels + + +@dataclass(slots=True) +class Node: + index: int + name: str + node_type: str + target: str + args: Any + kwargs: Any + users: tuple[str, ...] + tensor_output_meta: Any | None + module_stack: Any | None + value_kind: str + attrs: dict[str, Any] = field(default_factory=dict) + parents: tuple["Node", ...] = field(default_factory=tuple, repr=False) + children: tuple["Node", ...] = field(default_factory=tuple, repr=False) + + @classmethod + def from_layer(cls, record: CModels.LayerRecord) -> "Node": + return generate_node(record) + + @property + def is_placeholder(self) -> bool: + return self.node_type == "placeholder" + + @property + def is_operation(self) -> bool: + return self.node_type == "call_function" + + @property + def is_output(self) -> bool: + return self.node_type == "output" + + +@dataclass(slots=True) +class Graph: + source: Node | None + sources: tuple[Node, ...] + outputs: tuple[Node, ...] + nodes: tuple[Node, ...] + nodes_map: dict[str, Node] + fusions: dict[str, Any] = field(default_factory=dict) + + @classmethod + def from_map(cls, layer_map: CModels.LayerMap) -> "Graph": + return generate_graph(cls, layer_map) + + +################################################# Model Utils!!!!!!! ################################################# + + +def extract_node_refs(value: Any) -> tuple[str, ...]: + if isinstance(value, dict) and set(value.keys()) == {"node"} and isinstance(value["node"], str): + return (value["node"],) + + if isinstance(value, (list, tuple)): + refs: list[str] = [] + + for item in value: + refs.extend(extract_node_refs(item)) + + return tuple(refs) + + if isinstance(value, dict): + refs: list[str] = [] + + for item in value.values(): + refs.extend(extract_node_refs(item)) + + return tuple(refs) + + return () + + +def infer_value_kind(record: CModels.LayerRecord) -> str: + if record.node_type == "output": + return FModels.ValueKind.OUTPUT + if record.node_type != "placeholder": + return FModels.ValueKind.ACTIVATION + + target = record.target + + if target.startswith("p_"): + return FModels.ValueKind.PARAMETER + if target.startswith("b_"): + return FModels.ValueKind.BUFFER + if target.startswith("c_") or "lifted_tensor" in target or "lifted" in target: + return FModels.ValueKind.LIFTED_CONSTANT + if "past_key_values" in target or target in {"cache_position"}: + return FModels.ValueKind.CACHE_INPUT + + return FModels.ValueKind.USER_INPUT + + +def extract_attrs(record: CModels.LayerRecord) -> dict[str, Any]: + attrs: dict[str, Any] = {} + + if isinstance(record.kwargs, dict): + attrs.update(record.kwargs) + + args = record.args if isinstance(record.args, list) else [] + positional_attrs = constants.POSITIONAL_ATTRS_MAP.get(record.target, {}) + + for index, attr_name in positional_attrs.items(): + if index >= len(args): + continue + + value = args[index] + if extract_node_refs(value): + continue + + attrs.setdefault(attr_name, value) + + for index, value in enumerate(args): + if extract_node_refs(value): + continue + + attrs.setdefault(f"arg_{index}", value) + + return attrs + + +def generate_node(record: CModels.LayerRecord) -> Node: + return Node( + index=record.index, + name=record.name, + node_type=record.node_type, + target=record.target, + args=record.args, + kwargs=record.kwargs, + users=tuple(record.users), + tensor_output_meta=record.tensor_output_meta, + module_stack=record.module_stack, + value_kind=infer_value_kind(record), + attrs=extract_attrs(record), + ) + + +def generate_graph(cls: type[Graph], layer_map: CModels.LayerMap) -> Graph: + nodes = tuple(Node.from_layer(record) for record in layer_map.nodes) + nodes_map = {node.name: node for node in nodes} + parent_lists: dict[str, list[Node]] = {node.name: [] for node in nodes} + child_lists: dict[str, list[Node]] = {node.name: [] for node in nodes} + + for record in layer_map.nodes: + node = nodes_map[record.name] + parent_names = extract_node_refs((record.args, record.kwargs)) + + for parent_name in parent_names: + if parent_name not in nodes_map: + raise ValueError(f"{record.name} references missing parent node {parent_name}") + + parent = nodes_map[parent_name] + parent_lists[node.name].append(parent) + + if node not in child_lists[parent.name]: + child_lists[parent.name].append(node) + + for node in nodes: + node.parents = tuple(parent_lists[node.name]) + node.children = tuple(child_lists[node.name]) + + sources = tuple(node for node in nodes if not node.parents) + outputs = tuple(node for node in nodes if node.is_output) + + if not outputs: + outputs = tuple(node for node in nodes if not node.children) + + return cls(source=sources[0] if sources else None, sources=sources, outputs=outputs, nodes=nodes, nodes_map=nodes_map) diff --git a/python/cactus/transpiler/IR/simplify_ir.py b/python/cactus/transpiler/IR/simplify_ir.py new file mode 100644 index 000000000..e69de29bb diff --git a/python/cactus/transpiler/IR/utils.py b/python/cactus/transpiler/IR/utils.py new file mode 100644 index 000000000..e69de29bb From 04b9057f3865558c5ca531f1dd80f1c58af1d660 Mon Sep 17 00:00:00 2001 From: Parkirat Date: Sat, 25 Jul 2026 16:59:02 -0400 Subject: [PATCH 44/84] Added fusion results object --- python/cactus/transpiler/IR/models.py | 50 ++++++++++++++++++++++++++- python/cactus/transpiler/IR/utils.py | 7 ++++ 2 files changed, 56 insertions(+), 1 deletion(-) diff --git a/python/cactus/transpiler/IR/models.py b/python/cactus/transpiler/IR/models.py index 3bd5aa8b8..f0154d51e 100644 --- a/python/cactus/transpiler/IR/models.py +++ b/python/cactus/transpiler/IR/models.py @@ -47,13 +47,61 @@ class Graph: outputs: tuple[Node, ...] nodes: tuple[Node, ...] nodes_map: dict[str, Node] - fusions: dict[str, Any] = field(default_factory=dict) + fusions: list["FusionResult"] = field(default_factory=list) @classmethod def from_map(cls, layer_map: CModels.LayerMap) -> "Graph": return generate_graph(cls, layer_map) +@dataclass(slots=True) +class FusionResult: + fusion: FModels.FusionDefinition + source: Node + matched_nodes: tuple[Node, ...] = field(default_factory=tuple) + bindings: dict[str, Node] = field(default_factory=dict) + external_inputs: tuple[Node, ...] = field(default_factory=tuple) + attrs: dict[str, Any] = field(default_factory=dict) + + @classmethod + def from_match( + cls, + fusion: FModels.FusionDefinition, + source: Node, + *, + matched_nodes: tuple[Node, ...] = (), + bindings: dict[str, Node] | None = None, + external_inputs: tuple[Node, ...] = (), + attrs: dict[str, Any] | None = None, + ) -> "FusionResult": + return cls( + fusion=fusion, + source=source, + matched_nodes=matched_nodes or (source,), + bindings=bindings or {}, + external_inputs=external_inputs, + attrs=attrs or {}, + ) + + @property + def fusion_name(self) -> str: + return self.fusion.name + + @property + def target(self) -> str: + return self.fusion.target + + @property + def cactus_op(self) -> str | None: + return self.fusion.cactus_op + + @property + def consumed_node_names(self) -> frozenset[str]: + return frozenset(node.name for node in self.matched_nodes) + + + + ################################################# Model Utils!!!!!!! ################################################# diff --git a/python/cactus/transpiler/IR/utils.py b/python/cactus/transpiler/IR/utils.py index e69de29bb..128da17b4 100644 --- a/python/cactus/transpiler/IR/utils.py +++ b/python/cactus/transpiler/IR/utils.py @@ -0,0 +1,7 @@ +from . import models, constants +from ..Fusions import models as FModels +from ..Fusions import fusions as f +from ..Converter import models as CModels + + +def match_nodes(synth_node: models.Node, graph: models.Graph, ) -> bool: From 9c7ed823a78faac40fa2e7418ff8a3b89ae88a4c Mon Sep 17 00:00:00 2001 From: Parkirat Date: Sat, 25 Jul 2026 16:59:30 -0400 Subject: [PATCH 45/84] Future instructions --- python/cactus/transpiler/Converter/cache_utils.py | 1 + 1 file changed, 1 insertion(+) diff --git a/python/cactus/transpiler/Converter/cache_utils.py b/python/cactus/transpiler/Converter/cache_utils.py index af6134d55..9b4f296e2 100644 --- a/python/cactus/transpiler/Converter/cache_utils.py +++ b/python/cactus/transpiler/Converter/cache_utils.py @@ -3,6 +3,7 @@ from typing import Any import torch +#TODO: Clean out and simplify code @dataclass(slots=True) class CacheLayerSpec: From 46ecc0d80fde8f00321c2dd40e054396edb1aae5 Mon Sep 17 00:00:00 2001 From: Parkirat Date: Mon, 27 Jul 2026 10:47:26 -0400 Subject: [PATCH 46/84] Matching nodes --- python/cactus/transpiler/IR/match.py | 75 +++++++ python/cactus/transpiler/IR/match_utils.py | 238 +++++++++++++++++++++ python/cactus/transpiler/IR/utils.py | 7 - 3 files changed, 313 insertions(+), 7 deletions(-) create mode 100644 python/cactus/transpiler/IR/match.py create mode 100644 python/cactus/transpiler/IR/match_utils.py delete mode 100644 python/cactus/transpiler/IR/utils.py diff --git a/python/cactus/transpiler/IR/match.py b/python/cactus/transpiler/IR/match.py new file mode 100644 index 000000000..20210aa88 --- /dev/null +++ b/python/cactus/transpiler/IR/match.py @@ -0,0 +1,75 @@ +from typing import Any +from . import models, match_utils +from ..Fusions import models as FModels + + + + + + +def match_node_op(synth_node: FModels.FusionNode, node: models.Node) -> bool: + if not synth_node.ops: + return True + + if node.node_type in synth_node.ops: + return True + + if node.is_placeholder: + return "placeholder" in synth_node.ops + + if node.is_output: + return "output" in synth_node.ops + + return node.is_operation and node.target in synth_node.ops + + +def match_value_kind(synth_node: FModels.FusionNode, node: models.Node) -> bool: + if not synth_node.allowed_value_kinds: + return True + + if node.value_kind in synth_node.allowed_value_kinds: + return True + + return ( + FModels.ValueKind.CACHE_STATE in synth_node.allowed_value_kinds + and node.value_kind in {FModels.ValueKind.CACHE_INPUT, FModels.ValueKind.CACHE_OUTPUT} + ) + + +def match_attr_constraints(synth_node: FModels.FusionNode, graph: models.Graph, node: models.Node, bindings: dict[str, models.Node]) -> bool: + for constraint in synth_node.attrs: + if not match_utils.match_attr_constraint(constraint, graph, node, bindings): + return False + + return True + + + +def match_tensor_constraints( + synth_node: FModels.FusionNode, + graph: models.Graph, + node: models.Node, + bindings: dict[str, models.Node], +) -> bool: + for constraint in synth_node.tensor_constraints: + if not match_utils.match_tensor_constraint(constraint, graph, node, bindings): + return False + + return True + + + + + +def match_metadata_constraints(synth_node: FModels.FusionNode, node: models.Node) -> bool: + for key, expected in synth_node.metadata.items(): + actual = match_utils.get_metadata_constraint_value(node, key) + + if actual is match_utils.MISSING or not match_utils.compare_values(actual, expected): + return False + + return True + + + + diff --git a/python/cactus/transpiler/IR/match_utils.py b/python/cactus/transpiler/IR/match_utils.py new file mode 100644 index 000000000..f9d966ba0 --- /dev/null +++ b/python/cactus/transpiler/IR/match_utils.py @@ -0,0 +1,238 @@ +from typing import Any +from . import models +from ..Fusions import models as FModels + +MISSING = object() +WILDCARD_DIM_VALUES = (None, "*", "?", -1) +METADATA_ATTR_ALIASES = { + "approximation": ("approximation", "approximate"), +} + +def get_tensor_meta_value(tensor_meta: Any, key: str, default: Any = None) -> Any: + if isinstance(tensor_meta, dict): + return tensor_meta.get(key, default) + + return getattr(tensor_meta, key, default) + +def resolve_bound_node( + node_name: str, + graph: models.Graph, + bindings: dict[str, models.Node], +) -> models.Node | None: + if node_name in bindings: + return bindings[node_name] + + return graph.nodes_map.get(node_name) + + + +def get_metadata_constraint_value(node: models.Node, key: str) -> Any: + for attr_name in METADATA_ATTR_ALIASES.get(key, (key,)): + if attr_name in node.attrs: + return node.attrs[attr_name] + + if key == "approximation" and node.target == "aten.gelu.default": + return "erf" + + return MISSING + + +def match_shape(actual_shape: Any, expected_shape: tuple[Any, ...]) -> bool: + if not isinstance(actual_shape, (list, tuple)): + return False + + if len(actual_shape) != len(expected_shape): + return False + + return all(compare_dimensions(actual, expected) for actual, expected in zip(actual_shape, expected_shape)) + + +def get_dim(shape: Any, index: int) -> Any: + if not isinstance(shape, (list, tuple)): + return MISSING + + normalized_index = index if index >= 0 else len(shape) + index + + if normalized_index < 0 or normalized_index >= len(shape): + return MISSING + + return shape[normalized_index] + + +def compare_dimensions(actual: Any, expected: Any) -> bool: + if expected in WILDCARD_DIM_VALUES: + return True + + return compare_values(actual, expected) + + +def compare_dtypes(actual: Any, expected: Any) -> bool: + return normalize_dtype(actual) == normalize_dtype(expected) + + +def normalize_dtype(value: Any) -> str: + return str(value).lower().replace("torch.", "").strip() + + +def compare_values(actual: Any, expected: Any, comparator: str = "eq") -> bool: + comparator = comparator.lower() + + if comparator in {"eq", "=="}: + return values_equal(actual, expected) + + if comparator in {"ne", "!="}: + return not values_equal(actual, expected) + + if comparator == "in": + return value_in(actual, expected) + + if comparator == "not_in": + return not value_in(actual, expected) + + if comparator in {"gt", ">"}: + return actual > expected + + if comparator in {"gte", "ge", ">="}: + return actual >= expected + + if comparator in {"lt", "<"}: + return actual < expected + + if comparator in {"lte", "le", "<="}: + return actual <= expected + + if comparator == "contains": + return expected in actual + + if comparator == "truthy": + return bool(actual) + + if comparator == "falsy": + return not bool(actual) + + if comparator == "is_none": + return actual is None + + if comparator == "not_none": + return actual is not None + + raise ValueError(f"Unknown attr comparator: {comparator}") + + +def value_in(actual: Any, expected: Any) -> bool: + if isinstance(expected, (list, tuple, set, frozenset)): + return any(values_equal(actual, item) for item in expected) + + return values_equal(actual, expected) + + +def values_equal(actual: Any, expected: Any) -> bool: + actual = unwrap_singleton(actual, expected) + expected = unwrap_singleton(expected, actual) + + if isinstance(actual, dict) and isinstance(expected, dict): + if actual.keys() != expected.keys(): + return False + + return all(values_equal(actual[key], expected[key]) for key in actual) + + if isinstance(actual, (list, tuple)) and isinstance(expected, (list, tuple)): + if len(actual) != len(expected): + return False + + return all(values_equal(left, right) for left, right in zip(actual, expected)) + + if expected == "erf" and actual == "none": + return True + + return actual == expected + + +def unwrap_singleton(value: Any, other: Any) -> Any: + if isinstance(value, (list, tuple)) and len(value) == 1 and not isinstance(other, (list, tuple)): + return value[0] + + return value + +def match_attr_constraint( + constraint: FModels.AttrConstraint, + graph: models.Graph, + node: models.Node, + bindings: dict[str, models.Node], +) -> bool: + actual = node.attrs.get(constraint.name, MISSING) + + if actual is MISSING: + return not constraint.required + + expected = constraint.value + + if constraint.source_node is not None and constraint.source_attr is not None: + source_node = resolve_bound_node(constraint.source_node, graph, bindings) + + if source_node is None: + return True + + expected = source_node.attrs.get(constraint.source_attr, MISSING) + + if expected is MISSING: + return not constraint.required + + return compare_values(actual, expected, constraint.comparator) + +def match_tensor_constraint( + constraint: FModels.TensorConstraint, + graph: models.Graph, + node: models.Node, + bindings: dict[str, models.Node], +) -> bool: + tensor_meta = node.tensor_output_meta + + if tensor_meta is None: + return False + + shape = get_tensor_meta_value(tensor_meta, "shape", ()) + dtype = get_tensor_meta_value(tensor_meta, "dtype", None) + rank = len(shape) if isinstance(shape, (list, tuple)) else None + + if constraint.rank is not None and rank != constraint.rank: + return False + + if constraint.min_rank is not None and (rank is None or rank < constraint.min_rank): + return False + + if constraint.max_rank is not None and (rank is None or rank > constraint.max_rank): + return False + + if constraint.dtype is not None and not compare_dtypes(dtype, constraint.dtype): + return False + + if constraint.shape and not match_shape(shape, constraint.shape): + return False + + for dim_index, expected in constraint.dim_equals: + actual = get_dim(shape, dim_index) + + if actual is MISSING or not compare_dimensions(actual, expected): + return False + + for dim_index, other_node_name, other_dim_index in constraint.same_dim_as: + actual = get_dim(shape, dim_index) + other_node = resolve_bound_node(other_node_name, graph, bindings) + + if other_node is None: + continue + + other_shape = get_tensor_meta_value(other_node.tensor_output_meta, "shape", ()) + other_dim = get_dim(other_shape, other_dim_index) + + if actual is MISSING or other_dim is MISSING or not compare_dimensions(actual, other_dim): + return False + + for key, expected in constraint.metadata.items(): + actual = get_tensor_meta_value(tensor_meta, key, MISSING) + + if actual is MISSING or not compare_values(actual, expected): + return False + + return True diff --git a/python/cactus/transpiler/IR/utils.py b/python/cactus/transpiler/IR/utils.py deleted file mode 100644 index 128da17b4..000000000 --- a/python/cactus/transpiler/IR/utils.py +++ /dev/null @@ -1,7 +0,0 @@ -from . import models, constants -from ..Fusions import models as FModels -from ..Fusions import fusions as f -from ..Converter import models as CModels - - -def match_nodes(synth_node: models.Node, graph: models.Graph, ) -> bool: From 338b5a5483680e0cebcb78381949f64c59d62ef7 Mon Sep 17 00:00:00 2001 From: Parkirat Date: Mon, 27 Jul 2026 10:47:53 -0400 Subject: [PATCH 47/84] Converter testing --- python/cactus/transpiler/Converter/convert.py | 28 ++++++++++++++++--- python/cactus/transpiler/Converter/models.py | 10 +++++-- 2 files changed, 32 insertions(+), 6 deletions(-) diff --git a/python/cactus/transpiler/Converter/convert.py b/python/cactus/transpiler/Converter/convert.py index e2e8e1060..be1c42bf9 100644 --- a/python/cactus/transpiler/Converter/convert.py +++ b/python/cactus/transpiler/Converter/convert.py @@ -4,14 +4,34 @@ from ..ModelProfiles import models as MP_Models from ..ModelProfiles import profiles as MP_Profiles -#Runs the full converter pipeline and writes the exported LayerMap JSON. -def convert(model_id: str, input_modalities: tuple[str, ...], output_path: str, custom_profile: MP_Models.ModelProfile | None = None, inference_mode: str = "prefill_no_cache") -> models.LayerMap: + +def export_layer_map( + model_id: str, + input_modalities: tuple[str, ...], + custom_profile: MP_Models.ModelProfile | None = None, + inference_mode: str = "prefill_no_cache", +) -> models.LayerMap: profile = custom_profile or MP_Profiles.MODEL_ID_MAP[model_id] model = models.create_model(mp=profile, input_modalities=input_modalities, model_id=model_id, inference_mode=inference_mode) - layer_map = model.export(model.input) + return model.export(model.input) + + +#Runs the full converter pipeline and writes the exported LayerMap JSON. +def convert( + model_id: str, + input_modalities: tuple[str, ...], + output_path: str, + custom_profile: MP_Models.ModelProfile | None = None, + inference_mode: str = "prefill_no_cache", +) -> models.LayerMap: + layer_map = export_layer_map(model_id=model_id, input_modalities=input_modalities, custom_profile=custom_profile, inference_mode=inference_mode) Path(output_path).write_text(layer_map.model_dump_json(indent=4), encoding="utf-8") return layer_map if __name__ == "__main__": - convert(model_id="google/gemma-4-E2B", input_modalities=("text", "vision", "audio")) + convert( + model_id="google/gemma-4-E2B", + input_modalities=("text", "vision", "audio"), + output_path="jsons/output.json", + ) diff --git a/python/cactus/transpiler/Converter/models.py b/python/cactus/transpiler/Converter/models.py index a28d776cd..a930ec4db 100644 --- a/python/cactus/transpiler/Converter/models.py +++ b/python/cactus/transpiler/Converter/models.py @@ -77,7 +77,13 @@ def jsonable(x: Any) -> Any: return {"node": x.name} if isinstance(x, torch.Tensor): - return {"shape": x if isinstance(x, int) else str(x), "dtype": x.dtype} + return { + "shape": [jsonable(dim) for dim in x.shape], + "dtype": str(x.dtype), + "device": str(x.device), + "requires_grad": bool(x.requires_grad), + "stride": [jsonable(dim) for dim in x.stride()], + } if isinstance(x, torch.Size): return list(x) @@ -107,7 +113,7 @@ def jsonable(x: Any) -> Any: def extract_tensor_meta(node: torch.fx.Node) -> Any | None: - return jsonable(node.meta["val"]) if "val" not in node.meta else None + return jsonable(node.meta["val"]) if "val" in node.meta else None def extract_module_stack(node: torch.fx.Node) -> Any | None: From 1e5297b0889275361ee5a100e54aa730db4475ad Mon Sep 17 00:00:00 2001 From: Parkirat Date: Mon, 27 Jul 2026 22:04:59 -0400 Subject: [PATCH 48/84] Simplification of logic --- python/cactus/transpiler/Fusions/edges.py | 30 ++- python/cactus/transpiler/Fusions/fusions.py | 27 ++- python/cactus/transpiler/Fusions/models.py | 23 +- python/cactus/transpiler/Fusions/nodes.py | 35 ++-- python/cactus/transpiler/IR/match.py | 29 +-- python/cactus/transpiler/IR/match_utils.py | 220 ++++++-------------- python/cactus/transpiler/main.py | 18 +- 7 files changed, 150 insertions(+), 232 deletions(-) diff --git a/python/cactus/transpiler/Fusions/edges.py b/python/cactus/transpiler/Fusions/edges.py index 1a313fc49..56ac199a4 100644 --- a/python/cactus/transpiler/Fusions/edges.py +++ b/python/cactus/transpiler/Fusions/edges.py @@ -9,8 +9,6 @@ def _edge( dest_input_index: int | None = 0, *, source_output_index: int | None = None, - optional: bool = False, - transparent: bool = False, repeated: bool = False, repeated_group: str | None = None, metadata: dict | None = None, @@ -20,8 +18,6 @@ def _edge( dest=dest, dest_input_index=dest_input_index, source_output_index=source_output_index, - optional=optional, - transparent=transparent, repeated=repeated, repeated_group=repeated_group, metadata=metadata or {}, @@ -75,40 +71,40 @@ def _index_edges_by_dest(edges: dict[str, M.FusionEdge]) -> dict[str, tuple[str, "ln_center_to_norm": _edge("ln_center", "ln_norm", 0), "ln_inv_to_norm": _edge("ln_inv", "ln_norm", 1), "ln_norm_to_weight_mul": _edge("ln_norm", "ln_weight_mul", 0), - "ln_weight_mul_to_bias_add": _edge("ln_weight_mul", "ln_bias_add", 0, optional=True), + "ln_weight_mul_to_bias_add": _edge("ln_weight_mul", "ln_bias_add", 0), "gate_proj_to_activation": _edge("gate_proj", "mlp_activation", 0), "mlp_activation_to_product": _edge("mlp_activation", "mlp_product", 0), "up_proj_to_product": _edge("up_proj", "mlp_product", 1), "mlp_product_to_down_proj": _edge("mlp_product", "down_proj", 0), "up_proj_to_gelu": _edge("up_proj", "gelu", 0), "gelu_to_down_proj": _edge("gelu", "down_proj", 0), - "q_proj_to_q_view": _edge("q_proj", "q_view", 0, optional=True, transparent=True), - "k_proj_to_k_view": _edge("k_proj", "k_view", 0, optional=True, transparent=True), - "v_proj_to_v_view": _edge("v_proj", "v_view", 0, optional=True, transparent=True), - "q_view_to_q_rope": _edge("q_view", "q_rope", 0, optional=True, transparent=True), - "k_view_to_k_rope": _edge("k_view", "k_rope", 0, optional=True, transparent=True), + "q_proj_to_q_view": _edge("q_proj", "q_view", 0), + "k_proj_to_k_view": _edge("k_proj", "k_view", 0), + "v_proj_to_v_view": _edge("v_proj", "v_view", 0), + "q_view_to_q_rope": _edge("q_view", "q_rope", 0), + "k_view_to_k_rope": _edge("k_view", "k_rope", 0), "q_rope_to_qk": _edge("q_rope", "attn_qk", 0), "k_rope_to_qk": _edge("k_rope", "attn_qk", 1), "attn_qk_to_scale": _edge("attn_qk", "attn_scale", 0), - "attn_scale_to_mask_add": _edge("attn_scale", "attn_mask_add", 0, optional=True), + "attn_scale_to_mask_add": _edge("attn_scale", "attn_mask_add", 0), "attn_scale_to_softmax": _edge("attn_scale", "attn_softmax", 0), - "attn_mask_add_to_softmax": _edge("attn_mask_add", "attn_softmax", 0, optional=True), + "attn_mask_add_to_softmax": _edge("attn_mask_add", "attn_softmax", 0), "attn_softmax_to_value": _edge("attn_softmax", "attn_value", 0), - "v_view_to_value": _edge("v_view", "attn_value", 1, optional=True, transparent=True), + "v_view_to_value": _edge("v_view", "attn_value", 1), "rope_slice_odd_to_neg": _edge("rope_slice_odd", "rope_neg", 0), "rope_neg_to_rotate_cat": _edge("rope_neg", "rope_rotate_cat", 0), "rope_slice_even_to_rotate_cat": _edge("rope_slice_even", "rope_rotate_cat", 1), "rope_cos_mul_to_add": _edge("rope_cos_mul", "rope_add", 0), "rope_sin_mul_to_add": _edge("rope_sin_mul", "rope_add", 1), "rope_rotate_cat_to_sin_mul": _edge("rope_rotate_cat", "rope_sin_mul", 0), - "conv_to_bias_add": _edge("conv", "conv_bias_add", 0, optional=True), + "conv_to_bias_add": _edge("conv", "conv_bias_add", 0), "lstm_gate_mm_to_add": _edge("lstm_gate_mm", "lstm_gate_add", 0), "lstm_recurrent_mm_to_add": _edge("lstm_recurrent_mm", "lstm_gate_add", 1), "lstm_add_to_sigmoid": _edge("lstm_gate_add", "lstm_sigmoid", 0), "lstm_add_to_tanh": _edge("lstm_gate_add", "lstm_tanh", 0), - "delta_q_to_decode": _edge("delta_q", "delta_gate", 0, optional=True), - "delta_k_to_decode": _edge("delta_k", "delta_gate", 1, optional=True), - "delta_v_to_decode": _edge("delta_v", "delta_gate", 2, optional=True), + "delta_q_to_decode": _edge("delta_q", "delta_gate", 0), + "delta_k_to_decode": _edge("delta_k", "delta_gate", 1), + "delta_v_to_decode": _edge("delta_v", "delta_gate", 2), "router_logits_to_probs": _edge("router_logits", "routing_probs", 0), "routing_probs_to_topk": _edge("routing_probs", "topk", 0), "expert_gate_to_activation": _expert_edge("expert_gate", "expert_activation", 0), diff --git a/python/cactus/transpiler/Fusions/fusions.py b/python/cactus/transpiler/Fusions/fusions.py index f513eb962..8a56c923e 100644 --- a/python/cactus/transpiler/Fusions/fusions.py +++ b/python/cactus/transpiler/Fusions/fusions.py @@ -107,7 +107,6 @@ def _graph( shared_inputs: tuple[tuple[M.NodeRef, M.NodeRef], ...] = (), outputs: tuple[M.FusionOutput, ...] | None = None, attr_captures: tuple[M.AttrCapture, ...] = (), - optional_nodes: tuple[str, ...] = (), repeated_subgraphs: tuple[M.RepeatedSubgraph, ...] = (), cache_inputs: tuple[M.CacheInput, ...] = (), cache_outputs: tuple[M.CacheOutput, ...] = (), @@ -126,7 +125,6 @@ def _graph( shared_inputs=shared_inputs, outputs=outputs or (_output(root),), attr_captures=attr_captures, - optional_nodes=optional_nodes, repeated_subgraphs=repeated_subgraphs, cache_inputs=cache_inputs, cache_outputs=cache_outputs, @@ -440,15 +438,34 @@ def _index_by_root_op( constraints=("Mean dimension must be the hidden dimension.", "Final multiply must consume the RMS-normalized activation and norm weight."), ) +LAYERNORM_NO_BIAS_GRAPH = _graph( + "layernorm_no_bias", + "ln_weight_mul", + ("ln_mean", "ln_center", "ln_square", "ln_var", "ln_eps_add", "ln_inv", "ln_norm", "ln_weight_mul"), + edge_names=( + "ln_mean_to_center", + "ln_center_to_square", + "ln_square_to_var", + "ln_var_to_eps_add", + "ln_eps_add_to_inv", + "ln_center_to_norm", + "ln_inv_to_norm", + "ln_norm_to_weight_mul", + ), + inputs=(_input("x", "ln_mean", 0), _input("weight", "ln_weight_mul", 1)), + shared_inputs=(_shared_input("ln_mean", 0, "ln_center", 0),), + attr_captures=(M.AttrCapture("epsilon", "ln_eps_add", "other", default=1e-5, required=False),), + constraints=("Mean/variance dimensions must match the normalized feature axis.",), +) + LAYERNORM_GRAPH = _graph( "layernorm", "ln_bias_add", ("ln_mean", "ln_center", "ln_square", "ln_var", "ln_eps_add", "ln_inv", "ln_norm", "ln_weight_mul", "ln_bias_add"), edge_names=E.EDGE_GROUPS["layernorm"], - inputs=(_input("x", "ln_mean", 0), _input("weight", "ln_weight_mul", 1), _input("bias", "ln_bias_add", 1, optional=True)), + inputs=(_input("x", "ln_mean", 0), _input("weight", "ln_weight_mul", 1), _input("bias", "ln_bias_add", 1)), shared_inputs=(_shared_input("ln_mean", 0, "ln_center", 0),), attr_captures=(M.AttrCapture("epsilon", "ln_eps_add", "other", default=1e-5, required=False),), - optional_nodes=("ln_bias_add",), constraints=("Mean/variance dimensions must match the normalized feature axis.",), ) @@ -720,6 +737,7 @@ def _index_by_root_op( "linear": LINEAR_GRAPH, "linear_bias": LINEAR_BIAS_GRAPH, "rms_norm": RMS_NORM_GRAPH, + "layernorm_no_bias": LAYERNORM_NO_BIAS_GRAPH, "layernorm": LAYERNORM_GRAPH, "swiglu_mlp": SWIGLU_MLP_GRAPH, "gelu_mlp": GELU_MLP_GRAPH, @@ -765,6 +783,7 @@ def _index_by_root_op( "linear": _definition("linear", "matmul", LINEAR_GRAPH, fusion_fields=("generic", "linear")), "linear_bias": _definition("linear_bias", "linear", LINEAR_BIAS_GRAPH, fusion_fields=("generic", "linear")), "rms_norm": _definition("rms_norm", "rms_norm", RMS_NORM_GRAPH, fusion_fields=("generic", "rmsnorm", "gemma4_rmsnorm", "qwen2_5_rmsnorm")), + "layernorm_no_bias": _definition("layernorm_no_bias", "layernorm", LAYERNORM_NO_BIAS_GRAPH, fusion_fields=("generic", "normalization")), "layernorm": _definition("layernorm", "layernorm", LAYERNORM_GRAPH, fusion_fields=("generic", "normalization")), "swiglu_mlp": _definition("swiglu_mlp", "dense_mlp_tq_fused", SWIGLU_MLP_GRAPH, fusion_fields=("generic", "mlp", "gemma4_mlp", "qwen2_5_mlp", "lfm_mlp")), "gelu_mlp": _definition("gelu_mlp", "matmul", GELU_MLP_GRAPH, fusion_fields=("generic", "mlp")), diff --git a/python/cactus/transpiler/Fusions/models.py b/python/cactus/transpiler/Fusions/models.py index bcd3d0fe3..4eacdafed 100644 --- a/python/cactus/transpiler/Fusions/models.py +++ b/python/cactus/transpiler/Fusions/models.py @@ -100,8 +100,8 @@ class AttrConstraint: name: str value: Any = None - source_node: str | None = None - source_attr: str | None = None + source_node: str | None = None #Which node to pull attr from + source_attr: str | None = None #Which attr from external node to pull comparator: str = "eq" required: bool = True metadata: dict[str, Any] = field(default_factory=dict) @@ -145,16 +145,14 @@ class FusionNode: """ Describes one synthetic node inside a fusion pattern graph. - It stores the allowed raw op targets, optional node-level attr/tensor - constraints, and structural hints such as optional, transparent, or repeated. + It stores the allowed raw op targets, node-level attr/tensor constraints, + and structural hints for repeated pattern regions. IR matchers bind these descriptions to real exported nodes during fusion. """ name: str ops: tuple[str, ...] = () attrs: tuple[AttrConstraint, ...] = () - optional: bool = False - transparent: bool = False repeated: bool = False repeated_group: str | None = None allowed_value_kinds: tuple[str, ...] = () @@ -168,16 +166,14 @@ class FusionEdge: Describes one directed producer-to-consumer edge in a fusion pattern. source/dest identify synthetic nodes, dest_input_index preserves meaningful - input ordering from FX args, and optional/transparent/repeated flags let IR - match common graph variants without changing the core FusionGraph shape. + input ordering from FX args, and repeated flags let IR match repeated graph + regions such as MoE expert branches. """ source: str dest: str dest_input_index: int | None = None source_output_index: int | None = None - optional: bool = False - transparent: bool = False repeated: bool = False repeated_group: str | None = None metadata: dict[str, Any] = field(default_factory=dict) @@ -308,9 +304,9 @@ class FusionGraph: Describes the complete synthetic DAG for one candidate fusion. It stores pattern nodes, required edges, exposed inputs/outputs, attrs to - capture, optional nodes, repeated subgraphs, cache boundaries, and extra - textual constraints that IR matchers must enforce before replacing the - matched real subgraph with a fused op. + capture, repeated subgraphs, cache boundaries, and extra textual constraints + that IR matchers must enforce before replacing the matched real subgraph + with a fused op. """ name: str @@ -321,7 +317,6 @@ class FusionGraph: shared_inputs: tuple[tuple[NodeRef, NodeRef], ...] = () outputs: tuple[FusionOutput, ...] = () attr_captures: tuple[AttrCapture, ...] = () - optional_nodes: tuple[str, ...] = () repeated_subgraphs: tuple[RepeatedSubgraph, ...] = () cache_inputs: tuple[CacheInput, ...] = () cache_outputs: tuple[CacheOutput, ...] = () diff --git a/python/cactus/transpiler/Fusions/nodes.py b/python/cactus/transpiler/Fusions/nodes.py index a3e8287e0..4297412cc 100644 --- a/python/cactus/transpiler/Fusions/nodes.py +++ b/python/cactus/transpiler/Fusions/nodes.py @@ -1,6 +1,11 @@ from . import models as M EXPERT_GROUP = "experts" +CACHE_STATE_VALUE_KINDS = ( + M.ValueKind.CACHE_STATE, + M.ValueKind.CACHE_INPUT, + M.ValueKind.CACHE_OUTPUT, +) ATEN_OPS: dict[str, tuple[str, ...]] = { @@ -73,8 +78,6 @@ def _node( op_key: str, *, attrs: tuple[M.AttrConstraint, ...] = (), - optional: bool = False, - transparent: bool = False, repeated: bool = False, repeated_group: str | None = None, allowed_value_kinds: tuple[str, ...] = (), @@ -85,8 +88,6 @@ def _node( name=name, ops=ATEN_OPS[op_key], attrs=attrs, - optional=optional, - transparent=transparent, repeated=repeated, repeated_group=repeated_group, allowed_value_kinds=allowed_value_kinds, @@ -134,10 +135,10 @@ def _expert_node(name: str, op_key: str) -> M.FusionNode: "max": _node("max", "max"), "softmax": _node("softmax", "softmax"), "topk_direct": _node("topk_direct", "topk"), - "view": _node("view", "view", transparent=True), - "reshape": _node("reshape", "view", transparent=True), - "flatten": _node("flatten", "view", transparent=True), - "transpose": _node("transpose", "transpose", transparent=True), + "view": _node("view", "view"), + "reshape": _node("reshape", "view"), + "flatten": _node("flatten", "view"), + "transpose": _node("transpose", "transpose"), "slice": _node("slice", "slice"), "index": _node("index", "index"), "cat": _node("cat", "cat"), @@ -161,7 +162,7 @@ def _expert_node(name: str, op_key: str) -> M.FusionNode: "ln_inv": _node("ln_inv", "sqrt"), "ln_norm": _node("ln_norm", "mul"), "ln_weight_mul": _node("ln_weight_mul", "mul"), - "ln_bias_add": _node("ln_bias_add", "add", optional=True), + "ln_bias_add": _node("ln_bias_add", "add"), "groupnorm_mean": _node("groupnorm_mean", "mean"), "groupnorm_var": _node("groupnorm_var", "var"), "groupnorm": _node("groupnorm", "groupnorm"), @@ -185,14 +186,14 @@ def _expert_node(name: str, op_key: str) -> M.FusionNode: "q_proj": _node("q_proj", "matmul"), "k_proj": _node("k_proj", "matmul"), "v_proj": _node("v_proj", "matmul"), - "q_view": _node("q_view", "view", transparent=True, optional=True), - "k_view": _node("k_view", "view", transparent=True, optional=True), - "v_view": _node("v_view", "view", transparent=True, optional=True), + "q_view": _node("q_view", "view"), + "k_view": _node("k_view", "view"), + "v_view": _node("v_view", "view"), "q_rope": _node("q_rope", "add"), "k_rope": _node("k_rope", "add"), "attn_qk": _node("attn_qk", "matmul"), "attn_scale": _node("attn_scale", "mul"), - "attn_mask_add": _node("attn_mask_add", "add", optional=True), + "attn_mask_add": _node("attn_mask_add", "add"), "attn_softmax": _node("attn_softmax", "softmax"), "attn_value": _node("attn_value", "matmul"), "attn_sdp": _node("attn_sdp", "scaled_dot_product_attention"), @@ -205,13 +206,13 @@ def _expert_node(name: str, op_key: str) -> M.FusionNode: "rope_sin_mul": _node("rope_sin_mul", "mul"), "rope_add": _node("rope_add", "add"), "conv": _node("conv", "conv"), - "conv_bias_add": _node("conv_bias_add", "add", optional=True), + "conv_bias_add": _node("conv_bias_add", "add"), "conv_cache_append": _node("conv_cache_append", "cat"), "conv_cache_initialize": _node("conv_cache_initialize", "copy"), - "conv_cache_state": _node("conv_cache_state", "placeholder", allowed_value_kinds=(M.ValueKind.CACHE_STATE,)), + "conv_cache_state": _node("conv_cache_state", "placeholder", allowed_value_kinds=CACHE_STATE_VALUE_KINDS), "kv_cache_append": _node("kv_cache_append", "cat"), - "kv_cache_state": _node("kv_cache_state", "placeholder", allowed_value_kinds=(M.ValueKind.CACHE_STATE,)), - "recurrent_cache_state": _node("recurrent_cache_state", "placeholder", allowed_value_kinds=(M.ValueKind.CACHE_STATE,)), + "kv_cache_state": _node("kv_cache_state", "placeholder", allowed_value_kinds=CACHE_STATE_VALUE_KINDS), + "recurrent_cache_state": _node("recurrent_cache_state", "placeholder", allowed_value_kinds=CACHE_STATE_VALUE_KINDS), "recurrent_cache_write": _node("recurrent_cache_write", "copy"), "lstm_gate_mm": _node("lstm_gate_mm", "matmul"), "lstm_recurrent_mm": _node("lstm_recurrent_mm", "matmul"), diff --git a/python/cactus/transpiler/IR/match.py b/python/cactus/transpiler/IR/match.py index 20210aa88..ae8f2dbd6 100644 --- a/python/cactus/transpiler/IR/match.py +++ b/python/cactus/transpiler/IR/match.py @@ -3,10 +3,6 @@ from ..Fusions import models as FModels - - - - def match_node_op(synth_node: FModels.FusionNode, node: models.Node) -> bool: if not synth_node.ops: return True @@ -27,13 +23,7 @@ def match_value_kind(synth_node: FModels.FusionNode, node: models.Node) -> bool: if not synth_node.allowed_value_kinds: return True - if node.value_kind in synth_node.allowed_value_kinds: - return True - - return ( - FModels.ValueKind.CACHE_STATE in synth_node.allowed_value_kinds - and node.value_kind in {FModels.ValueKind.CACHE_INPUT, FModels.ValueKind.CACHE_OUTPUT} - ) + return node.value_kind in synth_node.allowed_value_kinds def match_attr_constraints(synth_node: FModels.FusionNode, graph: models.Graph, node: models.Node, bindings: dict[str, models.Node]) -> bool: @@ -45,12 +35,7 @@ def match_attr_constraints(synth_node: FModels.FusionNode, graph: models.Graph, -def match_tensor_constraints( - synth_node: FModels.FusionNode, - graph: models.Graph, - node: models.Node, - bindings: dict[str, models.Node], -) -> bool: +def match_tensor_constraints(synth_node: FModels.FusionNode, graph: models.Graph, node: models.Node, bindings: dict[str, models.Node]) -> bool: for constraint in synth_node.tensor_constraints: if not match_utils.match_tensor_constraint(constraint, graph, node, bindings): return False @@ -71,5 +56,13 @@ def match_metadata_constraints(synth_node: FModels.FusionNode, node: models.Node return True +NODE_MATCHERS = [ + match_node_op, + match_value_kind, + match_attr_constraints, + match_tensor_constraints, + match_metadata_constraints +] - +def match_nodes(node: models.node, synth_node: FModels.FusionNode) -> bool: + return all(MATCHER(synth_node, node) for MATCHER in NODE_MATCHERS) diff --git a/python/cactus/transpiler/IR/match_utils.py b/python/cactus/transpiler/IR/match_utils.py index f9d966ba0..69922d7ab 100644 --- a/python/cactus/transpiler/IR/match_utils.py +++ b/python/cactus/transpiler/IR/match_utils.py @@ -1,238 +1,138 @@ from typing import Any + from . import models from ..Fusions import models as FModels -MISSING = object() -WILDCARD_DIM_VALUES = (None, "*", "?", -1) -METADATA_ATTR_ALIASES = { - "approximation": ("approximation", "approximate"), -} - -def get_tensor_meta_value(tensor_meta: Any, key: str, default: Any = None) -> Any: - if isinstance(tensor_meta, dict): - return tensor_meta.get(key, default) - - return getattr(tensor_meta, key, default) - -def resolve_bound_node( - node_name: str, - graph: models.Graph, - bindings: dict[str, models.Node], -) -> models.Node | None: - if node_name in bindings: - return bindings[node_name] - return graph.nodes_map.get(node_name) +# MISSING = object() +def shape_matches(actual_shape: list[Any], expected_shape: tuple[Any, ...]) -> bool: + if len(actual_shape) != len(expected_shape): + return False + return all(expected is None or values_equal(actual, expected) for actual, expected in zip(actual_shape, expected_shape)) -def get_metadata_constraint_value(node: models.Node, key: str) -> Any: - for attr_name in METADATA_ATTR_ALIASES.get(key, (key,)): - if attr_name in node.attrs: - return node.attrs[attr_name] - if key == "approximation" and node.target == "aten.gelu.default": - return "erf" +def get_dim(shape: list[Any], index: int) -> Any: + if index < 0: + index += len(shape) - return MISSING + if index < 0 or index >= len(shape): + return None + return shape[index] -def match_shape(actual_shape: Any, expected_shape: tuple[Any, ...]) -> bool: - if not isinstance(actual_shape, (list, tuple)): - return False - if len(actual_shape) != len(expected_shape): - return False +def normalize_dtype(dtype: Any) -> str: + return str(dtype).removeprefix("torch.") - return all(compare_dimensions(actual, expected) for actual, expected in zip(actual_shape, expected_shape)) +def normalize_gelu_approximation(value: Any) -> Any: + if value == "none": + return "erf" -def get_dim(shape: Any, index: int) -> Any: - if not isinstance(shape, (list, tuple)): - return MISSING + return value - normalized_index = index if index >= 0 else len(shape) + index - if normalized_index < 0 or normalized_index >= len(shape): - return MISSING +def values_equal(actual: Any, expected: Any) -> bool: + if isinstance(actual, tuple): + actual = list(actual) - return shape[normalized_index] + if isinstance(expected, tuple): + expected = list(expected) + if isinstance(actual, list) and len(actual) == 1 and not isinstance(expected, (list, tuple)): + actual = actual[0] -def compare_dimensions(actual: Any, expected: Any) -> bool: - if expected in WILDCARD_DIM_VALUES: - return True + if isinstance(expected, list) and len(expected) == 1 and not isinstance(actual, (list, tuple)): + expected = expected[0] - return compare_values(actual, expected) + return actual == expected -def compare_dtypes(actual: Any, expected: Any) -> bool: - return normalize_dtype(actual) == normalize_dtype(expected) +def get_metadata_constraint_value(node: models.Node, key: str) -> Any: + if key == "approximation" and node.target == "aten.gelu.default": + return normalize_gelu_approximation(node.attrs.get("approximate", "erf")) -def normalize_dtype(value: Any) -> str: - return str(value).lower().replace("torch.", "").strip() + return node.attrs.get(key, None) def compare_values(actual: Any, expected: Any, comparator: str = "eq") -> bool: - comparator = comparator.lower() - - if comparator in {"eq", "=="}: + if comparator == "eq": return values_equal(actual, expected) - if comparator in {"ne", "!="}: + if comparator == "ne": return not values_equal(actual, expected) if comparator == "in": - return value_in(actual, expected) - - if comparator == "not_in": - return not value_in(actual, expected) - - if comparator in {"gt", ">"}: - return actual > expected - - if comparator in {"gte", "ge", ">="}: - return actual >= expected - - if comparator in {"lt", "<"}: - return actual < expected - - if comparator in {"lte", "le", "<="}: - return actual <= expected - - if comparator == "contains": - return expected in actual - - if comparator == "truthy": - return bool(actual) - - if comparator == "falsy": - return not bool(actual) - - if comparator == "is_none": - return actual is None - - if comparator == "not_none": - return actual is not None + return any(values_equal(actual, candidate) for candidate in expected) raise ValueError(f"Unknown attr comparator: {comparator}") -def value_in(actual: Any, expected: Any) -> bool: - if isinstance(expected, (list, tuple, set, frozenset)): - return any(values_equal(actual, item) for item in expected) - - return values_equal(actual, expected) - - -def values_equal(actual: Any, expected: Any) -> bool: - actual = unwrap_singleton(actual, expected) - expected = unwrap_singleton(expected, actual) - - if isinstance(actual, dict) and isinstance(expected, dict): - if actual.keys() != expected.keys(): - return False - - return all(values_equal(actual[key], expected[key]) for key in actual) - - if isinstance(actual, (list, tuple)) and isinstance(expected, (list, tuple)): - if len(actual) != len(expected): - return False - - return all(values_equal(left, right) for left, right in zip(actual, expected)) - - if expected == "erf" and actual == "none": - return True - - return actual == expected - - -def unwrap_singleton(value: Any, other: Any) -> Any: - if isinstance(value, (list, tuple)) and len(value) == 1 and not isinstance(other, (list, tuple)): - return value[0] - - return value - -def match_attr_constraint( - constraint: FModels.AttrConstraint, - graph: models.Graph, - node: models.Node, - bindings: dict[str, models.Node], -) -> bool: - actual = node.attrs.get(constraint.name, MISSING) - - if actual is MISSING: +def match_attr_constraint(constraint: FModels.AttrConstraint, node: models.Node, bindings: dict[str, models.Node]) -> bool: + actual = node.attrs.get(constraint.name, None) + if actual is None: return not constraint.required expected = constraint.value if constraint.source_node is not None and constraint.source_attr is not None: - source_node = resolve_bound_node(constraint.source_node, graph, bindings) + source_node = bindings.get(constraint.source_node) if source_node is None: - return True + return False - expected = source_node.attrs.get(constraint.source_attr, MISSING) + expected = source_node.attrs.get(constraint.source_attr, None) - if expected is MISSING: + if expected is None: return not constraint.required return compare_values(actual, expected, constraint.comparator) -def match_tensor_constraint( - constraint: FModels.TensorConstraint, - graph: models.Graph, - node: models.Node, - bindings: dict[str, models.Node], -) -> bool: - tensor_meta = node.tensor_output_meta - if tensor_meta is None: +def match_tensor_constraint(constraint: FModels.TensorConstraint, node: models.Node, bindings: dict[str, models.Node]) -> bool: + if node.tensor_output_meta is None: return False - shape = get_tensor_meta_value(tensor_meta, "shape", ()) - dtype = get_tensor_meta_value(tensor_meta, "dtype", None) - rank = len(shape) if isinstance(shape, (list, tuple)) else None + shape = node.tensor_output_meta.get("shape", ()) + dtype = node.tensor_output_meta.get("dtype") + rank = len(shape) if constraint.rank is not None and rank != constraint.rank: return False - if constraint.min_rank is not None and (rank is None or rank < constraint.min_rank): + if constraint.min_rank is not None and rank < constraint.min_rank: return False - if constraint.max_rank is not None and (rank is None or rank > constraint.max_rank): + if constraint.max_rank is not None and rank > constraint.max_rank: return False - if constraint.dtype is not None and not compare_dtypes(dtype, constraint.dtype): + if constraint.dtype is not None and normalize_dtype(dtype) != normalize_dtype(constraint.dtype): return False - if constraint.shape and not match_shape(shape, constraint.shape): + if constraint.shape and not shape_matches(shape, constraint.shape): return False for dim_index, expected in constraint.dim_equals: - actual = get_dim(shape, dim_index) - - if actual is MISSING or not compare_dimensions(actual, expected): + if not values_equal(get_dim(shape, dim_index), expected): return False for dim_index, other_node_name, other_dim_index in constraint.same_dim_as: - actual = get_dim(shape, dim_index) - other_node = resolve_bound_node(other_node_name, graph, bindings) - - if other_node is None: - continue + other_node = bindings.get(other_node_name) - other_shape = get_tensor_meta_value(other_node.tensor_output_meta, "shape", ()) - other_dim = get_dim(other_shape, other_dim_index) + if other_node is None or other_node.tensor_output_meta is None: + return False - if actual is MISSING or other_dim is MISSING or not compare_dimensions(actual, other_dim): + if not values_equal(get_dim(shape, dim_index), get_dim(other_node.tensor_output_meta.get("shape", ()), other_dim_index)): return False for key, expected in constraint.metadata.items(): - actual = get_tensor_meta_value(tensor_meta, key, MISSING) - - if actual is MISSING or not compare_values(actual, expected): + if not values_equal(node.tensor_output_meta.get(key, None), expected): return False return True + + + diff --git a/python/cactus/transpiler/main.py b/python/cactus/transpiler/main.py index 339f5ebae..103615b34 100644 --- a/python/cactus/transpiler/main.py +++ b/python/cactus/transpiler/main.py @@ -3,13 +3,27 @@ sys.path.append(str(Path(__file__).resolve().parents[1])) +from transpiler.Converter import constants from transpiler.Converter.convert import convert def main() -> None: + model_id = "google/gemma-4-E2B" + input_modalities = ("text", "vision", "audio") + output_dir = constants.CONVERTER_JSON_DIR + output_dir.mkdir(parents=True, exist_ok=True) + + convert( + model_id=model_id, + input_modalities=input_modalities, + output_path=str(output_dir / "output_prefill_with_cache.json"), + inference_mode=constants.PREFILL_WITH_CACHE_MODE, + ) convert( - model_id="google/gemma-4-E2B", - input_modalities=("text", "vision", "audio"), + model_id=model_id, + input_modalities=input_modalities, + output_path=str(output_dir / "output_decode_with_cache.json"), + inference_mode=constants.DECODE_WITH_CACHE_MODE, ) From 53c759e5e1771afdbe0806c16e7a00b4a92c31d2 Mon Sep 17 00:00:00 2001 From: Parkirat Date: Tue, 28 Jul 2026 10:06:04 -0400 Subject: [PATCH 49/84] Added graph matchers --- python/cactus/transpiler/IR/match.py | 228 +++++++++++++++++++-- python/cactus/transpiler/IR/match_utils.py | 184 ++++++++++++++++- 2 files changed, 396 insertions(+), 16 deletions(-) diff --git a/python/cactus/transpiler/IR/match.py b/python/cactus/transpiler/IR/match.py index ae8f2dbd6..31c3b8b86 100644 --- a/python/cactus/transpiler/IR/match.py +++ b/python/cactus/transpiler/IR/match.py @@ -1,9 +1,9 @@ -from typing import Any from . import models, match_utils from ..Fusions import models as FModels -def match_node_op(synth_node: FModels.FusionNode, node: models.Node) -> bool: +####################################### Node Matching Logic!!!!! ####################################### +def match_node_op(node: models.Node, synth_node: FModels.FusionNode, bindings: dict[str, models.Node] | None = None) -> bool: if not synth_node.ops: return True @@ -19,34 +19,30 @@ def match_node_op(synth_node: FModels.FusionNode, node: models.Node) -> bool: return node.is_operation and node.target in synth_node.ops -def match_value_kind(synth_node: FModels.FusionNode, node: models.Node) -> bool: +def match_value_kind(node: models.Node, synth_node: FModels.FusionNode, bindings: dict[str, models.Node] | None = None) -> bool: if not synth_node.allowed_value_kinds: return True return node.value_kind in synth_node.allowed_value_kinds -def match_attr_constraints(synth_node: FModels.FusionNode, graph: models.Graph, node: models.Node, bindings: dict[str, models.Node]) -> bool: +def match_attr_constraints(node: models.Node, synth_node: FModels.FusionNode, bindings: dict[str, models.Node] | None = None) -> bool: for constraint in synth_node.attrs: - if not match_utils.match_attr_constraint(constraint, graph, node, bindings): + if not match_utils.match_attr_constraint(constraint, node, bindings or {}): return False return True - -def match_tensor_constraints(synth_node: FModels.FusionNode, graph: models.Graph, node: models.Node, bindings: dict[str, models.Node]) -> bool: +def match_tensor_constraints(node: models.Node, synth_node: FModels.FusionNode, bindings: dict[str, models.Node] | None = None) -> bool: for constraint in synth_node.tensor_constraints: - if not match_utils.match_tensor_constraint(constraint, graph, node, bindings): + if not match_utils.match_tensor_constraint(constraint, node, bindings or {}): return False return True - - - -def match_metadata_constraints(synth_node: FModels.FusionNode, node: models.Node) -> bool: +def match_metadata_constraints(node: models.Node, synth_node: FModels.FusionNode, bindings: dict[str, models.Node] | None = None) -> bool: for key, expected in synth_node.metadata.items(): actual = match_utils.get_metadata_constraint_value(node, key) @@ -61,8 +57,210 @@ def match_metadata_constraints(synth_node: FModels.FusionNode, node: models.Node match_value_kind, match_attr_constraints, match_tensor_constraints, - match_metadata_constraints + match_metadata_constraints, ] -def match_nodes(node: models.node, synth_node: FModels.FusionNode) -> bool: - return all(MATCHER(synth_node, node) for MATCHER in NODE_MATCHERS) + +def match_nodes(node: models.Node, synth_node: FModels.FusionNode, bindings: dict[str, models.Node] | None = None) -> bool: + return all(matcher(node, synth_node, bindings or {}) for matcher in NODE_MATCHERS) + + +####################################### Other Matching Logic!!!!! ####################################### +def match_subgraph(node: models.Node, graph: models.Graph, subgraph: FModels.RepeatedSubgraph) -> bool: + match_count = 0 + + for candidate in graph.nodes: + if match_fusion(candidate, graph, subgraph.graph): + match_count += 1 + + if match_count < subgraph.min_count or (match_count is not None and match_count > subgraph.max_count): + return False + + return True + + +def match_repeated_subgraphs(node: models.Node, graph: models.Graph, fusion: FModels.FusionGraph) -> bool: + bindings = match_utils.bind_fusion_graph(node, fusion, match_nodes) + + if bindings is None: + return False + + for subgraph in fusion.repeated_subgraphs: + if subgraph.anchor_node is not None and subgraph.anchor_node not in bindings: + return False + + if not match_subgraph(node, graph, subgraph): + return False + + return True + + +def match_source_node(node: models.Node, graph: models.Graph, fusion: FModels.FusionGraph) -> bool: + return match_utils.bind_fusion_graph(node, fusion, match_nodes) is not None + + +def match_edges(node: models.Node, graph: models.Graph, fusion: FModels.FusionGraph) -> bool: + bindings = match_utils.bind_fusion_graph(node, fusion, match_nodes) + + if bindings is None: + return False + + for edge in fusion.edges: + if not match_utils.edge_matches(edge, bindings): + return False + + return True + + +def match_inputs(node: models.Node, graph: models.Graph, fusion: FModels.FusionGraph) -> bool: + bindings = match_utils.bind_fusion_graph(node, fusion, match_nodes) + + if bindings is None: + return False + + declared_inputs = set() + + for input_spec in fusion.inputs: + input_nodes = match_utils.get_fusion_input_nodes(input_spec, bindings) + + if not input_nodes and not input_spec.optional: + return False + + for input_node in input_nodes: + if not match_utils.match_boundary_value(input_node, input_spec.allowed_value_kinds, input_spec.tensor_constraints): + return False + + declared_inputs.add(id(input_node)) + + for cache_input in fusion.cache_inputs: + cache_node = match_utils.get_node_ref_parent(cache_input.source, bindings) + + if cache_node is not None: + declared_inputs.add(id(cache_node)) + + return match_utils.all_external_parents_declared(bindings, declared_inputs) + + +def match_shared_inputs(node: models.Node, graph: models.Graph, fusion: FModels.FusionGraph) -> bool: + bindings = match_utils.bind_fusion_graph(node, fusion, match_nodes) + + if bindings is None: + return False + + for left_ref, right_ref in fusion.shared_inputs: + left_parent = match_utils.get_node_ref_parent(left_ref, bindings) + right_parent = match_utils.get_node_ref_parent(right_ref, bindings) + + if left_parent is None or right_parent is None: + return False + + if left_parent is not right_parent: + return False + + return True + + +def match_constrains(node: models.Node, graph: models.Graph, fusion: FModels.FusionGraph) -> bool: + return True + + +def match_cache_inputs(node: models.Node, graph: models.Graph, fusion: FModels.FusionGraph) -> bool: + bindings = match_utils.bind_fusion_graph(node, fusion, match_nodes) + + if bindings is None: + return False + + for cache_input in fusion.cache_inputs: + cache_node = match_utils.get_node_ref_parent(cache_input.source, bindings) + + if cache_node is None: + if cache_input.optional: + continue + + return False + + if not match_utils.match_boundary_value(cache_node, (), cache_input.tensor_constraints): + return False + + return True + + +def match_cache_mutations(node: models.Node, graph: models.Graph, fusion: FModels.FusionGraph) -> bool: + input_roles = {cache_input.role for cache_input in fusion.cache_inputs} + output_roles = {cache_output.role for cache_output in fusion.cache_outputs} + + for mutation in fusion.cache_mutations: + if mutation.required and not set(mutation.read_roles).issubset(input_roles): + return False + + if mutation.required and not set(mutation.write_roles).issubset(output_roles): + return False + + return True + + +def match_cache_outputs(node: models.Node, graph: models.Graph, fusion: FModels.FusionGraph) -> bool: + bindings = match_utils.bind_fusion_graph(node, fusion, match_nodes) + + if bindings is None: + return False + + for cache_output in fusion.cache_outputs: + output_node = bindings.get(cache_output.node) + + if output_node is None: + if cache_output.optional: + continue + + return False + + for constraint in cache_output.tensor_constraints: + if not match_utils.match_tensor_constraint(constraint, output_node, bindings): + return False + + return True + + +def match_metadata(node: models.Node, graph: models.Graph, fusion: FModels.FusionGraph) -> bool: + return True + + +def match_output(node: models.Node, graph: models.Graph, fusion: FModels.FusionGraph) -> bool: + bindings = match_utils.bind_fusion_graph(node, fusion, match_nodes) + + if bindings is None: + return False + + for output in fusion.outputs: + output_node = bindings.get(output.node) + + if output_node is None: + return False + + for constraint in output.tensor_constraints: + if not match_utils.match_tensor_constraint(constraint, output_node, bindings): + return False + + return match_utils.external_children_are_valid(bindings, fusion) + + +FUSION_MATCHERS = [ + match_source_node, + match_edges, + match_inputs, + match_shared_inputs, + match_repeated_subgraphs, + match_constrains, + match_cache_inputs, + match_cache_mutations, + match_cache_outputs, + match_metadata, + match_output, +] + + +def match_fusion(source: models.Node, graph: models.Graph, fusion: FModels.FusionGraph) -> bool: + return all(matcher(source, graph, fusion) for matcher in FUSION_MATCHERS) + + +####################################### Top. Sort + Graph updating!!!!! ####################################### diff --git a/python/cactus/transpiler/IR/match_utils.py b/python/cactus/transpiler/IR/match_utils.py index 69922d7ab..5b921d40f 100644 --- a/python/cactus/transpiler/IR/match_utils.py +++ b/python/cactus/transpiler/IR/match_utils.py @@ -4,7 +4,9 @@ from ..Fusions import models as FModels -# MISSING = object() +MISSING = object() + +####################################### Node Matching Utils!!!!! ####################################### def shape_matches(actual_shape: list[Any], expected_shape: tuple[Any, ...]) -> bool: if len(actual_shape) != len(expected_shape): @@ -135,4 +137,184 @@ def match_tensor_constraint(constraint: FModels.TensorConstraint, node: models.N return True +####################################### Other matching utils!!!!! ####################################### + +def bind_fusion_graph( + source: models.Node, + fusion: FModels.FusionGraph, + node_matcher: Any, +) -> dict[str, models.Node] | None: + if fusion.root not in fusion.nodes: + return None + + bindings: dict[str, models.Node] = {} + + if not bind_node(fusion.root, source, fusion, bindings, node_matcher): + return None + + changed = True + + while changed: + changed = False + + for edge in fusion.edges: + if edge.dest not in bindings: + continue + + parent = get_edge_parent(edge, bindings) + + if parent is None: + return None + + if edge.source not in bindings: + if not bind_node(edge.source, parent, fusion, bindings, node_matcher): + return None + + changed = True + continue + + if bindings[edge.source] is not parent: + return None + + for synth_node_name in fusion.nodes: + if synth_node_name not in bindings: + return None + + return bindings + + +def bind_node( + synth_node_name: str, + node: models.Node, + fusion: FModels.FusionGraph, + bindings: dict[str, models.Node], + node_matcher: Any, +) -> bool: + synth_node = fusion.nodes.get(synth_node_name) + + if synth_node is None: + return False + + if synth_node_name in bindings: + return bindings[synth_node_name] is node + + if any(bound_node is node for bound_node in bindings.values()): + return False + + if not node_matcher(node, synth_node, bindings): + return False + + bindings[synth_node_name] = node + return True + + +def edge_matches(edge: FModels.FusionEdge, bindings: dict[str, models.Node]) -> bool: + source = bindings.get(edge.source) + dest = bindings.get(edge.dest) + + if source is None or dest is None: + return False + + parent = get_edge_parent(edge, bindings) + return parent is source + + +def get_edge_parent(edge: FModels.FusionEdge, bindings: dict[str, models.Node]) -> models.Node | None: + dest = bindings.get(edge.dest) + + if dest is None: + return None + + if edge.dest_input_index is None: + if edge.source not in bindings: + return None + + source = bindings[edge.source] + return source if any(parent is source for parent in dest.parents) else None + + return get_parent(dest, edge.dest_input_index) + + +def get_node_ref_parent(ref: FModels.NodeRef, bindings: dict[str, models.Node]) -> models.Node | None: + node = bindings.get(ref.node) + + if node is None or ref.parent_index is None: + return None + + return get_parent(node, ref.parent_index) + + +def get_parent(node: models.Node, parent_index: int) -> models.Node | None: + if parent_index < 0 or parent_index >= len(node.parents): + return None + + return node.parents[parent_index] + + +def get_fusion_input_nodes(input_spec: FModels.FusionInput, bindings: dict[str, models.Node]) -> tuple[models.Node, ...]: + source_node = bindings.get(input_spec.source.node) + + if source_node is None or input_spec.source.parent_index is None: + return () + + if not input_spec.variadic: + parent = get_parent(source_node, input_spec.source.parent_index) + return (parent,) if parent is not None else () + + end_parent_index = input_spec.end_parent_index + + if end_parent_index is None: + end_parent_index = len(source_node.parents) + + parents = source_node.parents[input_spec.source.parent_index:end_parent_index] + + if len(parents) < input_spec.min_count: + return () + + if input_spec.max_count is not None and len(parents) > input_spec.max_count: + return () + + return parents + + +def match_boundary_value( + node: models.Node, + allowed_value_kinds: tuple[str, ...], + tensor_constraints: tuple[FModels.TensorConstraint, ...], +) -> bool: + if allowed_value_kinds and node.value_kind not in allowed_value_kinds: + return False + + return all(match_tensor_constraint(constraint, node, {}) for constraint in tensor_constraints) + + +def all_external_parents_declared(bindings: dict[str, models.Node], declared_input_ids: set[int]) -> bool: + internal_ids = {id(node) for node in bindings.values()} + + for node in bindings.values(): + for parent in node.parents: + if id(parent) in internal_ids: + continue + + if id(parent) not in declared_input_ids: + return False + + return True + + +def external_children_are_valid(bindings: dict[str, models.Node], fusion: FModels.FusionGraph) -> bool: + internal_ids = {id(node) for node in bindings.values()} + exposed_output_names = {output.node for output in fusion.outputs} + if not fusion.allow_root_external_children: + exposed_output_names.discard(fusion.root) + + for synth_node_name, node in bindings.items(): + for child in node.children: + if id(child) in internal_ids: + continue + + if synth_node_name not in exposed_output_names: + return False + + return True From 344edc9a36d2a99f16d5ce36e138ae6fbdb75f84 Mon Sep 17 00:00:00 2001 From: Parkirat Date: Tue, 28 Jul 2026 14:19:06 -0400 Subject: [PATCH 50/84] Simplifications to fusion matcher functions --- python/cactus/transpiler/IR/match.py | 74 ++++++---------------- python/cactus/transpiler/IR/match_utils.py | 21 +----- 2 files changed, 24 insertions(+), 71 deletions(-) diff --git a/python/cactus/transpiler/IR/match.py b/python/cactus/transpiler/IR/match.py index 31c3b8b86..ce0953bf1 100644 --- a/python/cactus/transpiler/IR/match.py +++ b/python/cactus/transpiler/IR/match.py @@ -1,6 +1,7 @@ from . import models, match_utils from ..Fusions import models as FModels +ANY = () ####################################### Node Matching Logic!!!!! ####################################### def match_node_op(node: models.Node, synth_node: FModels.FusionNode, bindings: dict[str, models.Node] | None = None) -> bool: @@ -70,21 +71,18 @@ def match_subgraph(node: models.Node, graph: models.Graph, subgraph: FModels.Rep match_count = 0 for candidate in graph.nodes: + if subgraph.max_count is not None and match_count > subgraph.max_count: + return False if match_fusion(candidate, graph, subgraph.graph): match_count += 1 - if match_count < subgraph.min_count or (match_count is not None and match_count > subgraph.max_count): + if match_count < subgraph.min_count: return False return True -def match_repeated_subgraphs(node: models.Node, graph: models.Graph, fusion: FModels.FusionGraph) -> bool: - bindings = match_utils.bind_fusion_graph(node, fusion, match_nodes) - - if bindings is None: - return False - +def match_repeated_subgraphs(node: models.Node, graph: models.Graph, fusion: FModels.FusionGraph, bindings: dict[str, models.Node]) -> bool: for subgraph in fusion.repeated_subgraphs: if subgraph.anchor_node is not None and subgraph.anchor_node not in bindings: return False @@ -95,16 +93,7 @@ def match_repeated_subgraphs(node: models.Node, graph: models.Graph, fusion: FMo return True -def match_source_node(node: models.Node, graph: models.Graph, fusion: FModels.FusionGraph) -> bool: - return match_utils.bind_fusion_graph(node, fusion, match_nodes) is not None - - -def match_edges(node: models.Node, graph: models.Graph, fusion: FModels.FusionGraph) -> bool: - bindings = match_utils.bind_fusion_graph(node, fusion, match_nodes) - - if bindings is None: - return False - +def match_edges(node: models.Node, graph: models.Graph, fusion: FModels.FusionGraph, bindings: dict[str, models.Node]) -> bool: for edge in fusion.edges: if not match_utils.edge_matches(edge, bindings): return False @@ -112,12 +101,7 @@ def match_edges(node: models.Node, graph: models.Graph, fusion: FModels.FusionGr return True -def match_inputs(node: models.Node, graph: models.Graph, fusion: FModels.FusionGraph) -> bool: - bindings = match_utils.bind_fusion_graph(node, fusion, match_nodes) - - if bindings is None: - return False - +def match_inputs(node: models.Node, graph: models.Graph, fusion: FModels.FusionGraph, bindings: dict[str, models.Node]) -> bool: declared_inputs = set() for input_spec in fusion.inputs: @@ -141,12 +125,7 @@ def match_inputs(node: models.Node, graph: models.Graph, fusion: FModels.FusionG return match_utils.all_external_parents_declared(bindings, declared_inputs) -def match_shared_inputs(node: models.Node, graph: models.Graph, fusion: FModels.FusionGraph) -> bool: - bindings = match_utils.bind_fusion_graph(node, fusion, match_nodes) - - if bindings is None: - return False - +def match_shared_inputs(node: models.Node, graph: models.Graph, fusion: FModels.FusionGraph, bindings: dict[str, models.Node]) -> bool: for left_ref, right_ref in fusion.shared_inputs: left_parent = match_utils.get_node_ref_parent(left_ref, bindings) right_parent = match_utils.get_node_ref_parent(right_ref, bindings) @@ -160,16 +139,11 @@ def match_shared_inputs(node: models.Node, graph: models.Graph, fusion: FModels. return True -def match_constrains(node: models.Node, graph: models.Graph, fusion: FModels.FusionGraph) -> bool: +def match_constrains(node: models.Node, graph: models.Graph, fusion: FModels.FusionGraph, bindings: dict[str, models.Node]) -> bool: return True -def match_cache_inputs(node: models.Node, graph: models.Graph, fusion: FModels.FusionGraph) -> bool: - bindings = match_utils.bind_fusion_graph(node, fusion, match_nodes) - - if bindings is None: - return False - +def match_cache_inputs(node: models.Node, graph: models.Graph, fusion: FModels.FusionGraph, bindings: dict[str, models.Node]) -> bool: for cache_input in fusion.cache_inputs: cache_node = match_utils.get_node_ref_parent(cache_input.source, bindings) @@ -179,13 +153,13 @@ def match_cache_inputs(node: models.Node, graph: models.Graph, fusion: FModels.F return False - if not match_utils.match_boundary_value(cache_node, (), cache_input.tensor_constraints): + if not match_utils.match_boundary_value(cache_node, ANY, cache_input.tensor_constraints): return False return True -def match_cache_mutations(node: models.Node, graph: models.Graph, fusion: FModels.FusionGraph) -> bool: +def match_cache_mutations(node: models.Node, graph: models.Graph, fusion: FModels.FusionGraph, bindings: dict[str, models.Node]) -> bool: input_roles = {cache_input.role for cache_input in fusion.cache_inputs} output_roles = {cache_output.role for cache_output in fusion.cache_outputs} @@ -199,12 +173,7 @@ def match_cache_mutations(node: models.Node, graph: models.Graph, fusion: FModel return True -def match_cache_outputs(node: models.Node, graph: models.Graph, fusion: FModels.FusionGraph) -> bool: - bindings = match_utils.bind_fusion_graph(node, fusion, match_nodes) - - if bindings is None: - return False - +def match_cache_outputs(node: models.Node, graph: models.Graph, fusion: FModels.FusionGraph, bindings: dict[str, models.Node]) -> bool: for cache_output in fusion.cache_outputs: output_node = bindings.get(cache_output.node) @@ -221,16 +190,11 @@ def match_cache_outputs(node: models.Node, graph: models.Graph, fusion: FModels. return True -def match_metadata(node: models.Node, graph: models.Graph, fusion: FModels.FusionGraph) -> bool: +def match_metadata(node: models.Node, graph: models.Graph, fusion: FModels.FusionGraph, bindings: dict[str, models.Node]) -> bool: return True -def match_output(node: models.Node, graph: models.Graph, fusion: FModels.FusionGraph) -> bool: - bindings = match_utils.bind_fusion_graph(node, fusion, match_nodes) - - if bindings is None: - return False - +def match_output(node: models.Node, graph: models.Graph, fusion: FModels.FusionGraph, bindings: dict[str, models.Node]) -> bool: for output in fusion.outputs: output_node = bindings.get(output.node) @@ -245,7 +209,6 @@ def match_output(node: models.Node, graph: models.Graph, fusion: FModels.FusionG FUSION_MATCHERS = [ - match_source_node, match_edges, match_inputs, match_shared_inputs, @@ -260,7 +223,12 @@ def match_output(node: models.Node, graph: models.Graph, fusion: FModels.FusionG def match_fusion(source: models.Node, graph: models.Graph, fusion: FModels.FusionGraph) -> bool: - return all(matcher(source, graph, fusion) for matcher in FUSION_MATCHERS) + bindings = match_utils.bind_fusion_graph(source, fusion, match_nodes) + + if bindings is None: + return False + + return all(matcher(source, graph, fusion, bindings) for matcher in FUSION_MATCHERS) ####################################### Top. Sort + Graph updating!!!!! ####################################### diff --git a/python/cactus/transpiler/IR/match_utils.py b/python/cactus/transpiler/IR/match_utils.py index 5b921d40f..23ebbed45 100644 --- a/python/cactus/transpiler/IR/match_utils.py +++ b/python/cactus/transpiler/IR/match_utils.py @@ -139,11 +139,7 @@ def match_tensor_constraint(constraint: FModels.TensorConstraint, node: models.N ####################################### Other matching utils!!!!! ####################################### -def bind_fusion_graph( - source: models.Node, - fusion: FModels.FusionGraph, - node_matcher: Any, -) -> dict[str, models.Node] | None: +def bind_fusion_graph(source: models.Node, fusion: FModels.FusionGraph, node_matcher: Any) -> dict[str, models.Node] | None: if fusion.root not in fusion.nodes: return None @@ -183,13 +179,7 @@ def bind_fusion_graph( return bindings -def bind_node( - synth_node_name: str, - node: models.Node, - fusion: FModels.FusionGraph, - bindings: dict[str, models.Node], - node_matcher: Any, -) -> bool: +def bind_node(synth_node_name: str, node: models.Node, fusion: FModels.FusionGraph, bindings: dict[str, models.Node], node_matcher: Any) -> bool: synth_node = fusion.nodes.get(synth_node_name) if synth_node is None: @@ -237,7 +227,6 @@ def get_edge_parent(edge: FModels.FusionEdge, bindings: dict[str, models.Node]) def get_node_ref_parent(ref: FModels.NodeRef, bindings: dict[str, models.Node]) -> models.Node | None: node = bindings.get(ref.node) - if node is None or ref.parent_index is None: return None @@ -277,11 +266,7 @@ def get_fusion_input_nodes(input_spec: FModels.FusionInput, bindings: dict[str, return parents -def match_boundary_value( - node: models.Node, - allowed_value_kinds: tuple[str, ...], - tensor_constraints: tuple[FModels.TensorConstraint, ...], -) -> bool: +def match_boundary_value(node: models.Node, allowed_value_kinds: tuple[str, ...], tensor_constraints: tuple[FModels.TensorConstraint, ...]) -> bool: if allowed_value_kinds and node.value_kind not in allowed_value_kinds: return False From aa231af4d7309317d21d1fcd451dcfa06ec57b76 Mon Sep 17 00:00:00 2001 From: Parkirat Date: Tue, 28 Jul 2026 14:30:31 -0400 Subject: [PATCH 51/84] Added fusion definition matchers --- python/cactus/transpiler/IR/match.py | 63 +++++++++++-- python/cactus/transpiler/IR/match_utils.py | 103 +++++++++++++++++++-- 2 files changed, 154 insertions(+), 12 deletions(-) diff --git a/python/cactus/transpiler/IR/match.py b/python/cactus/transpiler/IR/match.py index ce0953bf1..549fb2b4a 100644 --- a/python/cactus/transpiler/IR/match.py +++ b/python/cactus/transpiler/IR/match.py @@ -66,16 +66,17 @@ def match_nodes(node: models.Node, synth_node: FModels.FusionNode, bindings: dic return all(matcher(node, synth_node, bindings or {}) for matcher in NODE_MATCHERS) -####################################### Other Matching Logic!!!!! ####################################### +####################################### Graph Matching Logic!!!!! ####################################### def match_subgraph(node: models.Node, graph: models.Graph, subgraph: FModels.RepeatedSubgraph) -> bool: match_count = 0 for candidate in graph.nodes: - if subgraph.max_count is not None and match_count > subgraph.max_count: - return False if match_fusion(candidate, graph, subgraph.graph): match_count += 1 + if subgraph.max_count is not None and match_count > subgraph.max_count: + return False + if match_count < subgraph.min_count: return False @@ -139,7 +140,7 @@ def match_shared_inputs(node: models.Node, graph: models.Graph, fusion: FModels. return True -def match_constrains(node: models.Node, graph: models.Graph, fusion: FModels.FusionGraph, bindings: dict[str, models.Node]) -> bool: +def match_constraints(node: models.Node, graph: models.Graph, fusion: FModels.FusionGraph, bindings: dict[str, models.Node]) -> bool: return True @@ -213,7 +214,7 @@ def match_output(node: models.Node, graph: models.Graph, fusion: FModels.FusionG match_inputs, match_shared_inputs, match_repeated_subgraphs, - match_constrains, + match_constraints, match_cache_inputs, match_cache_mutations, match_cache_outputs, @@ -228,7 +229,57 @@ def match_fusion(source: models.Node, graph: models.Graph, fusion: FModels.Fusio if bindings is None: return False + return match_fusion_bindings(source, graph, fusion, bindings) + + +def match_fusion_bindings(source: models.Node, graph: models.Graph, fusion: FModels.FusionGraph, bindings: dict[str, models.Node]) -> bool: return all(matcher(source, graph, fusion, bindings) for matcher in FUSION_MATCHERS) -####################################### Top. Sort + Graph updating!!!!! ####################################### +####################################### Fusion Definition Matching Logic!!!!! ####################################### + +def match_definition_graph(source: models.Node, graph: models.Graph, fusion: FModels.FusionDefinition, bindings: dict[str, models.Node], inference_mode: str | None = None, input_modalities: tuple[str, ...] = ANY, fusion_fields: tuple[str, ...] = ANY) -> bool: + return match_fusion_bindings(source, graph, fusion.graph, bindings) + + +def match_definition_metadata(source: models.Node, graph: models.Graph, fusion: FModels.FusionDefinition, bindings: dict[str, models.Node], inference_mode: str | None = None, input_modalities: tuple[str, ...] = ANY, fusion_fields: tuple[str, ...] = ANY) -> bool: + return match_utils.match_definition_metadata(fusion, bindings) + + +def match_definition_inference_mode(source: models.Node, graph: models.Graph, fusion: FModels.FusionDefinition, bindings: dict[str, models.Node], inference_mode: str | None = None, input_modalities: tuple[str, ...] = ANY, fusion_fields: tuple[str, ...] = ANY) -> bool: + if not fusion.supported_inference_modes or inference_mode is None: + return True + + return inference_mode in fusion.supported_inference_modes + + +def match_definition_modalities(source: models.Node, graph: models.Graph, fusion: FModels.FusionDefinition, bindings: dict[str, models.Node], inference_mode: str | None = None, input_modalities: tuple[str, ...] = ANY, fusion_fields: tuple[str, ...] = ANY) -> bool: + if not fusion.supported_modalities or not input_modalities: + return True + + return set(fusion.supported_modalities).issubset(input_modalities) + + +def match_definition_fusion_fields(source: models.Node, graph: models.Graph, fusion: FModels.FusionDefinition, bindings: dict[str, models.Node], inference_mode: str | None = None, input_modalities: tuple[str, ...] = ANY, fusion_fields: tuple[str, ...] = ANY) -> bool: + if not fusion.fusion_fields or not fusion_fields: + return True + + return not set(fusion.fusion_fields).isdisjoint(fusion_fields) + + +FUSION_DEFINITION_MATCHERS = [ + match_definition_graph, + match_definition_metadata, + match_definition_inference_mode, + match_definition_modalities, + match_definition_fusion_fields, +] + + +def match_fusion_definition(source: models.Node, graph: models.Graph, fusion: FModels.FusionDefinition, inference_mode: str | None = None, input_modalities: tuple[str, ...] = ANY, fusion_fields: tuple[str, ...] = ANY) -> bool: + bindings = match_utils.bind_fusion_graph(source, fusion.graph, match_nodes) + + if bindings is None: + return False + + return all(matcher(source, graph, fusion, bindings, inference_mode, input_modalities, fusion_fields) for matcher in FUSION_DEFINITION_MATCHERS) diff --git a/python/cactus/transpiler/IR/match_utils.py b/python/cactus/transpiler/IR/match_utils.py index 23ebbed45..9e13641bf 100644 --- a/python/cactus/transpiler/IR/match_utils.py +++ b/python/cactus/transpiler/IR/match_utils.py @@ -57,7 +57,7 @@ def get_metadata_constraint_value(node: models.Node, key: str) -> Any: if key == "approximation" and node.target == "aten.gelu.default": return normalize_gelu_approximation(node.attrs.get("approximate", "erf")) - return node.attrs.get(key, None) + return node.attrs.get(key, MISSING) def compare_values(actual: Any, expected: Any, comparator: str = "eq") -> bool: @@ -74,8 +74,8 @@ def compare_values(actual: Any, expected: Any, comparator: str = "eq") -> bool: def match_attr_constraint(constraint: FModels.AttrConstraint, node: models.Node, bindings: dict[str, models.Node]) -> bool: - actual = node.attrs.get(constraint.name, None) - if actual is None: + actual = node.attrs.get(constraint.name, MISSING) + if actual is MISSING: return not constraint.required expected = constraint.value @@ -86,9 +86,9 @@ def match_attr_constraint(constraint: FModels.AttrConstraint, node: models.Node, if source_node is None: return False - expected = source_node.attrs.get(constraint.source_attr, None) + expected = source_node.attrs.get(constraint.source_attr, MISSING) - if expected is None: + if expected is MISSING: return not constraint.required return compare_values(actual, expected, constraint.comparator) @@ -131,7 +131,7 @@ def match_tensor_constraint(constraint: FModels.TensorConstraint, node: models.N return False for key, expected in constraint.metadata.items(): - if not values_equal(node.tensor_output_meta.get(key, None), expected): + if not values_equal(node.tensor_output_meta.get(key, MISSING), expected): return False return True @@ -303,3 +303,94 @@ def external_children_are_valid(bindings: dict[str, models.Node], fusion: FModel return False return True + + +####################################### Fusion Definition Matching utils!!!!! ####################################### + +def match_definition_metadata(fusion: FModels.FusionDefinition, bindings: dict[str, models.Node]) -> bool: + required_attrs = fusion.metadata.get("required_attrs", {}) + + if not required_attrs: + return True + + attrs = collect_definition_attrs(fusion.graph, bindings) + + for name, expected in required_attrs.items(): + actual = attrs.get(name, MISSING) + + if actual is MISSING or not values_equal(actual, expected): + return False + + return True + + +def collect_definition_attrs(fusion: FModels.FusionGraph, bindings: dict[str, models.Node]) -> dict[str, Any]: + attrs: dict[str, Any] = {} + + for capture in fusion.attr_captures: + value = get_attr_capture_value(capture, bindings) + + if value is not MISSING: + attrs[capture.name] = value + + attrs.update(infer_definition_attrs(fusion, bindings)) + return attrs + + +def get_attr_capture_value(capture: FModels.AttrCapture, bindings: dict[str, models.Node]) -> Any: + if capture.source_node is None or capture.source_attr is None: + return capture.default + + source_node = bindings.get(capture.source_node) + + if source_node is None: + return MISSING if capture.required else capture.default + + return source_node.attrs.get(capture.source_attr, MISSING if capture.required else capture.default) + + +def infer_definition_attrs(fusion: FModels.FusionGraph, bindings: dict[str, models.Node]) -> dict[str, Any]: + attrs: dict[str, Any] = {} + + conv_node = bindings.get("conv") + + if conv_node is not None: + attrs.update(infer_conv_attrs(fusion, bindings, conv_node)) + + return attrs + + +def infer_conv_attrs(fusion: FModels.FusionGraph, bindings: dict[str, models.Node], conv_node: models.Node) -> dict[str, Any]: + attrs: dict[str, Any] = {} + weight_node = get_first_input_by_role(fusion, bindings, "weight") + weight_shape = get_tensor_shape(weight_node) + groups = conv_node.attrs.get("groups", 1) + + if len(weight_shape) >= 3: + kernel_shape = weight_shape[2:] + attrs["kernel_size"] = kernel_shape[0] if all(values_equal(dim, kernel_shape[0]) for dim in kernel_shape) else kernel_shape + + if len(weight_shape) >= 2: + attrs["depthwise"] = not values_equal(groups, 1) and values_equal(weight_shape[1], 1) + + return attrs + + +def get_first_input_by_role(fusion: FModels.FusionGraph, bindings: dict[str, models.Node], role: str) -> models.Node | None: + for input_spec in fusion.inputs: + if input_spec.role != role: + continue + + input_nodes = get_fusion_input_nodes(input_spec, bindings) + + if input_nodes: + return input_nodes[0] + + return None + + +def get_tensor_shape(node: models.Node | None) -> list[Any]: + if node is None or node.tensor_output_meta is None: + return [] + + return node.tensor_output_meta.get("shape", []) From 841bdaeba6c163207b5181396c68c2f4276f532c Mon Sep 17 00:00:00 2001 From: Parkirat Date: Tue, 28 Jul 2026 14:56:46 -0400 Subject: [PATCH 52/84] Added extra matchers for extraneous constraints for fusions --- python/cactus/transpiler/Fusions/fusions.py | 50 +++-- python/cactus/transpiler/Fusions/models.py | 12 +- python/cactus/transpiler/IR/extra_matchers.py | 205 ++++++++++++++++++ python/cactus/transpiler/IR/match.py | 4 +- 4 files changed, 249 insertions(+), 22 deletions(-) create mode 100644 python/cactus/transpiler/IR/extra_matchers.py diff --git a/python/cactus/transpiler/Fusions/fusions.py b/python/cactus/transpiler/Fusions/fusions.py index 8a56c923e..478091042 100644 --- a/python/cactus/transpiler/Fusions/fusions.py +++ b/python/cactus/transpiler/Fusions/fusions.py @@ -97,6 +97,10 @@ def _required_attrs(**attrs) -> dict: return {"required_attrs": attrs} +def _note(*items: str) -> dict[str, M.ConstraintSpec]: + return {"note": {"items": items}} + + def _graph( name: str, root: str, @@ -111,7 +115,7 @@ def _graph( cache_inputs: tuple[M.CacheInput, ...] = (), cache_outputs: tuple[M.CacheOutput, ...] = (), cache_mutations: tuple[M.CacheMutation, ...] = (), - constraints: tuple[str, ...] = (), + constraints: dict[str, M.ConstraintValue] | None = None, variants: tuple[str, ...] = (), metadata: dict | None = None, allow_root_external_children: bool = True, @@ -129,7 +133,7 @@ def _graph( cache_inputs=cache_inputs, cache_outputs=cache_outputs, cache_mutations=cache_mutations, - constraints=constraints, + constraints=constraints or {}, variants=variants, metadata=metadata or {}, allow_root_external_children=allow_root_external_children, @@ -414,7 +418,7 @@ def _index_by_root_op( ("linear_mm",), inputs=(_input("x", "linear_mm", 0), _input("weight", "linear_mm", 1)), attr_captures=(M.AttrCapture("pretransposed_rhs", default=False, required=False),), - constraints=("Weight layout/pretranspose must be verified against exported attrs or weight metadata.",), + constraints=_note("Weight layout/pretranspose must be verified against exported attrs or weight metadata."), ) LINEAR_BIAS_GRAPH = _graph( @@ -424,7 +428,9 @@ def _index_by_root_op( edge_names=("linear_mm_to_bias_add",), inputs=(_input("x", "linear_mm", 0), _input("weight", "linear_mm", 1), _input("bias", "linear_bias_add", 1)), attr_captures=(M.AttrCapture("pretransposed_rhs", default=False, required=False),), - constraints=("Bias add must consume only the linear output plus a parameter/buffer bias.",), + constraints={ + "input_value_kind": {"role": "bias", "allowed_value_kinds": (M.ValueKind.PARAMETER, M.ValueKind.BUFFER)}, + }, ) RMS_NORM_GRAPH = _graph( @@ -435,7 +441,9 @@ def _index_by_root_op( inputs=(_input("x", "rms_square", 0), _input("weight", "rms_weight_mul", 1)), shared_inputs=(_shared_input("rms_square", 0, "rms_scale", 0),), attr_captures=(M.AttrCapture("epsilon", "rms_eps_add", "other", default=1e-6, required=False),), - constraints=("Mean dimension must be the hidden dimension.", "Final multiply must consume the RMS-normalized activation and norm weight."), + constraints={ + "node_attr_equals": {"node": "rms_mean", "attr": "dim", "value": -1}, + }, ) LAYERNORM_NO_BIAS_GRAPH = _graph( @@ -455,7 +463,9 @@ def _index_by_root_op( inputs=(_input("x", "ln_mean", 0), _input("weight", "ln_weight_mul", 1)), shared_inputs=(_shared_input("ln_mean", 0, "ln_center", 0),), attr_captures=(M.AttrCapture("epsilon", "ln_eps_add", "other", default=1e-5, required=False),), - constraints=("Mean/variance dimensions must match the normalized feature axis.",), + constraints={ + "node_attrs_equal": {"left_node": "ln_mean", "left_attr": "dim", "right_node": "ln_var", "right_attr": "dim"}, + }, ) LAYERNORM_GRAPH = _graph( @@ -466,7 +476,10 @@ def _index_by_root_op( inputs=(_input("x", "ln_mean", 0), _input("weight", "ln_weight_mul", 1), _input("bias", "ln_bias_add", 1)), shared_inputs=(_shared_input("ln_mean", 0, "ln_center", 0),), attr_captures=(M.AttrCapture("epsilon", "ln_eps_add", "other", default=1e-5, required=False),), - constraints=("Mean/variance dimensions must match the normalized feature axis.",), + constraints={ + "node_attrs_equal": {"left_node": "ln_mean", "left_attr": "dim", "right_node": "ln_var", "right_attr": "dim"}, + "input_value_kind": {"role": "bias", "allowed_value_kinds": (M.ValueKind.PARAMETER, M.ValueKind.BUFFER)}, + }, ) SWIGLU_MLP_GRAPH = _graph( @@ -477,7 +490,6 @@ def _index_by_root_op( inputs=(_input("hidden", "gate_proj", 0), _input("gate_weight", "gate_proj", 1), _input("up_weight", "up_proj", 1), _input("down_weight", "down_proj", 1)), shared_inputs=(_shared_input("gate_proj", 0, "up_proj", 0),), attr_captures=(M.AttrCapture("product_scale", default=1.0, required=False),), - constraints=("gate_proj and up_proj must consume the same hidden input.", "down_proj must consume silu(gate_proj(hidden)) * up_proj(hidden)."), ) GELU_MLP_GRAPH = _graph( @@ -486,7 +498,6 @@ def _index_by_root_op( ("up_proj", "gelu", "down_proj"), inputs=(_input("hidden", "up_proj", 0), _input("up_weight", "up_proj", 1), _input("down_weight", "down_proj", 1)), edge_names=E.EDGE_GROUPS["gelu_mlp"], - constraints=("This pattern uses gelu as the activation before the down projection; IR should bind the activation output to down_proj input 0.",), ) ATTENTION_DIRECT_GRAPH = _graph( @@ -504,7 +515,10 @@ def _index_by_root_op( edge_names=("attn_qk_to_scale", "attn_scale_to_softmax", "attn_softmax_to_value"), inputs=(_input("query", "attn_qk", 0), _input("key", "attn_qk", 1), _input("value", "attn_value", 1)), attr_captures=(M.AttrCapture("scale", "attn_scale", "other", required=False), M.AttrCapture("is_causal", default=True, required=False), M.AttrCapture("window_size", default=0, required=False)), - constraints=("QK matmul must be over query/key head dimensions.", "Softmax axis must be the attention score axis.", "Value matmul must consume attention probabilities and value states."), + constraints={ + "node_attr_equals": {"node": "attn_softmax", "attr": "dim", "value": -1}, + "parent_tensor_dim_equal": {"left_node": "attn_qk", "left_parent_index": 0, "left_dim": -1, "right_node": "attn_qk", "right_parent_index": 1, "right_dim": -1, "allow_missing": True}, + }, ) ATTENTION_MASKED_GRAPH = _graph( @@ -514,7 +528,9 @@ def _index_by_root_op( edge_names=("attn_qk_to_scale", "attn_scale_to_mask_add", "attn_mask_add_to_softmax", "attn_softmax_to_value"), inputs=(_input("query", "attn_qk", 0), _input("key", "attn_qk", 1), _input("mask", "attn_mask_add", 1), _input("value", "attn_value", 1)), attr_captures=(M.AttrCapture("scale", "attn_scale", "other", required=False), M.AttrCapture("use_mask", default=True), M.AttrCapture("additive_mask", default=True, required=False)), - constraints=("Mask input must be broadcast-compatible with attention scores.",), + constraints={ + "input_broadcastable_to_node": {"role": "mask", "node": "attn_qk", "allow_missing": True}, + }, ) ROPE_GRAPH = _graph( @@ -528,7 +544,7 @@ def _index_by_root_op( _shared_input("rope_slice_even", 0, "rope_cos_mul", 0), ), attr_captures=(M.AttrCapture("theta", default=None, required=False), M.AttrCapture("position_offset", default=0, required=False)), - constraints=("RoPE cos/sin tables or generated constants must correspond to the same theta and position offset.",), + constraints=_note("RoPE cos/sin tables or generated constants must correspond to the same theta and position offset."), ) CONV_GRAPH = _graph( @@ -569,7 +585,9 @@ def _index_by_root_op( _cache_input("value_cache_state", "attn_value", 1, tensor_role=M.CacheTensorRole.VALUE, optional=True), ), attr_captures=(M.AttrCapture("scale", "attn_scale", "other", required=False), M.AttrCapture("position_offset", default=0, required=False), M.AttrCapture("window_size", default=0, required=False)), - constraints=("Cached attention must verify that key/value cache states belong to the same decoder layer.",), + constraints={ + "same_layer": {"nodes": ("attn_qk", "attn_value"), "allow_missing": True}, + }, ) EXPERT_BRANCH_GRAPH = _graph( @@ -579,7 +597,7 @@ def _index_by_root_op( edge_names=("expert_gate_to_activation", "expert_activation_to_product", "expert_up_to_product", "expert_product_to_down", "expert_down_to_weighted"), inputs=(_input("hidden", "expert_gate", 0), _input("gate_weight", "expert_gate", 1), _input("up_weight", "expert_up", 1), _input("down_weight", "expert_down", 1), _input("routing_weight", "expert_weighted", 1)), shared_inputs=(_shared_input("expert_gate", 0, "expert_up", 0),), - constraints=("Each repeated expert branch must correspond to the same expert index selected by routing.",), + constraints=_note("Each repeated expert branch must correspond to the same expert index selected by routing."), allow_root_external_children=False, ) @@ -591,7 +609,7 @@ def _index_by_root_op( inputs=(_input("hidden", "router_logits", 0), _input("router_weight", "router_logits", 1)), repeated_subgraphs=(M.RepeatedSubgraph("experts", EXPERT_BRANCH_GRAPH, min_count=1, anchor_node="topk"),), attr_captures=(M.AttrCapture("num_experts", default=None, required=False), M.AttrCapture("num_experts_per_tok", "topk", "k", required=False), M.AttrCapture("normalize_routing", default=True, required=False), M.AttrCapture("epsilon", default=1e-6, required=False), M.AttrCapture("routed_scaling_factor", default=1.0, required=False)), - constraints=("Top-k indices must select the repeated expert branches.", "Routing probabilities must be the weights used for expert output combination."), + constraints=_note("Top-k indices must select the repeated expert branches.", "Routing probabilities must be the weights used for expert output combination."), ) LSTM_CELL_GRAPH = _graph( @@ -600,7 +618,7 @@ def _index_by_root_op( ("lstm_gate_mm", "lstm_recurrent_mm", "lstm_gate_add", "lstm_sigmoid", "lstm_tanh"), edge_names=E.EDGE_GROUPS["lstm_cell"], inputs=(_input("input", "lstm_gate_mm", 0), _input("h_prev", "lstm_recurrent_mm", 0), _input("weight_ih", "lstm_gate_mm", 1), _input("weight_hh", "lstm_recurrent_mm", 1)), - constraints=("IR must verify the four LSTM gate splits and c_prev update path before using this fusion.",), + constraints=_note("IR must verify the four LSTM gate splits and c_prev update path before using this fusion."), ) DELTANET_DECODE_GRAPH = _graph( diff --git a/python/cactus/transpiler/Fusions/models.py b/python/cactus/transpiler/Fusions/models.py index 4eacdafed..b5090034f 100644 --- a/python/cactus/transpiler/Fusions/models.py +++ b/python/cactus/transpiler/Fusions/models.py @@ -2,6 +2,10 @@ from typing import Any +ConstraintSpec = dict[str, Any] +ConstraintValue = ConstraintSpec | tuple[ConstraintSpec, ...] | list[ConstraintSpec] + + class ValueKind: """ Names the semantic role of a real IR value. @@ -304,9 +308,9 @@ class FusionGraph: Describes the complete synthetic DAG for one candidate fusion. It stores pattern nodes, required edges, exposed inputs/outputs, attrs to - capture, repeated subgraphs, cache boundaries, and extra textual constraints - that IR matchers must enforce before replacing the matched real subgraph - with a fused op. + capture, repeated subgraphs, cache boundaries, and extra structured + constraints that IR matchers must enforce before replacing the matched real + subgraph with a fused op. """ name: str @@ -321,7 +325,7 @@ class FusionGraph: cache_inputs: tuple[CacheInput, ...] = () cache_outputs: tuple[CacheOutput, ...] = () cache_mutations: tuple[CacheMutation, ...] = () - constraints: tuple[str, ...] = () + constraints: dict[str, ConstraintValue] = field(default_factory=dict) variants: tuple[str, ...] = () metadata: dict[str, Any] = field(default_factory=dict) allow_root_external_children: bool = True diff --git a/python/cactus/transpiler/IR/extra_matchers.py b/python/cactus/transpiler/IR/extra_matchers.py new file mode 100644 index 000000000..2192faf05 --- /dev/null +++ b/python/cactus/transpiler/IR/extra_matchers.py @@ -0,0 +1,205 @@ +from typing import Any, Callable + +from . import models, match_utils +from ..Fusions import models as FModels + + +ExtraMatcher = Callable[[models.Node, models.Graph, FModels.FusionGraph, dict[str, models.Node], dict[str, Any]], bool] + + +############################################ Extra Matcher Utils!!!!! ############################################ +def get_bound_parent(bindings: dict[str, models.Node], node_name: str, parent_index: int) -> models.Node | None: + node = bindings.get(node_name) + + if node is None: + return None + + return match_utils.get_parent(node, parent_index) + + +def get_node_dim(node: models.Node, dim_index: int) -> Any: + shape = match_utils.get_tensor_shape(node) + + if not shape: + return match_utils.MISSING + + if dim_index < 0: + dim_index += len(shape) + + if dim_index < 0 or dim_index >= len(shape): + return match_utils.MISSING + + return shape[dim_index] + + +def shapes_are_broadcastable(left_shape: list[Any], right_shape: list[Any]) -> bool: + for left_dim, right_dim in zip(reversed(left_shape), reversed(right_shape)): + if match_utils.values_equal(left_dim, 1) or match_utils.values_equal(right_dim, 1): + continue + + if not match_utils.values_equal(left_dim, right_dim): + return False + + return True + + +def get_layer_key(node: models.Node) -> str | None: + for value in (module_stack_text(node), node.name, node.target): + layer_key = first_layer_key(value) + + if layer_key is not None: + return layer_key + + return None + + +def module_stack_text(node: models.Node) -> str: + return "" if node.module_stack is None else repr(node.module_stack) + + +def first_layer_key(value: str) -> str | None: + parts = value.replace("/", ".").replace("_", ".").split(".") + + for index, part in enumerate(parts[:-1]): + if part in {"layers", "layer", "h", "blocks", "block"} and parts[index + 1].isdigit(): + return f"{part}.{parts[index + 1]}" + + return None + + + +############################################ Extra Matchers!!!!! ############################################ +def match_extra_constraints(source: models.Node, graph: models.Graph, fusion: FModels.FusionGraph, bindings: dict[str, models.Node]) -> bool: + for constraint_name, raw_spec in fusion.constraints.items(): + matcher = EXTRA_MATCHERS.get(constraint_name) + + if matcher is None: + return False + + specs = normalize_constraint_specs(raw_spec) + + if specs is None: + return False + + for spec in specs: + if not matcher(source, graph, fusion, bindings, spec): + return False + + return True + + +def normalize_constraint_specs(raw_spec: Any) -> tuple[dict[str, Any], ...] | None: + if isinstance(raw_spec, dict): + return (raw_spec,) + + if isinstance(raw_spec, (list, tuple)) and all(isinstance(spec, dict) for spec in raw_spec): + return tuple(raw_spec) + + return None + + +def match_note(source: models.Node, graph: models.Graph, fusion: FModels.FusionGraph, bindings: dict[str, models.Node], spec: dict[str, Any]) -> bool: + return True + + +def match_node_attr_equals(source: models.Node, graph: models.Graph, fusion: FModels.FusionGraph, bindings: dict[str, models.Node], spec: dict[str, Any]) -> bool: + node = bindings.get(spec["node"]) + + if node is None: + return False + + actual = node.attrs.get(spec["attr"], match_utils.MISSING) + + if actual is match_utils.MISSING: + return bool(spec.get("allow_missing", False)) + + return match_utils.values_equal(actual, spec["value"]) + + +def match_node_attrs_equal(source: models.Node, graph: models.Graph, fusion: FModels.FusionGraph, bindings: dict[str, models.Node], spec: dict[str, Any]) -> bool: + left_node = bindings.get(spec["left_node"]) + right_node = bindings.get(spec["right_node"]) + + if left_node is None or right_node is None: + return False + + left_value = left_node.attrs.get(spec["left_attr"], match_utils.MISSING) + right_value = right_node.attrs.get(spec["right_attr"], match_utils.MISSING) + + if left_value is match_utils.MISSING or right_value is match_utils.MISSING: + return bool(spec.get("allow_missing", False)) + + return match_utils.values_equal(left_value, right_value) + + +def match_input_value_kind(source: models.Node, graph: models.Graph, fusion: FModels.FusionGraph, bindings: dict[str, models.Node], spec: dict[str, Any]) -> bool: + input_node = match_utils.get_first_input_by_role(fusion, bindings, spec["role"]) + + if input_node is None: + return bool(spec.get("allow_missing", False)) + + return input_node.value_kind in spec["allowed_value_kinds"] + + +def match_parent_tensor_dim_equal(source: models.Node, graph: models.Graph, fusion: FModels.FusionGraph, bindings: dict[str, models.Node], spec: dict[str, Any]) -> bool: + left_parent = get_bound_parent(bindings, spec["left_node"], spec["left_parent_index"]) + right_parent = get_bound_parent(bindings, spec["right_node"], spec["right_parent_index"]) + + if left_parent is None or right_parent is None: + return bool(spec.get("allow_missing", False)) + + left_dim = get_node_dim(left_parent, spec["left_dim"]) + right_dim = get_node_dim(right_parent, spec["right_dim"]) + + if left_dim is match_utils.MISSING or right_dim is match_utils.MISSING: + return bool(spec.get("allow_missing", False)) + + return match_utils.values_equal(left_dim, right_dim) + + +def match_input_broadcastable_to_node(source: models.Node, graph: models.Graph, fusion: FModels.FusionGraph, bindings: dict[str, models.Node], spec: dict[str, Any]) -> bool: + input_node = match_utils.get_first_input_by_role(fusion, bindings, spec["role"]) + target_node = bindings.get(spec["node"]) + + if input_node is None or target_node is None: + return bool(spec.get("allow_missing", False)) + + input_shape = match_utils.get_tensor_shape(input_node) + target_shape = match_utils.get_tensor_shape(target_node) + + if not input_shape or not target_shape: + return bool(spec.get("allow_missing", False)) + + return shapes_are_broadcastable(input_shape, target_shape) + + +def match_same_layer(source: models.Node, graph: models.Graph, fusion: FModels.FusionGraph, bindings: dict[str, models.Node], spec: dict[str, Any]) -> bool: + layer_keys = [] + + for node_name in spec["nodes"]: + node = bindings.get(node_name) + + if node is None: + return False + + layer_key = get_layer_key(node) + + if layer_key is None: + return bool(spec.get("allow_missing", False)) + + layer_keys.append(layer_key) + + return len(set(layer_keys)) == 1 + + + + +EXTRA_MATCHERS: dict[str, ExtraMatcher] = { + "note": match_note, + "node_attr_equals": match_node_attr_equals, + "node_attrs_equal": match_node_attrs_equal, + "input_value_kind": match_input_value_kind, + "parent_tensor_dim_equal": match_parent_tensor_dim_equal, + "input_broadcastable_to_node": match_input_broadcastable_to_node, + "same_layer": match_same_layer, +} diff --git a/python/cactus/transpiler/IR/match.py b/python/cactus/transpiler/IR/match.py index 549fb2b4a..de925d276 100644 --- a/python/cactus/transpiler/IR/match.py +++ b/python/cactus/transpiler/IR/match.py @@ -1,4 +1,4 @@ -from . import models, match_utils +from . import extra_matchers, models, match_utils from ..Fusions import models as FModels ANY = () @@ -141,7 +141,7 @@ def match_shared_inputs(node: models.Node, graph: models.Graph, fusion: FModels. def match_constraints(node: models.Node, graph: models.Graph, fusion: FModels.FusionGraph, bindings: dict[str, models.Node]) -> bool: - return True + return extra_matchers.match_extra_constraints(node, graph, fusion, bindings) def match_cache_inputs(node: models.Node, graph: models.Graph, fusion: FModels.FusionGraph, bindings: dict[str, models.Node]) -> bool: From 610d3b4566288a2e732be8d10b3bda9d150a8b20 Mon Sep 17 00:00:00 2001 From: Parkirat Date: Tue, 28 Jul 2026 15:15:10 -0400 Subject: [PATCH 53/84] Additional extra matchers for RoPE, linear projs, and LSTMs (yet to be confirmed) --- python/cactus/transpiler/Fusions/fusions.py | 13 +- python/cactus/transpiler/IR/extra_matchers.py | 323 ++++++++++++++++++ 2 files changed, 333 insertions(+), 3 deletions(-) diff --git a/python/cactus/transpiler/Fusions/fusions.py b/python/cactus/transpiler/Fusions/fusions.py index 478091042..7ad23f2f2 100644 --- a/python/cactus/transpiler/Fusions/fusions.py +++ b/python/cactus/transpiler/Fusions/fusions.py @@ -418,7 +418,9 @@ def _index_by_root_op( ("linear_mm",), inputs=(_input("x", "linear_mm", 0), _input("weight", "linear_mm", 1)), attr_captures=(M.AttrCapture("pretransposed_rhs", default=False, required=False),), - constraints=_note("Weight layout/pretranspose must be verified against exported attrs or weight metadata."), + constraints={ + "linear_weight_layout": {"input_role": "x", "weight_role": "weight", "output_node": "linear_mm"}, + }, ) LINEAR_BIAS_GRAPH = _graph( @@ -429,6 +431,7 @@ def _index_by_root_op( inputs=(_input("x", "linear_mm", 0), _input("weight", "linear_mm", 1), _input("bias", "linear_bias_add", 1)), attr_captures=(M.AttrCapture("pretransposed_rhs", default=False, required=False),), constraints={ + "linear_weight_layout": {"input_role": "x", "weight_role": "weight", "output_node": "linear_mm"}, "input_value_kind": {"role": "bias", "allowed_value_kinds": (M.ValueKind.PARAMETER, M.ValueKind.BUFFER)}, }, ) @@ -544,7 +547,9 @@ def _index_by_root_op( _shared_input("rope_slice_even", 0, "rope_cos_mul", 0), ), attr_captures=(M.AttrCapture("theta", default=None, required=False), M.AttrCapture("position_offset", default=0, required=False)), - constraints=_note("RoPE cos/sin tables or generated constants must correspond to the same theta and position offset."), + constraints={ + "rope_tables_compatible": {"x_role": "x", "cos_role": "cos", "sin_role": "sin"}, + }, ) CONV_GRAPH = _graph( @@ -618,7 +623,9 @@ def _index_by_root_op( ("lstm_gate_mm", "lstm_recurrent_mm", "lstm_gate_add", "lstm_sigmoid", "lstm_tanh"), edge_names=E.EDGE_GROUPS["lstm_cell"], inputs=(_input("input", "lstm_gate_mm", 0), _input("h_prev", "lstm_recurrent_mm", 0), _input("weight_ih", "lstm_gate_mm", 1), _input("weight_hh", "lstm_recurrent_mm", 1)), - constraints=_note("IR must verify the four LSTM gate splits and c_prev update path before using this fusion."), + constraints={ + "lstm_gate_structure": {"gate_node": "lstm_gate_add", "gate_count": 4, "require_explicit_gate_split": True}, + }, ) DELTANET_DECODE_GRAPH = _graph( diff --git a/python/cactus/transpiler/IR/extra_matchers.py b/python/cactus/transpiler/IR/extra_matchers.py index 2192faf05..65bfa1ab5 100644 --- a/python/cactus/transpiler/IR/extra_matchers.py +++ b/python/cactus/transpiler/IR/extra_matchers.py @@ -7,6 +7,34 @@ ExtraMatcher = Callable[[models.Node, models.Graph, FModels.FusionGraph, dict[str, models.Node], dict[str, Any]], bool] +WEIGHT_VALUE_KINDS = ( + FModels.ValueKind.PARAMETER, + FModels.ValueKind.BUFFER, + FModels.ValueKind.LIFTED_CONSTANT, +) + +COPY_TARGETS = { + "aten.clone.default", + "aten._to_copy.default", + "aten.contiguous.default", +} + +TRANSPOSE_TARGETS = { + "aten.t.default", + "aten.transpose.int", + "aten.permute.default", +} + +COS_TARGETS = {"aten.cos.default"} +SIN_TARGETS = {"aten.sin.default"} +SPLIT_TARGETS = { + "aten.chunk.default", + "aten.split.Tensor", + "aten.split_with_sizes.default", + "aten.unbind.int", +} + + ############################################ Extra Matcher Utils!!!!! ############################################ def get_bound_parent(bindings: dict[str, models.Node], node_name: str, parent_index: int) -> models.Node | None: node = bindings.get(node_name) @@ -32,6 +60,233 @@ def get_node_dim(node: models.Node, dim_index: int) -> Any: return shape[dim_index] +def classify_linear_rhs(node: models.Node) -> tuple[str, models.Node] | None: + return classify_linear_rhs_(node, set()) + + +def classify_linear_rhs_(node: models.Node, seen: set[int]) -> tuple[str, models.Node] | None: + if id(node) in seen: + return None + + seen.add(id(node)) + + if node.target in COPY_TARGETS and len(node.parents) == 1: + return classify_linear_rhs_(node.parents[0], seen) + + if node.value_kind in WEIGHT_VALUE_KINDS: + return ("direct", node) + + if node.target in TRANSPOSE_TARGETS and len(node.parents) == 1 and node.parents[0].value_kind in WEIGHT_VALUE_KINDS: + if linear_transpose_is_valid(node, node.parents[0]): + return ("transposed", node.parents[0]) + + return None + + +def linear_transpose_is_valid(node: models.Node, source_weight: models.Node) -> bool: + source_shape = match_utils.get_tensor_shape(source_weight) + + if len(source_shape) != 2: + return False + + if node.target == "aten.t.default": + return True + + if node.target == "aten.transpose.int": + dim0 = node.attrs.get("dim0", match_utils.MISSING) + dim1 = node.attrs.get("dim1", match_utils.MISSING) + return {normalize_dim(dim0, 2), normalize_dim(dim1, 2)} == {0, 1} + + permutation = node.attrs.get("permutation", match_utils.MISSING) + return match_utils.values_equal(permutation, [1, 0]) + + +def linear_shapes_match(input_node: models.Node, rhs_node: models.Node, output_node: models.Node, source_weight: models.Node, layout_kind: str, spec: dict[str, Any]) -> bool: + input_shape = match_utils.get_tensor_shape(input_node) + rhs_shape = match_utils.get_tensor_shape(rhs_node) + output_shape = match_utils.get_tensor_shape(output_node) + + if len(rhs_shape) < 2: + return bool(spec.get("allow_missing", False)) + + if not input_shape or not output_shape: + return bool(spec.get("allow_missing", False)) + + input_dim = input_shape[-1] + output_dim = output_shape[-1] + + if not match_utils.values_equal(rhs_shape[-2], input_dim): + return False + + if not match_utils.values_equal(rhs_shape[-1], output_dim): + return False + + source_shape = match_utils.get_tensor_shape(source_weight) + + if layout_kind == "transposed" and len(source_shape) >= 2: + if not match_utils.values_equal(source_shape[-2], output_dim): + return False + + if not match_utils.values_equal(source_shape[-1], input_dim): + return False + + return True + + +def compatible_shapes(left_shape: list[Any], right_shape: list[Any], spec: dict[str, Any]) -> bool: + if not left_shape or not right_shape: + return bool(spec.get("allow_missing", False)) + + return ( + match_utils.values_equal(left_shape, right_shape) + or shapes_are_broadcastable(left_shape, right_shape) + or shapes_are_broadcastable(right_shape, left_shape) + ) + + +def compatible_dtypes(left_node: models.Node, right_node: models.Node, spec: dict[str, Any]) -> bool: + if not spec.get("require_same_dtype", True): + return True + + left_dtype = get_node_dtype(left_node) + right_dtype = get_node_dtype(right_node) + + if left_dtype is match_utils.MISSING or right_dtype is match_utils.MISSING: + return bool(spec.get("allow_missing", False)) + + return match_utils.values_equal(left_dtype, right_dtype) + + +def compatible_trig_sources(cos_node: models.Node, sin_node: models.Node, spec: dict[str, Any]) -> bool: + cos_sources = collect_trig_angle_sources(cos_node, COS_TARGETS, int(spec.get("max_trig_depth", 8))) + sin_sources = collect_trig_angle_sources(sin_node, SIN_TARGETS, int(spec.get("max_trig_depth", 8))) + + if cos_sources or sin_sources: + return cos_sources == sin_sources + + return not spec.get("require_trig_sources", False) + + +def collect_trig_angle_sources(node: models.Node, trig_targets: set[str], max_depth: int) -> frozenset[str]: + sources: set[str] = set() + seen: set[int] = set() + + def visit(current: models.Node, depth: int) -> None: + if depth > max_depth or id(current) in seen: + return + + seen.add(id(current)) + + if current.target in trig_targets: + sources.add(current.parents[0].name if current.parents else current.name) + return + + for parent in current.parents: + visit(parent, depth + 1) + + visit(node, 0) + return frozenset(sources) + + +def rope_tables_match_input(x_node: models.Node, cos_shape: list[Any], sin_shape: list[Any], spec: dict[str, Any]) -> bool: + x_shape = match_utils.get_tensor_shape(x_node) + + if not x_shape: + return bool(spec.get("allow_missing", False)) + + return rope_table_shape_matches_x(cos_shape, x_shape) and rope_table_shape_matches_x(sin_shape, x_shape) + + +def rope_table_shape_matches_x(table_shape: list[Any], x_shape: list[Any]) -> bool: + if not table_shape: + return False + + if shapes_are_broadcastable(table_shape, x_shape): + return True + + table_last = table_shape[-1] + x_last = x_shape[-1] + + if isinstance(table_last, int) and isinstance(x_last, int) and table_last > 0: + return x_last in {table_last, table_last * 2} + + return False + + +def gate_shape_matches(gate_node: models.Node, gate_count: int, spec: dict[str, Any]) -> bool: + gate_shape = match_utils.get_tensor_shape(gate_node) + + if not gate_shape: + return bool(spec.get("allow_missing", False)) + + hidden_dim = gate_shape[-1] + + if not isinstance(hidden_dim, int): + return bool(spec.get("allow_symbolic", True)) + + return hidden_dim % gate_count == 0 + + +def has_explicit_gate_split(gate_node: models.Node, gate_count: int, max_depth: int) -> bool: + for node in walk_children(gate_node, max_depth): + if node.target not in SPLIT_TARGETS: + continue + + if split_node_has_gate_count(node, gate_count): + return True + + return False + + +def split_node_has_gate_count(node: models.Node, gate_count: int) -> bool: + if node.target == "aten.unbind.int": + return len(node.children) == gate_count + + if node.target == "aten.chunk.default": + return match_utils.values_equal(node.attrs.get("arg_1", match_utils.MISSING), gate_count) + + split_sizes = node.attrs.get("arg_1", match_utils.MISSING) + + if isinstance(split_sizes, (list, tuple)): + return len(split_sizes) == gate_count + + return len(node.children) == gate_count + + +def walk_children(node: models.Node, max_depth: int) -> tuple[models.Node, ...]: + found: list[models.Node] = [] + seen: set[int] = set() + + def visit(current: models.Node, depth: int) -> None: + if depth > max_depth: + return + + for child in current.children: + if id(child) in seen: + continue + + seen.add(id(child)) + found.append(child) + visit(child, depth + 1) + + visit(node, 0) + return tuple(found) + + +def get_node_dtype(node: models.Node) -> Any: + if not isinstance(node.tensor_output_meta, dict): + return match_utils.MISSING + + return node.tensor_output_meta.get("dtype", match_utils.MISSING) + + +def normalize_dim(dim: Any, rank: int) -> Any: + if not isinstance(dim, int): + return dim + + return dim + rank if dim < 0 else dim + + def shapes_are_broadcastable(left_shape: list[Any], right_shape: list[Any]) -> bool: for left_dim, right_dim in zip(reversed(left_shape), reversed(right_shape)): if match_utils.values_equal(left_dim, 1) or match_utils.values_equal(right_dim, 1): @@ -102,6 +357,71 @@ def match_note(source: models.Node, graph: models.Graph, fusion: FModels.FusionG return True +def match_linear_weight_layout(source: models.Node, graph: models.Graph, fusion: FModels.FusionGraph, bindings: dict[str, models.Node], spec: dict[str, Any]) -> bool: + input_node = match_utils.get_first_input_by_role(fusion, bindings, spec["input_role"]) + weight_node = match_utils.get_first_input_by_role(fusion, bindings, spec["weight_role"]) + output_node = bindings.get(spec["output_node"]) + + if input_node is None or weight_node is None or output_node is None: + return bool(spec.get("allow_missing", False)) + + layout = classify_linear_rhs(weight_node) + + if layout is None: + return False + + layout_kind, source_weight = layout + allowed_layouts = spec.get("allowed_layouts", ("direct", "transposed")) + + if layout_kind not in allowed_layouts: + return False + + return linear_shapes_match(input_node, weight_node, output_node, source_weight, layout_kind, spec) + + +def match_rope_tables_compatible(source: models.Node, graph: models.Graph, fusion: FModels.FusionGraph, bindings: dict[str, models.Node], spec: dict[str, Any]) -> bool: + x_node = match_utils.get_first_input_by_role(fusion, bindings, spec["x_role"]) + cos_node = match_utils.get_first_input_by_role(fusion, bindings, spec["cos_role"]) + sin_node = match_utils.get_first_input_by_role(fusion, bindings, spec["sin_role"]) + + if cos_node is None or sin_node is None: + return bool(spec.get("allow_missing", False)) + + cos_shape = match_utils.get_tensor_shape(cos_node) + sin_shape = match_utils.get_tensor_shape(sin_node) + + if not compatible_shapes(cos_shape, sin_shape, spec): + return False + + if not compatible_dtypes(cos_node, sin_node, spec): + return False + + if not compatible_trig_sources(cos_node, sin_node, spec): + return False + + if x_node is not None and not rope_tables_match_input(x_node, cos_shape, sin_shape, spec): + return False + + return True + + +def match_lstm_gate_structure(source: models.Node, graph: models.Graph, fusion: FModels.FusionGraph, bindings: dict[str, models.Node], spec: dict[str, Any]) -> bool: + gate_node = bindings.get(spec["gate_node"]) + + if gate_node is None: + return bool(spec.get("allow_missing", False)) + + gate_count = int(spec.get("gate_count", 4)) + + if not gate_shape_matches(gate_node, gate_count, spec): + return False + + if spec.get("require_explicit_gate_split", True): + return has_explicit_gate_split(gate_node, gate_count, int(spec.get("max_depth", 4))) + + return True + + def match_node_attr_equals(source: models.Node, graph: models.Graph, fusion: FModels.FusionGraph, bindings: dict[str, models.Node], spec: dict[str, Any]) -> bool: node = bindings.get(spec["node"]) @@ -195,6 +515,9 @@ def match_same_layer(source: models.Node, graph: models.Graph, fusion: FModels.F EXTRA_MATCHERS: dict[str, ExtraMatcher] = { + "linear_weight_layout": match_linear_weight_layout, + "rope_tables_compatible": match_rope_tables_compatible, + "lstm_gate_structure": match_lstm_gate_structure, "note": match_note, "node_attr_equals": match_node_attr_equals, "node_attrs_equal": match_node_attrs_equal, From 80243d3205e00382e58dd3917817dcf787361f3c Mon Sep 17 00:00:00 2001 From: Parkirat Date: Tue, 28 Jul 2026 15:54:40 -0400 Subject: [PATCH 54/84] Minor modifications to converter and input processing to support MOE models --- .../transpiler/Converter/input_processor.py | 13 +++++++ python/cactus/transpiler/Converter/models.py | 34 +++++++++++++------ .../cactus/transpiler/Converter/overrides.py | 8 ++++- 3 files changed, 44 insertions(+), 11 deletions(-) diff --git a/python/cactus/transpiler/Converter/input_processor.py b/python/cactus/transpiler/Converter/input_processor.py index 35330262b..095bb432c 100644 --- a/python/cactus/transpiler/Converter/input_processor.py +++ b/python/cactus/transpiler/Converter/input_processor.py @@ -75,10 +75,23 @@ def default_processor(model_id: str, configs: dict[str, dict[str, Any]], model_p raise RuntimeError(f"Unable to load processor for {model_id}") +def text_tokenizer_processor(model_id: str, configs: dict[str, dict[str, Any]], model_profile: str): + from transformers import PreTrainedTokenizerFast + + tokenizer_config = configs.get("tokenizer_config.json", {}) + return PreTrainedTokenizerFast( + tokenizer_file=_local_tokenizer_path(model_id=model_id, model_profile=model_profile), + bos_token=tokenizer_config.get("bos_token"), + eos_token=tokenizer_config.get("eos_token"), + pad_token=tokenizer_config.get("pad_token"), + ) + + PROCESSOR_MAP = { "google/gemma-4-E2B": gemma4_processor, "openai/whisper-tiny": default_processor, "nvidia/parakeet-tdt-0.6b-v3": default_processor, "LiquidAI/LFM2-VL-3B": default_processor, "Qwen/Qwen2.5-0.5B": default_processor, + "LiquidAI/LFM2.5-8B-A1B": text_tokenizer_processor, } diff --git a/python/cactus/transpiler/Converter/models.py b/python/cactus/transpiler/Converter/models.py index a930ec4db..73a7387ed 100644 --- a/python/cactus/transpiler/Converter/models.py +++ b/python/cactus/transpiler/Converter/models.py @@ -18,6 +18,7 @@ token = constants.token EXPORT_PATCHES = { "gemma4_audio_mask": OV.patch_gemma4_audio_mask_for_export, + "transformers_moe_grouped_mm_fallback": OV.patch_transformers_moe_grouped_mm_for_export, } @@ -303,25 +304,38 @@ def load_model(model_id: str, mp: MP_Models.ModelProfile | None = None) -> torch export_patches = mp.export_patches seen_classes: set[Any] = set() + base_kwargs: dict[str, Any] = {"trust_remote_code": True} + last_error: Exception | None = None + + if constants.token is not None: + base_kwargs["token"] = constants.token + + load_attempts = ( + {**base_kwargs, "dtype": "auto", "low_cpu_mem_usage": True}, + {**base_kwargs, "torch_dtype": "auto", "low_cpu_mem_usage": True}, + base_kwargs, + ) for model_class in candidate_classes: if model_class in seen_classes: continue seen_classes.add(model_class) - try: - model = model_class.from_pretrained(model_id) - model.eval() - except: - pass - for patch in export_patches: - patch_fn = EXPORT_PATCHES[patch] - patch_fn() + for load_kwargs in load_attempts: + try: + model = model_class.from_pretrained(model_id, **load_kwargs) + model.eval() + + for patch in export_patches: + patch_fn = EXPORT_PATCHES[patch] + patch_fn() - return model + return model + except Exception as e: + last_error = e - raise RuntimeError(f"Unable to load model {model_id}") + raise RuntimeError(f"Unable to load model {model_id}") from last_error #Builds a loaded model bundle from profile, modalities, model id, and inference mode. diff --git a/python/cactus/transpiler/Converter/overrides.py b/python/cactus/transpiler/Converter/overrides.py index 32338820c..b1c2fb111 100644 --- a/python/cactus/transpiler/Converter/overrides.py +++ b/python/cactus/transpiler/Converter/overrides.py @@ -22,4 +22,10 @@ def exportable_bidirectional_mask(config, inputs_embeds, attention_mask=None, an return mask.expand(batch_size, 1, seq_len, seq_len) - gemma4_modeling.create_bidirectional_mask = exportable_bidirectional_mask \ No newline at end of file + gemma4_modeling.create_bidirectional_mask = exportable_bidirectional_mask + + +def patch_transformers_moe_grouped_mm_for_export() -> None: + import transformers.integrations.moe as moe + + moe._can_use_grouped_mm = lambda input, weight, offs: False From 8aa76727470116ae8fb739d06fe4f92dd5030d69 Mon Sep 17 00:00:00 2001 From: Parkirat Date: Tue, 28 Jul 2026 15:58:50 -0400 Subject: [PATCH 55/84] Additional fusions + nodes for MOE ops --- python/cactus/transpiler/Fusions/edges.py | 118 ++++++++++++++++++++ python/cactus/transpiler/Fusions/fusions.py | 103 +++++++++++++++-- python/cactus/transpiler/Fusions/nodes.py | 113 +++++++++++++++++++ 3 files changed, 323 insertions(+), 11 deletions(-) diff --git a/python/cactus/transpiler/Fusions/edges.py b/python/cactus/transpiler/Fusions/edges.py index 56ac199a4..3c5d676dc 100644 --- a/python/cactus/transpiler/Fusions/edges.py +++ b/python/cactus/transpiler/Fusions/edges.py @@ -107,6 +107,64 @@ def _index_edges_by_dest(edges: dict[str, M.FusionEdge]) -> dict[str, tuple[str, "delta_v_to_decode": _edge("delta_v", "delta_gate", 2), "router_logits_to_probs": _edge("router_logits", "routing_probs", 0), "routing_probs_to_topk": _edge("routing_probs", "topk", 0), + "moe_hidden_to_router": _edge("moe_hidden_view", "moe_router_logits", 0), + "moe_router_weight_to_router": _edge("moe_router_weight_transpose", "moe_router_logits", 1), + "moe_router_logits_to_sigmoid": _edge("moe_router_logits", "moe_router_sigmoid", 0), + "moe_router_sigmoid_to_bias": _edge("moe_router_sigmoid", "moe_router_bias", 0), + "moe_router_bias_to_topk": _edge("moe_router_bias", "moe_topk", 0), + "moe_topk_to_indices": _edge("moe_topk", "moe_topk_indices", 0), + "moe_router_sigmoid_to_gather": _edge("moe_router_sigmoid", "moe_routing_gather", 0), + "moe_topk_indices_to_gather": _edge("moe_topk_indices", "moe_routing_gather", 1), + "moe_gather_to_routing_sum": _edge("moe_routing_gather", "moe_routing_sum", 0), + "moe_routing_sum_to_eps": _edge("moe_routing_sum", "moe_routing_eps", 0), + "moe_gather_to_norm": _edge("moe_routing_gather", "moe_routing_norm", 0), + "moe_eps_to_norm": _edge("moe_routing_eps", "moe_routing_norm", 1), + "moe_norm_to_scale": _edge("moe_routing_norm", "moe_routing_scale", 0), + "moe_token_arange_to_unsqueeze": _edge("moe_token_arange", "moe_token_unsqueeze", 0), + "moe_token_unsqueeze_to_expand": _edge("moe_token_unsqueeze", "moe_token_expand", 0), + "moe_token_expand_to_clone": _edge("moe_token_expand", "moe_token_clone", 0), + "moe_token_clone_to_flat": _edge("moe_token_clone", "moe_token_flat", 0), + "moe_scale_to_routing_flat": _edge("moe_routing_scale", "moe_routing_flat", 0), + "moe_topk_indices_to_flat": _edge("moe_topk_indices", "moe_topk_flat", 0), + "moe_topk_flat_to_sort": _edge("moe_topk_flat", "moe_sort", 0), + "moe_sort_to_indices": _edge("moe_sort", "moe_sort_indices", 0), + "moe_unsort_empty_to_index_put": _edge("moe_unsort_empty", "moe_unsort_index_put", 0), + "moe_sort_indices_to_index_put": _edge("moe_sort_indices", "moe_unsort_index_put", 1), + "moe_unsort_arange_to_index_put": _edge("moe_unsort_arange", "moe_unsort_index_put", 2), + "moe_topk_flat_to_sorted_expert_index": _edge("moe_topk_flat", "moe_sorted_expert_index", 0), + "moe_sort_indices_to_sorted_expert_index": _edge("moe_sort_indices", "moe_sorted_expert_index", 1), + "moe_sorted_expert_index_to_copy": _edge("moe_sorted_expert_index", "moe_sorted_expert_copy", 0), + "moe_sorted_expert_copy_to_histc": _edge("moe_sorted_expert_copy", "moe_histc", 0), + "moe_histc_to_offsets": _edge("moe_histc", "moe_offsets", 0), + "moe_routing_flat_to_sorted_routing_index": _edge("moe_routing_flat", "moe_sorted_routing_index", 0), + "moe_sort_indices_to_sorted_routing_index": _edge("moe_sort_indices", "moe_sorted_routing_index", 1), + "moe_token_flat_to_positions_index": _edge("moe_token_flat", "moe_token_positions_index", 0), + "moe_sort_indices_to_positions_index": _edge("moe_sort_indices", "moe_token_positions_index", 1), + "moe_hidden_to_sorted_hidden_index": _edge("moe_hidden_view", "moe_sorted_hidden_index", 0), + "moe_positions_to_sorted_hidden_index": _edge("moe_token_positions_index", "moe_sorted_hidden_index", 1), + "moe_gate_up_weight_to_grouped": _edge("moe_gate_up_weight_transpose", "moe_grouped_gate_up", 1), + "moe_sorted_hidden_to_grouped_gate_up": _edge("moe_sorted_hidden_index", "moe_grouped_gate_up", 0), + "moe_offsets_to_grouped_gate_up": _edge("moe_offsets", "moe_grouped_gate_up", 2), + "moe_grouped_gate_up_to_split": _edge("moe_grouped_gate_up", "moe_gate_up_split", 0), + "moe_split_to_gate": _edge("moe_gate_up_split", "moe_gate_tensor", 0), + "moe_split_to_up": _edge("moe_gate_up_split", "moe_up_tensor", 0), + "moe_gate_to_copy": _edge("moe_gate_tensor", "moe_gate_copy", 0), + "moe_gate_copy_to_sigmoid": _edge("moe_gate_copy", "moe_gate_sigmoid", 0), + "moe_gate_copy_to_activation": _edge("moe_gate_copy", "moe_gate_activation", 0), + "moe_gate_sigmoid_to_activation": _edge("moe_gate_sigmoid", "moe_gate_activation", 1), + "moe_gate_activation_to_copy": _edge("moe_gate_activation", "moe_gate_activation_copy", 0), + "moe_gate_activation_copy_to_product": _edge("moe_gate_activation_copy", "moe_expert_product", 0), + "moe_up_to_product": _edge("moe_up_tensor", "moe_expert_product", 1), + "moe_product_to_grouped_down": _edge("moe_expert_product", "moe_grouped_down", 0), + "moe_down_weight_to_grouped": _edge("moe_down_weight_transpose", "moe_grouped_down", 1), + "moe_offsets_to_grouped_down": _edge("moe_offsets", "moe_grouped_down", 2), + "moe_sorted_routing_to_unsqueeze": _edge("moe_sorted_routing_index", "moe_routing_unsqueeze", 0), + "moe_grouped_down_to_weighted": _edge("moe_grouped_down", "moe_grouped_weighted", 0), + "moe_routing_unsqueeze_to_weighted": _edge("moe_routing_unsqueeze", "moe_grouped_weighted", 1), + "moe_weighted_to_unsort_index": _edge("moe_grouped_weighted", "moe_unsort_index", 0), + "moe_unsort_index_put_to_unsort_index": _edge("moe_unsort_index_put", "moe_unsort_index", 1), + "moe_unsort_index_to_output_view": _edge("moe_unsort_index", "moe_output_view", 0), + "moe_output_view_to_combine": _edge("moe_output_view", "moe_grouped_combine", 0), "expert_gate_to_activation": _expert_edge("expert_gate", "expert_activation", 0), "expert_activation_to_product": _expert_edge("expert_activation", "expert_product", 0), "expert_up_to_product": _expert_edge("expert_up", "expert_product", 1), @@ -191,6 +249,66 @@ def _index_edges_by_dest(edges: dict[str, M.FusionEdge]) -> dict[str, tuple[str, "expert_down_to_weighted", "expert_weighted_to_combine", ), + "lfm_grouped_moe": ( + "moe_hidden_to_router", + "moe_router_weight_to_router", + "moe_router_logits_to_sigmoid", + "moe_router_sigmoid_to_bias", + "moe_router_bias_to_topk", + "moe_topk_to_indices", + "moe_router_sigmoid_to_gather", + "moe_topk_indices_to_gather", + "moe_gather_to_routing_sum", + "moe_routing_sum_to_eps", + "moe_gather_to_norm", + "moe_eps_to_norm", + "moe_norm_to_scale", + "moe_token_arange_to_unsqueeze", + "moe_token_unsqueeze_to_expand", + "moe_token_expand_to_clone", + "moe_token_clone_to_flat", + "moe_scale_to_routing_flat", + "moe_topk_indices_to_flat", + "moe_topk_flat_to_sort", + "moe_sort_to_indices", + "moe_unsort_empty_to_index_put", + "moe_sort_indices_to_index_put", + "moe_unsort_arange_to_index_put", + "moe_topk_flat_to_sorted_expert_index", + "moe_sort_indices_to_sorted_expert_index", + "moe_sorted_expert_index_to_copy", + "moe_sorted_expert_copy_to_histc", + "moe_histc_to_offsets", + "moe_routing_flat_to_sorted_routing_index", + "moe_sort_indices_to_sorted_routing_index", + "moe_token_flat_to_positions_index", + "moe_sort_indices_to_positions_index", + "moe_hidden_to_sorted_hidden_index", + "moe_positions_to_sorted_hidden_index", + "moe_gate_up_weight_to_grouped", + "moe_sorted_hidden_to_grouped_gate_up", + "moe_offsets_to_grouped_gate_up", + "moe_grouped_gate_up_to_split", + "moe_split_to_gate", + "moe_split_to_up", + "moe_gate_to_copy", + "moe_gate_copy_to_sigmoid", + "moe_gate_copy_to_activation", + "moe_gate_sigmoid_to_activation", + "moe_gate_activation_to_copy", + "moe_gate_activation_copy_to_product", + "moe_up_to_product", + "moe_product_to_grouped_down", + "moe_down_weight_to_grouped", + "moe_offsets_to_grouped_down", + "moe_sorted_routing_to_unsqueeze", + "moe_grouped_down_to_weighted", + "moe_routing_unsqueeze_to_weighted", + "moe_weighted_to_unsort_index", + "moe_unsort_index_put_to_unsort_index", + "moe_unsort_index_to_output_view", + "moe_output_view_to_combine", + ), } diff --git a/python/cactus/transpiler/Fusions/fusions.py b/python/cactus/transpiler/Fusions/fusions.py index 7ad23f2f2..ca2edf035 100644 --- a/python/cactus/transpiler/Fusions/fusions.py +++ b/python/cactus/transpiler/Fusions/fusions.py @@ -97,10 +97,6 @@ def _required_attrs(**attrs) -> dict: return {"required_attrs": attrs} -def _note(*items: str) -> dict[str, M.ConstraintSpec]: - return {"note": {"items": items}} - - def _graph( name: str, root: str, @@ -602,19 +598,102 @@ def _index_by_root_op( edge_names=("expert_gate_to_activation", "expert_activation_to_product", "expert_up_to_product", "expert_product_to_down", "expert_down_to_weighted"), inputs=(_input("hidden", "expert_gate", 0), _input("gate_weight", "expert_gate", 1), _input("up_weight", "expert_up", 1), _input("down_weight", "expert_down", 1), _input("routing_weight", "expert_weighted", 1)), shared_inputs=(_shared_input("expert_gate", 0, "expert_up", 0),), - constraints=_note("Each repeated expert branch must correspond to the same expert index selected by routing."), - allow_root_external_children=False, ) MOE_GATED_GRAPH = _graph( "moe_layer_gated", "moe_combine", - ("router_logits", "routing_probs", "topk", "moe_combine"), - edge_names=("router_logits_to_probs", "routing_probs_to_topk"), - inputs=(_input("hidden", "router_logits", 0), _input("router_weight", "router_logits", 1)), - repeated_subgraphs=(M.RepeatedSubgraph("experts", EXPERT_BRANCH_GRAPH, min_count=1, anchor_node="topk"),), + ("moe_combine",), + inputs=(_variadic_input("expert_outputs", "moe_combine", 0),), + repeated_subgraphs=(M.RepeatedSubgraph("experts", EXPERT_BRANCH_GRAPH, min_count=1),), attr_captures=(M.AttrCapture("num_experts", default=None, required=False), M.AttrCapture("num_experts_per_tok", "topk", "k", required=False), M.AttrCapture("normalize_routing", default=True, required=False), M.AttrCapture("epsilon", default=1e-6, required=False), M.AttrCapture("routed_scaling_factor", default=1.0, required=False)), - constraints=_note("Top-k indices must select the repeated expert branches.", "Routing probabilities must be the weights used for expert output combination."), + constraints={ + "moe_expert_branch_routing": {"repeated_subgraph": "experts", "combine_node": "moe_combine", "routing_weight_role": "routing_weight", "max_depth": 64}, + "moe_routing_weights_combine": {"repeated_subgraph": "experts", "combine_node": "moe_combine", "weighted_node": "expert_weighted", "routing_weight_role": "routing_weight", "max_depth": 64}, + }, +) + +LFM_GROUPED_MOE_GRAPH = _graph( + "lfm_grouped_moe", + "moe_grouped_combine", + ( + "moe_hidden_view", + "moe_router_weight_transpose", + "moe_router_logits", + "moe_router_sigmoid", + "moe_router_bias", + "moe_topk", + "moe_topk_indices", + "moe_routing_gather", + "moe_routing_sum", + "moe_routing_eps", + "moe_routing_norm", + "moe_routing_scale", + "moe_token_arange", + "moe_token_unsqueeze", + "moe_token_expand", + "moe_token_clone", + "moe_token_flat", + "moe_routing_flat", + "moe_topk_flat", + "moe_sort", + "moe_sort_indices", + "moe_unsort_empty", + "moe_unsort_arange", + "moe_unsort_index_put", + "moe_sorted_expert_index", + "moe_sorted_expert_copy", + "moe_histc", + "moe_offsets", + "moe_sorted_routing_index", + "moe_token_positions_index", + "moe_sorted_hidden_index", + "moe_gate_up_weight_transpose", + "moe_grouped_gate_up", + "moe_gate_up_split", + "moe_gate_tensor", + "moe_up_tensor", + "moe_gate_copy", + "moe_gate_sigmoid", + "moe_gate_activation", + "moe_gate_activation_copy", + "moe_expert_product", + "moe_down_weight_transpose", + "moe_grouped_down", + "moe_routing_unsqueeze", + "moe_grouped_weighted", + "moe_unsort_index", + "moe_output_view", + "moe_grouped_combine", + ), + edge_names=E.EDGE_GROUPS["lfm_grouped_moe"], + inputs=( + _input("hidden", "moe_hidden_view", 0), + _input("router_weight", "moe_router_weight_transpose", 0, allowed_value_kinds=(M.ValueKind.PARAMETER, M.ValueKind.BUFFER)), + _input("router_bias", "moe_router_bias", 1, allowed_value_kinds=(M.ValueKind.PARAMETER, M.ValueKind.BUFFER)), + _input("gate_up_weight", "moe_gate_up_weight_transpose", 0, allowed_value_kinds=(M.ValueKind.PARAMETER, M.ValueKind.BUFFER)), + _input("down_weight", "moe_down_weight_transpose", 0, allowed_value_kinds=(M.ValueKind.PARAMETER, M.ValueKind.BUFFER)), + ), + attr_captures=( + M.AttrCapture("num_experts", "moe_histc", "bins", required=False), + M.AttrCapture("num_experts_per_tok", "moe_topk", "k", required=True), + M.AttrCapture("epsilon", "moe_routing_eps", "other", default=1e-6, required=False), + M.AttrCapture("routed_scaling_factor", "moe_routing_scale", "other", default=1.0, required=False), + ), + constraints={ + "grouped_moe_structure": { + "topk_node": "moe_topk", + "topk_indices_node": "moe_topk_indices", + "router_bias_role": "router_bias", + "gate_up_weight_role": "gate_up_weight", + "down_weight_role": "down_weight", + "gate_up_grouped_node": "moe_grouped_gate_up", + "down_grouped_node": "moe_grouped_down", + "split_node": "moe_gate_up_split", + "offsets_node": "moe_offsets", + "combine_node": "moe_grouped_combine", + }, + }, ) LSTM_CELL_GRAPH = _graph( @@ -775,6 +854,7 @@ def _index_by_root_op( "kv_cache_append": KV_CACHE_APPEND_GRAPH, "attention_cached": ATTENTION_CACHED_GRAPH, "moe_layer_gated": MOE_GATED_GRAPH, + "lfm_grouped_moe": LFM_GROUPED_MOE_GRAPH, "lstm_cell": LSTM_CELL_GRAPH, "gated_deltanet_decode": DELTANET_DECODE_GRAPH, "gated_deltanet_prefill": DELTANET_PREFILL_GRAPH, @@ -829,6 +909,7 @@ def _index_by_root_op( "kv_cache_append": _definition("kv_cache_append", "kv_cache_append", KV_CACHE_APPEND_GRAPH, fusion_fields=("generic", "cache"), supported_inference_modes=("prefill_with_cache", "decode_with_cache")), "attention_cached": _definition("attention_cached", "attention_cached", ATTENTION_CACHED_GRAPH, fusion_fields=("generic", "attention", "cache"), supported_inference_modes=("decode_with_cache",)), "moe_layer_gated": _definition("moe_layer_gated", "moe_layer_gated", MOE_GATED_GRAPH, fusion_fields=("generic", "moe")), + "lfm_grouped_moe": _definition("lfm_grouped_moe", "moe_layer_gated", LFM_GROUPED_MOE_GRAPH, fusion_fields=("generic", "moe", "lfm_moe")), "lstm_cell": _definition("lstm_cell", "lstm_cell", LSTM_CELL_GRAPH, fusion_fields=("generic", "recurrent", "audio")), "gated_deltanet_decode": _definition("gated_deltanet_decode", "gated_deltanet_decode", DELTANET_DECODE_GRAPH, fusion_fields=("generic", "recurrent", "cache")), "gated_deltanet_prefill": _definition("gated_deltanet_prefill", "gated_deltanet_prefill", DELTANET_PREFILL_GRAPH, fusion_fields=("generic", "recurrent", "cache")), diff --git a/python/cactus/transpiler/Fusions/nodes.py b/python/cactus/transpiler/Fusions/nodes.py index 4297412cc..9aaf81bfe 100644 --- a/python/cactus/transpiler/Fusions/nodes.py +++ b/python/cactus/transpiler/Fusions/nodes.py @@ -32,12 +32,21 @@ "min": ("aten.min.dim", "aten.amin.default"), "softmax": ("aten.softmax.int", "aten._softmax.default"), "topk": ("aten.topk.default",), + "sort": ("aten.sort.default",), "view": ("aten.view.default", "aten.reshape.default", "aten.flatten.using_ints"), + "unsqueeze": ("aten.unsqueeze.default",), + "expand": ("aten.expand.default",), "transpose": ("aten.t.default", "aten.transpose.int", "aten.permute.default"), + "split": ("aten.split_with_sizes.default", "aten.chunk.default"), "slice": ("aten.slice.Tensor", "aten.select.int", "aten.narrow.default"), "index": ("aten.select.int", "aten.index.Tensor"), + "index_put": ("aten.index_put.default",), "cat": ("aten.cat.default",), "scatter": ("aten.scatter.src", "aten.scatter.value", "aten.scatter_reduce.two"), + "histc": ("aten.histc.default",), + "cumsum": ("aten.cumsum.default",), + "empty": ("aten.empty_permuted.default",), + "grouped_mm": ("transformers.grouped_mm_fallback.default",), "neg": ("aten.neg.default", "operator.neg", ""), "where": ("aten.where.self",), "eq": ("aten.eq.Tensor", "aten.eq.Scalar"), @@ -135,14 +144,22 @@ def _expert_node(name: str, op_key: str) -> M.FusionNode: "max": _node("max", "max"), "softmax": _node("softmax", "softmax"), "topk_direct": _node("topk_direct", "topk"), + "sort": _node("sort", "sort"), "view": _node("view", "view"), "reshape": _node("reshape", "view"), "flatten": _node("flatten", "view"), + "unsqueeze": _node("unsqueeze", "unsqueeze"), + "expand": _node("expand", "expand"), "transpose": _node("transpose", "transpose"), + "split": _node("split", "split"), "slice": _node("slice", "slice"), "index": _node("index", "index"), + "index_put": _node("index_put", "index_put"), "cat": _node("cat", "cat"), "gather": _node("gather", "embedding"), + "histc": _node("histc", "histc"), + "cumsum": _node("cumsum", "cumsum"), + "empty": _node("empty", "empty"), "not_equal": _node("not_equal", "ne"), "equal": _node("equal", "eq"), "embedding": _node("embedding", "embedding"), @@ -227,6 +244,54 @@ def _expert_node(name: str, op_key: str) -> M.FusionNode: "router_logits": _node("router_logits", "matmul"), "routing_probs": _node("routing_probs", "softmax"), "topk": _node("topk", "topk"), + "moe_hidden_view": _node("moe_hidden_view", "view"), + "moe_router_weight_transpose": _node("moe_router_weight_transpose", "transpose"), + "moe_router_logits": _node("moe_router_logits", "matmul"), + "moe_router_sigmoid": _node("moe_router_sigmoid", "sigmoid"), + "moe_router_bias": _node("moe_router_bias", "add"), + "moe_topk": _node("moe_topk", "topk"), + "moe_topk_indices": _node("moe_topk_indices", "getitem", attrs=(_attr("index", 1),)), + "moe_routing_gather": _node("moe_routing_gather", "embedding"), + "moe_routing_sum": _node("moe_routing_sum", "sum"), + "moe_routing_eps": _node("moe_routing_eps", "add"), + "moe_routing_norm": _node("moe_routing_norm", "div"), + "moe_routing_scale": _node("moe_routing_scale", "mul"), + "moe_token_arange": _node("moe_token_arange", "arange"), + "moe_token_unsqueeze": _node("moe_token_unsqueeze", "unsqueeze"), + "moe_token_expand": _node("moe_token_expand", "expand"), + "moe_token_clone": _node("moe_token_clone", "copy"), + "moe_token_flat": _node("moe_token_flat", "view"), + "moe_routing_flat": _node("moe_routing_flat", "view"), + "moe_topk_flat": _node("moe_topk_flat", "view"), + "moe_sort": _node("moe_sort", "sort"), + "moe_sort_indices": _node("moe_sort_indices", "getitem", attrs=(_attr("index", 1),)), + "moe_unsort_empty": _node("moe_unsort_empty", "empty"), + "moe_unsort_arange": _node("moe_unsort_arange", "arange"), + "moe_unsort_index_put": _node("moe_unsort_index_put", "index_put"), + "moe_sorted_expert_index": _node("moe_sorted_expert_index", "index"), + "moe_sorted_expert_copy": _node("moe_sorted_expert_copy", "copy"), + "moe_histc": _node("moe_histc", "histc"), + "moe_offsets": _node("moe_offsets", "cumsum"), + "moe_sorted_routing_index": _node("moe_sorted_routing_index", "index"), + "moe_token_positions_index": _node("moe_token_positions_index", "index"), + "moe_sorted_hidden_index": _node("moe_sorted_hidden_index", "index"), + "moe_gate_up_weight_transpose": _node("moe_gate_up_weight_transpose", "transpose"), + "moe_grouped_gate_up": _node("moe_grouped_gate_up", "grouped_mm"), + "moe_gate_up_split": _node("moe_gate_up_split", "split"), + "moe_gate_tensor": _node("moe_gate_tensor", "getitem", attrs=(_attr("index", 0),)), + "moe_up_tensor": _node("moe_up_tensor", "getitem", attrs=(_attr("index", 1),)), + "moe_gate_copy": _node("moe_gate_copy", "copy"), + "moe_gate_sigmoid": _node("moe_gate_sigmoid", "sigmoid"), + "moe_gate_activation": _node("moe_gate_activation", "mul"), + "moe_gate_activation_copy": _node("moe_gate_activation_copy", "copy"), + "moe_expert_product": _node("moe_expert_product", "mul"), + "moe_down_weight_transpose": _node("moe_down_weight_transpose", "transpose"), + "moe_grouped_down": _node("moe_grouped_down", "grouped_mm"), + "moe_routing_unsqueeze": _node("moe_routing_unsqueeze", "unsqueeze"), + "moe_grouped_weighted": _node("moe_grouped_weighted", "mul"), + "moe_unsort_index": _node("moe_unsort_index", "index"), + "moe_output_view": _node("moe_output_view", "view"), + "moe_grouped_combine": _node("moe_grouped_combine", "sum"), "expert_gate": _expert_node("expert_gate", "matmul"), "expert_up": _expert_node("expert_up", "matmul"), "expert_activation": _expert_node("expert_activation", "silu"), @@ -349,6 +414,54 @@ def _expert_node(name: str, op_key: str) -> M.FusionNode: "router_logits", "routing_probs", "topk", + "moe_hidden_view", + "moe_router_weight_transpose", + "moe_router_logits", + "moe_router_sigmoid", + "moe_router_bias", + "moe_topk", + "moe_topk_indices", + "moe_routing_gather", + "moe_routing_sum", + "moe_routing_eps", + "moe_routing_norm", + "moe_routing_scale", + "moe_token_arange", + "moe_token_unsqueeze", + "moe_token_expand", + "moe_token_clone", + "moe_token_flat", + "moe_routing_flat", + "moe_topk_flat", + "moe_sort", + "moe_sort_indices", + "moe_unsort_empty", + "moe_unsort_arange", + "moe_unsort_index_put", + "moe_sorted_expert_index", + "moe_sorted_expert_copy", + "moe_histc", + "moe_offsets", + "moe_sorted_routing_index", + "moe_token_positions_index", + "moe_sorted_hidden_index", + "moe_gate_up_weight_transpose", + "moe_grouped_gate_up", + "moe_gate_up_split", + "moe_gate_tensor", + "moe_up_tensor", + "moe_gate_copy", + "moe_gate_sigmoid", + "moe_gate_activation", + "moe_gate_activation_copy", + "moe_expert_product", + "moe_down_weight_transpose", + "moe_grouped_down", + "moe_routing_unsqueeze", + "moe_grouped_weighted", + "moe_unsort_index", + "moe_output_view", + "moe_grouped_combine", "expert_gate", "expert_up", "expert_activation", From ae567c624a98ea3ac13dd2801bb45156ec240edf Mon Sep 17 00:00:00 2001 From: Parkirat Date: Tue, 28 Jul 2026 16:07:07 -0400 Subject: [PATCH 56/84] Added LFM MOE model profile --- .../transpiler/ModelProfiles/profiles.py | 24 +++++++++++++++++++ 1 file changed, 24 insertions(+) diff --git a/python/cactus/transpiler/ModelProfiles/profiles.py b/python/cactus/transpiler/ModelProfiles/profiles.py index 3809af2fb..0bf2ffa83 100644 --- a/python/cactus/transpiler/ModelProfiles/profiles.py +++ b/python/cactus/transpiler/ModelProfiles/profiles.py @@ -126,10 +126,34 @@ load_strategy="causal_lm", ) +LFM_MOE_PROFILE = ModelProfile( + model_profiles="lfm_moe", + components=TEXT_COMPONENTS, + inference_type=QWEN2_5_INFERENCE_PATTERNS, + cache_type=("attention_kv", "conv_state"), + cache_policy=(), + files=("config.json", "generation_config.json", "tokenizer.json", "tokenizer_config.json"), + fusion_fields=( + "generic", + "embedding", + "lfm_attention", + "lfm_mlp", + "lfm_moe", + "moe", + "conv", + "linear", + ), + supported_modalties=("text",), + input_strategy="processor", + export_patches=("transformers_moe_grouped_mm_fallback",), + load_strategy="causal_lm", +) + MODEL_ID_MAP = { "google/gemma-4-E2B": GEMMA4_E2B_PROFILE, "openai/whisper-tiny": WHISPER_PROFILE, "nvidia/parakeet-tdt-0.6b-v3": PARAKEET_PROFILE, "LiquidAI/LFM2-VL-3B": LFM_VLM_PROFILE, "Qwen/Qwen2.5-0.5B": QWEN2_5_0_5B_PROFILE, + "LiquidAI/LFM2.5-8B-A1B": LFM_MOE_PROFILE, } From 0cdd3d6708d83fb166d570bc0a4d3874579f1f6b Mon Sep 17 00:00:00 2001 From: Parkirat Date: Tue, 28 Jul 2026 16:09:43 -0400 Subject: [PATCH 57/84] Additional matchers for lfm MOE --- python/cactus/transpiler/IR/constants.py | 10 +- python/cactus/transpiler/IR/extra_matchers.py | 747 +++++++++++++++++- python/cactus/transpiler/IR/match_utils.py | 8 +- 3 files changed, 756 insertions(+), 9 deletions(-) diff --git a/python/cactus/transpiler/IR/constants.py b/python/cactus/transpiler/IR/constants.py index 19b8dcb88..29a15c903 100644 --- a/python/cactus/transpiler/IR/constants.py +++ b/python/cactus/transpiler/IR/constants.py @@ -39,9 +39,12 @@ "aten.softmax.int": {1: "dim", 2: "dtype"}, "aten._softmax.default": {1: "dim", 2: "half_to_float"}, "aten.topk.default": {1: "k", 2: "dim", 3: "largest", 4: "sorted"}, + "aten.sort.default": {1: "dim", 2: "descending"}, "aten.view.default": {1: "shape"}, "aten.reshape.default": {1: "shape"}, "aten.flatten.using_ints": {1: "start_dim", 2: "end_dim"}, + "aten.unsqueeze.default": {1: "dim"}, + "aten.expand.default": {1: "shape", 2: "implicit"}, "aten.t.default": {}, "aten.transpose.int": {1: "dim0", 2: "dim1"}, "aten.permute.default": {1: "permutation"}, @@ -49,12 +52,17 @@ "aten.select.int": {1: "dim", 2: "index"}, "aten.narrow.default": {1: "dim", 2: "start", 3: "length"}, "aten.index.Tensor": {1: "indices"}, + "aten.index_put.default": {3: "accumulate"}, "aten.cat.default": {1: "dim"}, "aten.embedding.default": {2: "padding_idx", 3: "scale_grad_by_freq", 4: "sparse"}, "aten.gather.default": {1: "dim", 2: "index"}, + "aten.split_with_sizes.default": {1: "split_sizes", 2: "dim"}, "aten.scatter.src": {1: "dim", 2: "index"}, "aten.scatter.value": {1: "dim", 2: "index", 3: "value"}, "aten.scatter_reduce.two": {1: "dim", 2: "index", 4: "reduce", 5: "include_self"}, + "aten.histc.default": {1: "bins", 2: "min", 3: "max"}, + "aten.cumsum.default": {1: "dim", 2: "dtype"}, + "aten.empty_permuted.default": {0: "size", 1: "stride"}, "aten.clamp.default": {1: "min", 2: "max"}, "aten.clamp.Tensor": {1: "min", 2: "max"}, "aten.gelu.default": {1: "approximate"}, @@ -96,4 +104,4 @@ }, "aten.fft_rfft.default": {1: "n", 2: "dim", 3: "norm"}, "aten.fft_irfft.default": {1: "n", 2: "dim", 3: "norm"}, -} \ No newline at end of file +} diff --git a/python/cactus/transpiler/IR/extra_matchers.py b/python/cactus/transpiler/IR/extra_matchers.py index 65bfa1ab5..ae9a7b524 100644 --- a/python/cactus/transpiler/IR/extra_matchers.py +++ b/python/cactus/transpiler/IR/extra_matchers.py @@ -3,7 +3,7 @@ from . import models, match_utils from ..Fusions import models as FModels - +#TODO: Clean up this file ExtraMatcher = Callable[[models.Node, models.Graph, FModels.FusionGraph, dict[str, models.Node], dict[str, Any]], bool] @@ -37,6 +37,14 @@ ############################################ Extra Matcher Utils!!!!! ############################################ def get_bound_parent(bindings: dict[str, models.Node], node_name: str, parent_index: int) -> models.Node | None: + """ + Returns a real parent node from a bound synthetic node/input position. + + `bindings` maps synthetic fusion-node names to real IR nodes. This helper + first finds the real node bound to `node_name`, then asks match_utils for + the parent at `parent_index`. It returns None when the synthetic node was + not bound or the requested parent index does not exist. + """ node = bindings.get(node_name) if node is None: @@ -46,6 +54,14 @@ def get_bound_parent(bindings: dict[str, models.Node], node_name: str, parent_in def get_node_dim(node: models.Node, dim_index: int) -> Any: + """ + Returns one tensor dimension from a real IR node's output shape. + + This reads `tensor_output_meta["shape"]`, supports negative indices in the + Python style, and returns `match_utils.MISSING` instead of raising when the + shape is unavailable or the dimension index is invalid. Matchers use that + missing sentinel so they can either fail closed or honor `allow_missing`. + """ shape = match_utils.get_tensor_shape(node) if not shape: @@ -61,10 +77,26 @@ def get_node_dim(node: models.Node, dim_index: int) -> Any: def classify_linear_rhs(node: models.Node) -> tuple[str, models.Node] | None: + """ + Classifies the right-hand side of a linear/matmul as direct or transposed. + + The node passed here is the actual RHS input to the matmul. The helper + delegates to the recursive implementation with an empty `seen` set so we + can safely walk through layout-only wrapper nodes without infinite loops. + """ return classify_linear_rhs_(node, set()) def classify_linear_rhs_(node: models.Node, seen: set[int]) -> tuple[str, models.Node] | None: + """ + Walks backward through simple RHS wrapper ops to find the source weight. + + It accepts either a direct parameter/buffer/lifted constant weight, or a + valid transpose/permute/t of such a weight. Copy-like wrappers such as + `_to_copy` and `contiguous` are skipped recursively. The return value is + `(layout_kind, source_weight)`, where layout_kind is `direct` or + `transposed`; None means the RHS is not safely recognizable as a weight. + """ if id(node) in seen: return None @@ -84,6 +116,14 @@ def classify_linear_rhs_(node: models.Node, seen: set[int]) -> tuple[str, models def linear_transpose_is_valid(node: models.Node, source_weight: models.Node) -> bool: + """ + Checks whether a transpose-like op really swaps a 2D linear weight. + + Cactus needs to know whether the RHS has already been converted from + `[out_dim, in_dim]` to `[in_dim, out_dim]`. This accepts `t`, `transpose` + over dimensions 0/1, and `permute([1, 0])`, but rejects non-2D weights or + arbitrary permutations that would not be a standard linear transpose. + """ source_shape = match_utils.get_tensor_shape(source_weight) if len(source_shape) != 2: @@ -102,6 +142,15 @@ def linear_transpose_is_valid(node: models.Node, source_weight: models.Node) -> def linear_shapes_match(input_node: models.Node, rhs_node: models.Node, output_node: models.Node, source_weight: models.Node, layout_kind: str, spec: dict[str, Any]) -> bool: + """ + Verifies the dimensions around a linear/matmul are internally consistent. + + The RHS node is the tensor consumed by the real matmul, so its last two dims + should be `[input_dim, output_dim]`. If that RHS came from a transposed + original weight, the original source weight should be `[output_dim, + input_dim]`. Missing shapes are handled through the constraint spec's + `allow_missing` flag so callers can choose conservative or permissive mode. + """ input_shape = match_utils.get_tensor_shape(input_node) rhs_shape = match_utils.get_tensor_shape(rhs_node) output_shape = match_utils.get_tensor_shape(output_node) @@ -134,6 +183,13 @@ def linear_shapes_match(input_node: models.Node, rhs_node: models.Node, output_n def compatible_shapes(left_shape: list[Any], right_shape: list[Any], spec: dict[str, Any]) -> bool: + """ + Checks whether two tensor shapes are equal or broadcast-compatible. + + This is used for things like RoPE cos/sin tables, where shapes may be + exactly identical or broadcastable across singleton dimensions. If either + shape is missing, the result is controlled by `allow_missing` in the spec. + """ if not left_shape or not right_shape: return bool(spec.get("allow_missing", False)) @@ -145,6 +201,14 @@ def compatible_shapes(left_shape: list[Any], right_shape: list[Any], spec: dict[ def compatible_dtypes(left_node: models.Node, right_node: models.Node, spec: dict[str, Any]) -> bool: + """ + Checks whether two nodes have matching output dtypes when required. + + Most table-like values in one fusion should have the same dtype. The check + can be disabled with `require_same_dtype=False`; otherwise missing dtype + metadata follows `allow_missing`, and available dtype strings are compared + with the shared `values_equal` normalization behavior. + """ if not spec.get("require_same_dtype", True): return True @@ -158,6 +222,14 @@ def compatible_dtypes(left_node: models.Node, right_node: models.Node, spec: dic def compatible_trig_sources(cos_node: models.Node, sin_node: models.Node, spec: dict[str, Any]) -> bool: + """ + Checks whether RoPE cos and sin tables appear to come from the same angles. + + The helper walks backward from the cos/sin input nodes until it finds + `aten.cos` and `aten.sin`, then compares the node feeding those trig ops. + If no trig source is visible, the matcher can either allow that partial + evidence or fail closed with `require_trig_sources=True`. + """ cos_sources = collect_trig_angle_sources(cos_node, COS_TARGETS, int(spec.get("max_trig_depth", 8))) sin_sources = collect_trig_angle_sources(sin_node, SIN_TARGETS, int(spec.get("max_trig_depth", 8))) @@ -168,10 +240,25 @@ def compatible_trig_sources(cos_node: models.Node, sin_node: models.Node, spec: def collect_trig_angle_sources(node: models.Node, trig_targets: set[str], max_depth: int) -> frozenset[str]: + """ + Finds the angle-producing parent nodes for visible trig ops upstream. + + Starting from `node`, this walks backward through parents up to `max_depth`. + When it reaches a target such as `aten.cos.default` or `aten.sin.default`, + it records that trig node's first parent as the angle source. The result is + a frozenset so cos and sin paths can be compared independent of ordering. + """ sources: set[str] = set() seen: set[int] = set() def visit(current: models.Node, depth: int) -> None: + """ + Recursively walks parents while avoiding cycles and depth blowups. + + The exported graph should be a DAG, but the `seen` guard makes this + helper robust if malformed graph data or future graph rewrites ever + introduce repeated references. + """ if depth > max_depth or id(current) in seen: return @@ -189,6 +276,14 @@ def visit(current: models.Node, depth: int) -> None: def rope_tables_match_input(x_node: models.Node, cos_shape: list[Any], sin_shape: list[Any], spec: dict[str, Any]) -> bool: + """ + Checks whether RoPE cos/sin table shapes can apply to the rotated input. + + RoPE tables often have fewer broadcast dimensions than the activation, and + the final table dimension can represent either the full hidden/head dim or + half of it depending on the rotate-half implementation. This delegates the + actual table-vs-input rule to `rope_table_shape_matches_x` for both tables. + """ x_shape = match_utils.get_tensor_shape(x_node) if not x_shape: @@ -198,6 +293,14 @@ def rope_tables_match_input(x_node: models.Node, cos_shape: list[Any], sin_shape def rope_table_shape_matches_x(table_shape: list[Any], x_shape: list[Any]) -> bool: + """ + Checks one RoPE table shape against the activation being rotated. + + The table is accepted if it is broadcastable to the activation shape. If + not, we allow the common rotate-half form where the table's last dimension + is half of the activation's last dimension. Symbolic or missing dimensions + are intentionally conservative here. + """ if not table_shape: return False @@ -214,6 +317,14 @@ def rope_table_shape_matches_x(table_shape: list[Any], x_shape: list[Any]) -> bo def gate_shape_matches(gate_node: models.Node, gate_count: int, spec: dict[str, Any]) -> bool: + """ + Checks whether an LSTM combined-gate tensor can hold `gate_count` gates. + + LSTM gate projections commonly produce a last dimension equal to + `4 * hidden_size`. We do not know hidden_size directly here, so this checks + that the final dimension is divisible by the expected gate count. Symbolic + dimensions can be allowed through `allow_symbolic`. + """ gate_shape = match_utils.get_tensor_shape(gate_node) if not gate_shape: @@ -228,6 +339,14 @@ def gate_shape_matches(gate_node: models.Node, gate_count: int, spec: dict[str, def has_explicit_gate_split(gate_node: models.Node, gate_count: int, max_depth: int) -> bool: + """ + Looks for a visible split/chunk/unbind that separates combined LSTM gates. + + Shape divisibility alone is not enough to prove LSTM semantics. This walks + forward from the combined gate node and requires a split-like op whose attr + or fanout indicates the expected number of gates. If no such split is + visible, the LSTM fusion fails closed. + """ for node in walk_children(gate_node, max_depth): if node.target not in SPLIT_TARGETS: continue @@ -239,6 +358,14 @@ def has_explicit_gate_split(gate_node: models.Node, gate_count: int, max_depth: def split_node_has_gate_count(node: models.Node, gate_count: int) -> bool: + """ + Determines whether a split-like node creates the expected number of gates. + + Different export paths encode splits differently: `unbind` is inferred from + child count, `chunk` stores the chunk count as positional attr `arg_1`, and + `split_with_sizes` stores the explicit split sizes. If those attrs are not + available, child count is used as the fallback evidence. + """ if node.target == "aten.unbind.int": return len(node.children) == gate_count @@ -254,10 +381,24 @@ def split_node_has_gate_count(node: models.Node, gate_count: int) -> bool: def walk_children(node: models.Node, max_depth: int) -> tuple[models.Node, ...]: + """ + Returns all reachable child nodes within a bounded forward traversal. + + Extra matchers use this when they need local forward evidence, such as + proving an LSTM gate tensor is split later. The traversal tracks object ids + to avoid revisiting shared children in a DAG and to remain safe if graph + rewrites produce unexpected cycles. + """ found: list[models.Node] = [] seen: set[int] = set() def visit(current: models.Node, depth: int) -> None: + """ + Recursively collects children while respecting the maximum depth. + + The outer `walk_children` owns the result list and visited set; this + nested helper only handles traversal mechanics. + """ if depth > max_depth: return @@ -274,6 +415,13 @@ def visit(current: models.Node, depth: int) -> None: def get_node_dtype(node: models.Node) -> Any: + """ + Reads a node's output dtype from tensor metadata. + + The exported JSON stores tensor metadata as dictionaries. If metadata is + missing or has some non-dict form, the function returns the shared MISSING + sentinel so callers can decide whether to fail or allow missing metadata. + """ if not isinstance(node.tensor_output_meta, dict): return match_utils.MISSING @@ -281,13 +429,367 @@ def get_node_dtype(node: models.Node) -> Any: def normalize_dim(dim: Any, rank: int) -> Any: + """ + Converts a possibly negative dimension index into a non-negative one. + + This mirrors Python indexing for known integer dimensions. Non-integer + values are returned unchanged so symbolic or malformed values do not crash + the matcher; they simply fail later if an exact comparison is required. + """ if not isinstance(dim, int): return dim return dim + rank if dim < 0 else dim +def get_repeated_subgraph(fusion: FModels.FusionGraph, name: str) -> FModels.RepeatedSubgraph | None: + """ + Fetches a named repeated subgraph declaration from a fusion graph. + + MoE fusions store expert branches as repeated subgraphs. This helper finds + the declaration by name so the MoE-specific matchers can reuse its pattern + and constraints when collecting real expert-branch bindings. + """ + for subgraph in fusion.repeated_subgraphs: + if subgraph.name == name: + return subgraph + + return None + + +def collect_repeated_subgraph_bindings(graph: models.Graph, subgraph: FModels.RepeatedSubgraph) -> tuple[dict[str, models.Node], ...]: + """ + Finds every real subgraph binding that matches a repeated pattern. + + This scans each node as a possible root for `subgraph.graph`, runs the same + graph binder used by normal fusion matching, then validates the full set of + fusion matchers against that candidate. The result is not just a count: it + returns each binding dictionary so later matchers can inspect which real + nodes were bound to expert_gate, expert_weighted, routing_weight, etc. + """ + from . import match + + matches: list[dict[str, models.Node]] = [] + + for candidate in graph.nodes: + candidate_bindings = match_utils.bind_fusion_graph(candidate, subgraph.graph, match.match_nodes) + + if candidate_bindings is None: + continue + + if not match.match_fusion_bindings(candidate, graph, subgraph.graph, candidate_bindings): + continue + + matches.append(candidate_bindings) + + return tuple(matches) + + +def relevant_moe_branch_bindings(graph: models.Graph, fusion: FModels.FusionGraph, bindings: dict[str, models.Node], spec: dict[str, Any]) -> tuple[dict[str, models.Node], ...]: + """ + Filters repeated expert branches down to ones driven by MoE routing. + + The first step collects every expert-like branch in the whole graph. Then + this keeps only branches whose declared `routing_weight` input is reachable + from a top-k node or that top-k node's probability input. This is the key + safety check that prevents us from fusing arbitrary MLP-looking branches as + MoE experts when they are not actually selected by the router. + """ + subgraph = get_repeated_subgraph(fusion, spec["repeated_subgraph"]) + + if subgraph is None: + return () + + topk_node = bindings.get(spec.get("topk_node", "topk")) + combine_node = bindings.get(spec.get("combine_node", "moe_combine")) + topk_nodes = (topk_node,) if topk_node is not None else find_reachable_topk_nodes(graph, combine_node, spec) + + if not topk_nodes: + return () + + max_depth = int(spec.get("max_depth", 16)) + branch_bindings = collect_repeated_subgraph_bindings(graph, subgraph) + relevant = [] + + for branch_binding in branch_bindings: + routing_weight = match_utils.get_first_input_by_role(subgraph.graph, branch_binding, spec["routing_weight_role"]) + + if routing_weight is None: + continue + + routing_sources = tuple(source for topk in topk_nodes for source in (topk, first_parent(topk)) if source is not None) + + if not has_any_path(routing_sources, routing_weight, max_depth): + continue + + relevant.append(branch_binding) + + return tuple(relevant) + + +def required_moe_branch_count(subgraph: FModels.RepeatedSubgraph, topk_node: models.Node | None, spec: dict[str, Any]) -> int: + """ + Determines how many routed expert branches must be proven. + + The spec may override this directly with `min_count`. Otherwise, if a top-k + node exposes its `k` attr, that is the strongest graph-local value because + it tells us how many experts each token selects. If neither exists, the + repeated subgraph's own `min_count` is used. + """ + if "min_count" in spec: + return int(spec["min_count"]) + + if topk_node is not None and spec.get("min_count_from_topk", True): + k = topk_node.attrs.get("k", match_utils.MISSING) + + if isinstance(k, int): + return k + + return subgraph.min_count + + +def get_attr(node: models.Node, *names: str) -> Any: + """ + Returns the first available attr from a node. + + Some attrs are normalized with friendly names such as `split_sizes`, while + older/generated JSON may only expose fallback names such as `arg_1`. This + lets extra matchers accept either form without duplicating logic. + """ + for name in names: + value = node.attrs.get(name, match_utils.MISSING) + + if value is not match_utils.MISSING: + return value + + return match_utils.MISSING + + +def get_known_int(value: Any) -> int | None: + """ + Converts plain integer-like values into int, otherwise returns None. + + Shape metadata can contain symbolic values, so matchers use this helper + when a semantic check is only safe for known concrete dimensions. + """ + return value if isinstance(value, int) and not isinstance(value, bool) else None + + +def get_shape_dim(node: models.Node | None, dim_index: int) -> int | None: + """ + Returns a concrete int dimension from a node shape when available. + + The helper supports negative indexing and intentionally returns None for + missing/symbolic dimensions so callers can decide whether to fail or skip. + """ + if node is None: + return None + + return get_known_int(get_node_dim(node, dim_index)) + + +def get_num_experts_from_inputs(fusion: FModels.FusionGraph, bindings: dict[str, models.Node], spec: dict[str, Any]) -> int | None: + """ + Infers the MoE expert count from weight/bias shapes. + + For LFM grouped MoE, router bias is `[num_experts]`, gate/up expert weight + is `[num_experts, 2 * intermediate, hidden]`, and down expert weight is + `[num_experts, hidden, intermediate]`. Matching these first dimensions is + a strong safety check that the grouped matmuls are really expert weights. + """ + candidates: list[int] = [] + + for role in (spec["gate_up_weight_role"], spec["down_weight_role"], spec.get("router_bias_role")): + if role is None: + continue + + input_node = match_utils.get_first_input_by_role(fusion, bindings, role) + shape = match_utils.get_tensor_shape(input_node) + + if not shape: + continue + + dim = get_known_int(shape[0]) + + if dim is not None: + candidates.append(dim) + + if not candidates: + return None + + if len(set(candidates)) != 1: + return None + + return candidates[0] + + +def split_sizes_are_valid(split_node: models.Node, gate_up_node: models.Node, down_weight: models.Node | None, spec: dict[str, Any]) -> bool: + """ + Checks the grouped gate/up projection split. + + LFM projects gate and up together, then splits the final dimension into two + equal chunks. The second chunk should match the down-projection expert + input dimension when that shape is available. + """ + split_sizes = get_attr(split_node, "split_sizes", "arg_1") + + if not isinstance(split_sizes, (list, tuple)) or len(split_sizes) != 2: + return bool(spec.get("allow_missing", False)) + + left_size = get_known_int(split_sizes[0]) + right_size = get_known_int(split_sizes[1]) + + if left_size is None or right_size is None: + return bool(spec.get("allow_symbolic", True)) + + if left_size != right_size: + return False + + gate_up_output_dim = get_shape_dim(gate_up_node, -1) + + if gate_up_output_dim is not None and gate_up_output_dim != left_size + right_size: + return False + + down_input_dim = get_shape_dim(down_weight, -1) + + if down_input_dim is not None and down_input_dim != right_size: + return False + + return True + + +def grouped_moe_rows_match_topk(hidden_node: models.Node | None, gate_up_node: models.Node, topk_node: models.Node, spec: dict[str, Any]) -> bool: + """ + Checks that routed grouped-mm rows equal token count times top-k. + + This links the routing decision to the grouped expert matmul size. Missing + or symbolic dimensions can be allowed by `allow_missing` because shape + metadata may be incomplete in some exported graphs. + """ + token_count = get_shape_dim(hidden_node, 0) + grouped_rows = get_shape_dim(gate_up_node, 0) + topk = get_known_int(get_attr(topk_node, "k", "arg_1")) + + if token_count is None or grouped_rows is None or topk is None: + return bool(spec.get("allow_missing", False)) + + return grouped_rows == token_count * topk + + +def grouped_moe_weight_permute_is_valid(node: models.Node | None) -> bool: + """ + Checks the expert weight permutation used before grouped_mm. + + The LFM export stores expert weights with expert dimension first and then + permutes the last two dims to feed grouped_mm. That appears as + `permute([0, 2, 1])`. + """ + if node is None: + return False + + return match_utils.values_equal(get_attr(node, "permutation", "arg_1"), [0, 2, 1]) + + +def find_reachable_topk_nodes(graph: models.Graph, combine_node: models.Node | None, spec: dict[str, Any]) -> tuple[models.Node, ...]: + """ + Finds top-k router nodes that can flow into the MoE combine output. + + Some MoE fusion graphs do not bind top-k as a direct synthetic node because + the exported routing path can contain indexing/scatter/gather wrappers. + This helper searches the real graph for `aten.topk.default` nodes that have + a forward path to the matched combine node within the configured depth. + """ + if combine_node is None: + return () + + max_depth = int(spec.get("max_depth", 16)) + return tuple(node for node in graph.nodes if node.target == "aten.topk.default" and has_path(node, combine_node, max_depth)) + + +def first_parent(node: models.Node) -> models.Node | None: + """ + Returns the first parent of a node, if it has one. + + For top-k, parent 0 is usually the routing probabilities tensor. The MoE + routing matcher treats both the top-k node and its first parent as valid + routing sources because exported graphs may use either values or indices + through different wrapper chains. + """ + return node.parents[0] if node.parents else None + + +def branch_output_reaches_combine(branch_binding: dict[str, models.Node], subgraph: FModels.RepeatedSubgraph, combine_node: models.Node, spec: dict[str, Any]) -> bool: + """ + Checks whether one matched expert branch contributes to the combine node. + + The branch's exposed weighted output should eventually flow into the outer + MoE combine/sum/scatter-add node. This proves that a routed expert branch is + actually part of the candidate MoE output instead of being an unrelated + branch elsewhere in the graph. + """ + weighted_node_name = spec.get("weighted_node", subgraph.graph.root) + weighted_node = branch_binding.get(weighted_node_name) + + if weighted_node is None: + return False + + return has_path(weighted_node, combine_node, int(spec.get("max_depth", 16))) + + +def has_any_path(sources: tuple[models.Node | None, ...], dest: models.Node, max_depth: int) -> bool: + """ + Returns True if any non-None source has a forward path to `dest`. + + This is a small convenience wrapper around `has_path` for cases where a + matcher has multiple plausible routing sources, such as top-k itself and + the routing-probability parent of top-k. + """ + for source in sources: + if source is not None and has_path(source, dest, max_depth): + return True + + return False + + +def has_path(source: models.Node, dest: models.Node, max_depth: int) -> bool: + """ + Performs a bounded forward reachability check in the IR DAG. + + The function starts at `source`, walks through `children`, and succeeds if + it reaches `dest` within `max_depth`. The visited set prevents repeated + work on shared DAG nodes and protects against accidental cycles. + """ + if source is dest: + return True + + seen: set[int] = set() + + def visit(node: models.Node, depth: int) -> bool: + if depth > max_depth or id(node) in seen: + return False + + seen.add(id(node)) + + for child in node.children: + if child is dest: + return True + + if visit(child, depth + 1): + return True + + return False + + return visit(source, 0) + + def shapes_are_broadcastable(left_shape: list[Any], right_shape: list[Any]) -> bool: + """ + Implements standard trailing-dimension broadcasting for two shapes. + + It compares dimensions from right to left. Each pair is compatible if the + dims are equal or either side is 1. Extra leading dimensions on the longer + shape are allowed, matching normal tensor broadcasting rules. + """ for left_dim, right_dim in zip(reversed(left_shape), reversed(right_shape)): if match_utils.values_equal(left_dim, 1) or match_utils.values_equal(right_dim, 1): continue @@ -299,6 +801,14 @@ def shapes_are_broadcastable(left_shape: list[Any], right_shape: list[Any]) -> b def get_layer_key(node: models.Node) -> str | None: + """ + Extracts a coarse layer identifier from module stack, node name, or target. + + Cached attention checks need to know whether different matched nodes belong + to the same decoder layer. This tries the richer module stack first, then + falls back to name/target strings, and returns keys like `layers.0` when it + can identify a layer index. + """ for value in (module_stack_text(node), node.name, node.target): layer_key = first_layer_key(value) @@ -309,10 +819,25 @@ def get_layer_key(node: models.Node) -> str | None: def module_stack_text(node: models.Node) -> str: + """ + Converts module_stack metadata into searchable text. + + Exported module stacks may be None or structured objects/lists. The matcher + only needs a best-effort string to search for layer-like path fragments, so + None becomes an empty string and everything else uses repr. + """ return "" if node.module_stack is None else repr(node.module_stack) def first_layer_key(value: str) -> str | None: + """ + Finds the first layer/block path fragment in a string. + + The helper normalizes `/` and `_` into `.` so names like + `model_layers_12_attn` can be treated similarly to + `model.layers.12.attn`. It then looks for common containers such as + `layers`, `layer`, `blocks`, or `h` followed by a numeric index. + """ parts = value.replace("/", ".").replace("_", ".").split(".") for index, part in enumerate(parts[:-1]): @@ -325,6 +850,14 @@ def first_layer_key(value: str) -> str | None: ############################################ Extra Matchers!!!!! ############################################ def match_extra_constraints(source: models.Node, graph: models.Graph, fusion: FModels.FusionGraph, bindings: dict[str, models.Node]) -> bool: + """ + Dispatches every structured FusionGraph constraint to its registered matcher. + + `fusion.constraints` is a dict keyed by constraint name. This function looks + up each name in `EXTRA_MATCHERS`, normalizes each raw spec into one or more + spec dictionaries, and requires every matcher invocation to return True. + Unknown constraint names or malformed specs fail closed. + """ for constraint_name, raw_spec in fusion.constraints.items(): matcher = EXTRA_MATCHERS.get(constraint_name) @@ -344,6 +877,13 @@ def match_extra_constraints(source: models.Node, graph: models.Graph, fusion: FM def normalize_constraint_specs(raw_spec: Any) -> tuple[dict[str, Any], ...] | None: + """ + Normalizes one constraint value into a tuple of spec dictionaries. + + A constraint key can map to a single dict or a list/tuple of dicts. The + latter lets one matcher type be applied multiple times without duplicating + registry keys. Non-dict specs return None so the caller can fail safely. + """ if isinstance(raw_spec, dict): return (raw_spec,) @@ -353,11 +893,15 @@ def normalize_constraint_specs(raw_spec: Any) -> tuple[dict[str, Any], ...] | No return None -def match_note(source: models.Node, graph: models.Graph, fusion: FModels.FusionGraph, bindings: dict[str, models.Node], spec: dict[str, Any]) -> bool: - return True - - def match_linear_weight_layout(source: models.Node, graph: models.Graph, fusion: FModels.FusionGraph, bindings: dict[str, models.Node], spec: dict[str, Any]) -> bool: + """ + Verifies that a linear/matmul RHS is a compatible model weight layout. + + The matcher resolves the fusion's input, weight, and output nodes by role. + It then classifies the weight path as direct or transposed and checks that + dimensions line up with the matmul output. This is what prevents a random + activation matrix from being treated as a Cactus linear weight. + """ input_node = match_utils.get_first_input_by_role(fusion, bindings, spec["input_role"]) weight_node = match_utils.get_first_input_by_role(fusion, bindings, spec["weight_role"]) output_node = bindings.get(spec["output_node"]) @@ -380,6 +924,13 @@ def match_linear_weight_layout(source: models.Node, graph: models.Graph, fusion: def match_rope_tables_compatible(source: models.Node, graph: models.Graph, fusion: FModels.FusionGraph, bindings: dict[str, models.Node], spec: dict[str, Any]) -> bool: + """ + Verifies that RoPE cos/sin inputs look like compatible rotation tables. + + It gets the matched x/cos/sin values by role, checks cos/sin shape and dtype + compatibility, optionally proves they came from matching angle sources, and + then checks that both tables can broadcast or half-dim match the RoPE input. + """ x_node = match_utils.get_first_input_by_role(fusion, bindings, spec["x_role"]) cos_node = match_utils.get_first_input_by_role(fusion, bindings, spec["cos_role"]) sin_node = match_utils.get_first_input_by_role(fusion, bindings, spec["sin_role"]) @@ -406,6 +957,14 @@ def match_rope_tables_compatible(source: models.Node, graph: models.Graph, fusio def match_lstm_gate_structure(source: models.Node, graph: models.Graph, fusion: FModels.FusionGraph, bindings: dict[str, models.Node], spec: dict[str, Any]) -> bool: + """ + Verifies that an LSTM gate node has real four-gate structure evidence. + + The matcher first checks that the combined gate tensor's final dimension is + divisible by the expected gate count. By default it then requires an + explicit downstream split/chunk/unbind into that many gates, so LSTM fusion + fails closed when the graph only vaguely resembles an LSTM. + """ gate_node = bindings.get(spec["gate_node"]) if gate_node is None: @@ -422,7 +981,142 @@ def match_lstm_gate_structure(source: models.Node, graph: models.Graph, fusion: return True +def match_moe_expert_branch_routing(source: models.Node, graph: models.Graph, fusion: FModels.FusionGraph, bindings: dict[str, models.Node], spec: dict[str, Any]) -> bool: + """ + Verifies that enough matched expert branches are driven by router top-k. + + This locates the repeated expert-branch pattern, finds or discovers a top-k + node connected to the MoE combine output, then collects expert branch + bindings whose routing-weight input is reachable from that routing source. + The match succeeds only if the number of routed branches meets the expected + count from `topk.k`, an explicit spec override, or the subgraph minimum. + """ + subgraph = get_repeated_subgraph(fusion, spec["repeated_subgraph"]) + combine_node = bindings.get(spec.get("combine_node", "moe_combine")) + topk_node = bindings.get(spec.get("topk_node", "topk")) + + if subgraph is None: + return bool(spec.get("allow_missing", False)) + + if topk_node is None: + topk_nodes = find_reachable_topk_nodes(graph, combine_node, spec) + topk_node = topk_nodes[0] if topk_nodes else None + + if topk_node is None: + return bool(spec.get("allow_missing", False)) + + branch_bindings = relevant_moe_branch_bindings(graph, fusion, bindings, spec) + min_count = required_moe_branch_count(subgraph, topk_node, spec) + + return len(branch_bindings) >= min_count + + +def match_moe_routing_weights_combine(source: models.Node, graph: models.Graph, fusion: FModels.FusionGraph, bindings: dict[str, models.Node], spec: dict[str, Any]) -> bool: + """ + Verifies that routed expert outputs are actually combined into MoE output. + + This reuses the routed expert-branch filtering, then checks that each + branch's weighted expert output has a forward path to the matched combine + node. That proves both sides of MoE semantics: routing weights are used on + expert outputs, and those weighted outputs feed the final combine op. + """ + subgraph = get_repeated_subgraph(fusion, spec["repeated_subgraph"]) + topk_node = bindings.get(spec.get("topk_node", "topk")) + combine_node = bindings.get(spec["combine_node"]) + + if subgraph is None or combine_node is None: + return bool(spec.get("allow_missing", False)) + + if topk_node is None: + topk_nodes = find_reachable_topk_nodes(graph, combine_node, spec) + topk_node = topk_nodes[0] if topk_nodes else None + + if topk_node is None: + return bool(spec.get("allow_missing", False)) + + branch_bindings = relevant_moe_branch_bindings(graph, fusion, bindings, spec) + combined_branch_count = sum( + 1 for branch_binding in branch_bindings + if branch_output_reaches_combine(branch_binding, subgraph, combine_node, spec) + ) + + return combined_branch_count >= required_moe_branch_count(subgraph, topk_node, spec) + + +def match_grouped_moe_structure(source: models.Node, graph: models.Graph, fusion: FModels.FusionGraph, bindings: dict[str, models.Node], spec: dict[str, Any]) -> bool: + """ + Verifies the grouped-mm MoE structure exported by LFM-style MoE models. + + This matcher handles the compact representation where expert branches are + not expanded into separate per-expert matmuls. The graph edges prove the + local dataflow; this semantic check verifies top-k routing, expert-count + consistency, shared grouped-mm offsets, valid gate/up splitting, and final + top-k combine shape. + """ + topk_node = bindings.get(spec["topk_node"]) + topk_indices_node = bindings.get(spec["topk_indices_node"]) + gate_up_grouped_node = bindings.get(spec["gate_up_grouped_node"]) + down_grouped_node = bindings.get(spec["down_grouped_node"]) + split_node = bindings.get(spec["split_node"]) + offsets_node = bindings.get(spec["offsets_node"]) + combine_node = bindings.get(spec["combine_node"]) + hidden_node = bindings.get("moe_hidden_view") + gate_up_weight_transpose = bindings.get("moe_gate_up_weight_transpose") + down_weight_transpose = bindings.get("moe_down_weight_transpose") + down_weight = match_utils.get_first_input_by_role(fusion, bindings, spec["down_weight_role"]) + + if any(node is None for node in (topk_node, topk_indices_node, gate_up_grouped_node, down_grouped_node, split_node, offsets_node, combine_node)): + return bool(spec.get("allow_missing", False)) + + if get_known_int(get_attr(topk_node, "k", "arg_1")) is None: + return False + + if not match_utils.values_equal(get_attr(topk_indices_node, "index", "arg_1"), 1): + return False + + if not grouped_moe_weight_permute_is_valid(gate_up_weight_transpose): + return False + + if not grouped_moe_weight_permute_is_valid(down_weight_transpose): + return False + + gate_up_offsets = match_utils.get_parent(gate_up_grouped_node, 2) + down_offsets = match_utils.get_parent(down_grouped_node, 2) + + if gate_up_offsets is not offsets_node or down_offsets is not offsets_node: + return False + + if not match_utils.values_equal(get_attr(offsets_node, "dim", "arg_1"), 0): + return False + + if not split_sizes_are_valid(split_node, gate_up_grouped_node, down_weight, spec): + return False + + num_experts = get_num_experts_from_inputs(fusion, bindings, spec) + histc_node = bindings.get("moe_histc") + histc_bins = get_known_int(get_attr(histc_node, "bins", "arg_1")) if histc_node is not None else None + + if num_experts is not None and histc_bins is not None and num_experts != histc_bins: + return False + + if not grouped_moe_rows_match_topk(hidden_node, gate_up_grouped_node, topk_node, spec): + return False + + if not match_utils.values_equal(get_attr(combine_node, "dim", "arg_1"), 1): + return False + + return True + + def match_node_attr_equals(source: models.Node, graph: models.Graph, fusion: FModels.FusionGraph, bindings: dict[str, models.Node], spec: dict[str, Any]) -> bool: + """ + Checks that one matched node has a specific normalized attr value. + + The spec names a synthetic node, an attr key, and the expected value. The + function resolves the real node from bindings, reads its normalized attrs, + and compares with `values_equal`. Missing attrs fail unless + `allow_missing=True` is present in the spec. + """ node = bindings.get(spec["node"]) if node is None: @@ -437,6 +1131,14 @@ def match_node_attr_equals(source: models.Node, graph: models.Graph, fusion: FMo def match_node_attrs_equal(source: models.Node, graph: models.Graph, fusion: FModels.FusionGraph, bindings: dict[str, models.Node], spec: dict[str, Any]) -> bool: + """ + Checks that attrs on two matched nodes are equal to each other. + + This is useful for pattern checks like layernorm, where mean and variance + should reduce over the same dimension. The spec identifies left/right + synthetic nodes and attr names; both attrs must exist and compare equal + unless missing attrs are explicitly allowed. + """ left_node = bindings.get(spec["left_node"]) right_node = bindings.get(spec["right_node"]) @@ -453,6 +1155,14 @@ def match_node_attrs_equal(source: models.Node, graph: models.Graph, fusion: FMo def match_input_value_kind(source: models.Node, graph: models.Graph, fusion: FModels.FusionGraph, bindings: dict[str, models.Node], spec: dict[str, Any]) -> bool: + """ + Checks the semantic kind of an external input to a fusion. + + The fusion graph declares inputs by role, such as `bias` or `weight`. This + matcher resolves the first real node for that role and verifies its + `value_kind` is in the allowed set, preventing activations from being + mistaken for parameters/buffers. + """ input_node = match_utils.get_first_input_by_role(fusion, bindings, spec["role"]) if input_node is None: @@ -462,6 +1172,14 @@ def match_input_value_kind(source: models.Node, graph: models.Graph, fusion: FMo def match_parent_tensor_dim_equal(source: models.Node, graph: models.Graph, fusion: FModels.FusionGraph, bindings: dict[str, models.Node], spec: dict[str, Any]) -> bool: + """ + Compares one dimension from two parent tensors of matched nodes. + + The spec names the left/right synthetic nodes, which parent index to inspect + on each real node, and which shape dimension to compare. It is used for + checks such as attention query/key head-dim compatibility. Missing parents + or shapes follow the spec's `allow_missing` behavior. + """ left_parent = get_bound_parent(bindings, spec["left_node"], spec["left_parent_index"]) right_parent = get_bound_parent(bindings, spec["right_node"], spec["right_parent_index"]) @@ -478,6 +1196,13 @@ def match_parent_tensor_dim_equal(source: models.Node, graph: models.Graph, fusi def match_input_broadcastable_to_node(source: models.Node, graph: models.Graph, fusion: FModels.FusionGraph, bindings: dict[str, models.Node], spec: dict[str, Any]) -> bool: + """ + Checks that a declared fusion input can broadcast to a target node shape. + + This is mainly for mask-like inputs. It resolves an input by role, resolves + the target synthetic node, reads both output shapes, and applies standard + broadcasting. Missing nodes or shapes are controlled by `allow_missing`. + """ input_node = match_utils.get_first_input_by_role(fusion, bindings, spec["role"]) target_node = bindings.get(spec["node"]) @@ -494,6 +1219,14 @@ def match_input_broadcastable_to_node(source: models.Node, graph: models.Graph, def match_same_layer(source: models.Node, graph: models.Graph, fusion: FModels.FusionGraph, bindings: dict[str, models.Node], spec: dict[str, Any]) -> bool: + """ + Checks that multiple matched nodes appear to belong to the same layer. + + The spec provides synthetic node names. For each bound real node, we derive + a layer key from module stack/name/target text. The check passes only when + every derived key is identical, which helps avoid mixing cache/key/value + nodes from different decoder layers. + """ layer_keys = [] for node_name in spec["nodes"]: @@ -518,7 +1251,9 @@ def match_same_layer(source: models.Node, graph: models.Graph, fusion: FModels.F "linear_weight_layout": match_linear_weight_layout, "rope_tables_compatible": match_rope_tables_compatible, "lstm_gate_structure": match_lstm_gate_structure, - "note": match_note, + "moe_expert_branch_routing": match_moe_expert_branch_routing, + "moe_routing_weights_combine": match_moe_routing_weights_combine, + "grouped_moe_structure": match_grouped_moe_structure, "node_attr_equals": match_node_attr_equals, "node_attrs_equal": match_node_attrs_equal, "input_value_kind": match_input_value_kind, diff --git a/python/cactus/transpiler/IR/match_utils.py b/python/cactus/transpiler/IR/match_utils.py index 9e13641bf..08f28ee05 100644 --- a/python/cactus/transpiler/IR/match_utils.py +++ b/python/cactus/transpiler/IR/match_utils.py @@ -5,6 +5,7 @@ MISSING = object() +METADATA_ONLY_TARGETS = {"aten._assert_tensor_metadata.default"} ####################################### Node Matching Utils!!!!! ####################################### @@ -95,7 +96,7 @@ def match_attr_constraint(constraint: FModels.AttrConstraint, node: models.Node, def match_tensor_constraint(constraint: FModels.TensorConstraint, node: models.Node, bindings: dict[str, models.Node]) -> bool: - if node.tensor_output_meta is None: + if not isinstance(node.tensor_output_meta, dict): return False shape = node.tensor_output_meta.get("shape", ()) @@ -296,6 +297,9 @@ def external_children_are_valid(bindings: dict[str, models.Node], fusion: FModel for synth_node_name, node in bindings.items(): for child in node.children: + if child.target in METADATA_ONLY_TARGETS: + continue + if id(child) in internal_ids: continue @@ -390,7 +394,7 @@ def get_first_input_by_role(fusion: FModels.FusionGraph, bindings: dict[str, mod def get_tensor_shape(node: models.Node | None) -> list[Any]: - if node is None or node.tensor_output_meta is None: + if node is None or not isinstance(node.tensor_output_meta, dict): return [] return node.tensor_output_meta.get("shape", []) From 90837665541a2cc8bf7bde2b5467582f1a28dc13 Mon Sep 17 00:00:00 2001 From: Parkirat Date: Tue, 28 Jul 2026 16:37:46 -0400 Subject: [PATCH 58/84] LFM with cache (prefill + decode) --- .../transpiler/Converter/cache_utils.py | 139 +++++++++++++++--- python/cactus/transpiler/Converter/models.py | 3 + python/cactus/transpiler/Fusions/edges.py | 7 + python/cactus/transpiler/Fusions/fusions.py | 12 ++ .../transpiler/ModelProfiles/profiles.py | 2 +- 5 files changed, 139 insertions(+), 24 deletions(-) diff --git a/python/cactus/transpiler/Converter/cache_utils.py b/python/cactus/transpiler/Converter/cache_utils.py index 9b4f296e2..0676cb8f5 100644 --- a/python/cactus/transpiler/Converter/cache_utils.py +++ b/python/cactus/transpiler/Converter/cache_utils.py @@ -9,8 +9,10 @@ class CacheLayerSpec: index: int layer_type: str - key_shape: tuple[int, ...] - value_shape: tuple[int, ...] + key_shape: tuple[int, ...] | None = None + value_shape: tuple[int, ...] | None = None + conv_state_shape: tuple[int, ...] | None = None + recurrent_state_shape: tuple[int, ...] | None = None sliding_window: int | None = None @@ -80,11 +82,12 @@ def forward( "mm_token_type_ids": mm_token_type_ids, "inputs_embeds": inputs_embeds, "use_cache": True, + "return_dict": True, } model_kwargs = {key: value for key, value in model_kwargs.items() if value is not None} outputs = self.model(**filter_forward_kwargs(self.model, model_kwargs)) - return (primary_model_output(outputs), *flatten_dynamic_cache(getattr(outputs, "past_key_values", None))) + return (primary_model_output(outputs), *flatten_dynamic_cache(getattr(outputs, "past_key_values", None), self.cache_spec)) #How: asks HF configs for their decoder/text sub-config when available, otherwise returns the original config. @@ -216,6 +219,20 @@ def cache_spec_from_model(cls: type[CacheSpec], model: torch.nn.Module, batch_si layer_sliding_window = int(sliding_window) cache_sequence_length = min(past_sequence_length, max(1, layer_sliding_window - 1)) + if layer_type == "conv": + layers.append( + CacheLayerSpec( + index=index, + layer_type=str(layer_type), + conv_state_shape=(batch_size, hidden_size, int(getattr(text_config, "conv_L_cache", 1))), + ) + ) + continue + + if layer_type in ("mamba", "linear_attention", "moe"): + layers.append(CacheLayerSpec(index=index, layer_type=str(layer_type))) + continue + cache_shape = (batch_size, layer_num_key_value_heads, cache_sequence_length, layer_head_dim) layers.append( CacheLayerSpec( @@ -245,8 +262,14 @@ def cache_spec_empty_tensors( device = device or torch.device("cpu") for layer in cache_spec.layers: - tensors.append(torch.zeros(layer.key_shape, dtype=dtype, device=device)) - tensors.append(torch.zeros(layer.value_shape, dtype=dtype, device=device)) + if layer.key_shape is not None: + tensors.append(torch.zeros(layer.key_shape, dtype=dtype, device=device)) + if layer.value_shape is not None: + tensors.append(torch.zeros(layer.value_shape, dtype=dtype, device=device)) + if layer.conv_state_shape is not None: + tensors.append(torch.zeros(layer.conv_state_shape, dtype=dtype, device=device)) + if layer.recurrent_state_shape is not None: + tensors.append(torch.zeros(layer.recurrent_state_shape, dtype=dtype, device=device)) return tuple(tensors) @@ -256,44 +279,81 @@ def cache_spec_empty_tensors( def cache_spec_to_dynamic_cache(cache_spec: CacheSpec, flat_tensors: tuple[torch.Tensor, ...]): from transformers.cache_utils import DynamicCache - if len(flat_tensors) != len(cache_spec.layers) * 2: - raise ValueError(f"Expected {len(cache_spec.layers) * 2} cache tensors, got {len(flat_tensors)}") + expected_tensor_count = cache_spec_tensor_count(cache_spec) + + if len(flat_tensors) != expected_tensor_count: + raise ValueError(f"Expected {expected_tensor_count} cache tensors, got {len(flat_tensors)}") cache = DynamicCache(config=cache_spec.config) + flat_index = 0 - for index, layer_spec in enumerate(cache_spec.layers): - flat_index = index * 2 + for layer_spec in cache_spec.layers: layer = cache.layers[layer_spec.index] - key = flat_tensors[flat_index] - value = flat_tensors[flat_index + 1] - layer.keys = key - layer.values = value - layer.dtype = key.dtype - layer.device = key.device - layer.is_initialized = True - if hasattr(layer, "cumulative_length"): - layer.cumulative_length = cache_spec.past_sequence_length + if layer_spec.key_shape is not None and layer_spec.value_shape is not None: + key = flat_tensors[flat_index] + value = flat_tensors[flat_index + 1] + flat_index += 2 + layer.keys = key + layer.values = value + layer.dtype = key.dtype + layer.device = key.device + layer.is_initialized = True + + if hasattr(layer, "cumulative_length"): + layer.cumulative_length = cache_spec.past_sequence_length + + if layer_spec.conv_state_shape is not None: + conv_state = flat_tensors[flat_index].clone() + flat_index += 1 + layer.conv_states = conv_state + layer.dtype = conv_state.dtype + layer.device = conv_state.device + layer.max_batch_size = conv_state.shape[0] + layer.conv_kernel_size = conv_state.shape[-1] + layer.is_conv_states_initialized = True + layer.has_previous_state = cache_spec.past_sequence_length > 0 + + if layer_spec.recurrent_state_shape is not None: + recurrent_state = flat_tensors[flat_index].clone() + flat_index += 1 + layer.recurrent_states = recurrent_state + layer.dtype = recurrent_state.dtype + layer.device = recurrent_state.device + layer.is_recurrent_states_initialized = True + layer.has_previous_state = cache_spec.past_sequence_length > 0 return cache +def cache_spec_tensor_count(cache_spec: CacheSpec) -> int: + count = 0 + + for layer in cache_spec.layers: + count += int(layer.key_shape is not None) + count += int(layer.value_shape is not None) + count += int(layer.conv_state_shape is not None) + count += int(layer.recurrent_state_shape is not None) + + return count + + #How: checks common HF output containers for logits, last_hidden_state, or first tuple item. #Why: CacheExportWrapper.forward needs a stable first exported output before appending flat cache tensors. def primary_model_output(outputs: Any) -> torch.Tensor: if isinstance(outputs, dict): for key in ("logits", "last_hidden_state"): output = outputs.get(key) - if output is not None: + if isinstance(output, torch.Tensor): return output - if hasattr(outputs, "logits") and outputs.logits is not None: + if hasattr(outputs, "logits") and isinstance(outputs.logits, torch.Tensor): return outputs.logits - if hasattr(outputs, "last_hidden_state") and outputs.last_hidden_state is not None: + if hasattr(outputs, "last_hidden_state") and isinstance(outputs.last_hidden_state, torch.Tensor): return outputs.last_hidden_state - if isinstance(outputs, (tuple, list)) and len(outputs) > 0: + if isinstance(outputs, (tuple, list)) and len(outputs) > 0 and isinstance(outputs[0], torch.Tensor): return outputs[0] raise ValueError("Could not find a tensor output to export") @@ -301,7 +361,7 @@ def primary_model_output(outputs: Any) -> torch.Tensor: #How: walks either legacy tuple caches or DynamicCache layers and emits key/value tensors in flat order. #Why: CacheExportWrapper.forward returns this flat tuple so the layermap contains explicit cache output tensors. -def flatten_dynamic_cache(cache: Any) -> tuple[torch.Tensor, ...]: +def flatten_dynamic_cache(cache: Any, cache_spec: CacheSpec | None = None) -> tuple[torch.Tensor, ...]: if cache is None: return () @@ -319,11 +379,44 @@ def flatten_dynamic_cache(cache: Any) -> tuple[torch.Tensor, ...]: if layers is None: return () + if cache_spec is not None: + for layer_spec in cache_spec.layers: + layer = layers[layer_spec.index] + + if layer_spec.key_shape is not None: + key = getattr(layer, "keys", None) + if key is not None: + flat.append(key) + + if layer_spec.value_shape is not None: + value = getattr(layer, "values", None) + if value is not None: + flat.append(value) + + if layer_spec.conv_state_shape is not None: + conv_state = getattr(layer, "conv_states", None) + if conv_state is not None: + flat.append(conv_state) + + if layer_spec.recurrent_state_shape is not None: + recurrent_state = getattr(layer, "recurrent_states", None) + if recurrent_state is not None: + flat.append(recurrent_state) + + return tuple(flat) + for layer in layers: key = getattr(layer, "keys", None) value = getattr(layer, "values", None) if key is None or value is None: + conv_state = getattr(layer, "conv_states", None) + recurrent_state = getattr(layer, "recurrent_states", None) + + if conv_state is not None: + flat.append(conv_state) + if recurrent_state is not None: + flat.append(recurrent_state) continue flat.append(key) diff --git a/python/cactus/transpiler/Converter/models.py b/python/cactus/transpiler/Converter/models.py index 73a7387ed..1393c7764 100644 --- a/python/cactus/transpiler/Converter/models.py +++ b/python/cactus/transpiler/Converter/models.py @@ -311,6 +311,9 @@ def load_model(model_id: str, mp: MP_Models.ModelProfile | None = None) -> torch base_kwargs["token"] = constants.token load_attempts = ( + {**base_kwargs, "dtype": "auto", "low_cpu_mem_usage": True, "local_files_only": True}, + {**base_kwargs, "torch_dtype": "auto", "low_cpu_mem_usage": True, "local_files_only": True}, + {**base_kwargs, "local_files_only": True}, {**base_kwargs, "dtype": "auto", "low_cpu_mem_usage": True}, {**base_kwargs, "torch_dtype": "auto", "low_cpu_mem_usage": True}, base_kwargs, diff --git a/python/cactus/transpiler/Fusions/edges.py b/python/cactus/transpiler/Fusions/edges.py index 3c5d676dc..e8ed2d806 100644 --- a/python/cactus/transpiler/Fusions/edges.py +++ b/python/cactus/transpiler/Fusions/edges.py @@ -124,6 +124,7 @@ def _index_edges_by_dest(edges: dict[str, M.FusionEdge]) -> dict[str, tuple[str, "moe_token_unsqueeze_to_expand": _edge("moe_token_unsqueeze", "moe_token_expand", 0), "moe_token_expand_to_clone": _edge("moe_token_expand", "moe_token_clone", 0), "moe_token_clone_to_flat": _edge("moe_token_clone", "moe_token_flat", 0), + "moe_token_expand_to_flat": _edge("moe_token_expand", "moe_token_flat", 0), "moe_scale_to_routing_flat": _edge("moe_routing_scale", "moe_routing_flat", 0), "moe_topk_indices_to_flat": _edge("moe_topk_indices", "moe_topk_flat", 0), "moe_topk_flat_to_sort": _edge("moe_topk_flat", "moe_sort", 0), @@ -311,6 +312,12 @@ def _index_edges_by_dest(edges: dict[str, M.FusionEdge]) -> dict[str, tuple[str, ), } +EDGE_GROUPS["lfm_grouped_moe_no_token_clone"] = tuple( + "moe_token_expand_to_flat" if edge_name == "moe_token_expand_to_clone" else edge_name + for edge_name in EDGE_GROUPS["lfm_grouped_moe"] + if edge_name != "moe_token_clone_to_flat" +) + EDGE_NAMES_BY_SOURCE: dict[str, tuple[str, ...]] = _index_edges_by_source(EDGES) EDGE_NAMES_BY_DEST: dict[str, tuple[str, ...]] = _index_edges_by_dest(EDGES) diff --git a/python/cactus/transpiler/Fusions/fusions.py b/python/cactus/transpiler/Fusions/fusions.py index ca2edf035..d2705823d 100644 --- a/python/cactus/transpiler/Fusions/fusions.py +++ b/python/cactus/transpiler/Fusions/fusions.py @@ -696,6 +696,16 @@ def _index_by_root_op( }, ) +LFM_GROUPED_MOE_NO_TOKEN_CLONE_GRAPH = _graph( + "lfm_grouped_moe_no_token_clone", + "moe_grouped_combine", + tuple(node_name for node_name in LFM_GROUPED_MOE_GRAPH.nodes if node_name != "moe_token_clone"), + edge_names=E.EDGE_GROUPS["lfm_grouped_moe_no_token_clone"], + inputs=LFM_GROUPED_MOE_GRAPH.inputs, + attr_captures=LFM_GROUPED_MOE_GRAPH.attr_captures, + constraints=LFM_GROUPED_MOE_GRAPH.constraints, +) + LSTM_CELL_GRAPH = _graph( "lstm_cell", "lstm_tanh", @@ -855,6 +865,7 @@ def _index_by_root_op( "attention_cached": ATTENTION_CACHED_GRAPH, "moe_layer_gated": MOE_GATED_GRAPH, "lfm_grouped_moe": LFM_GROUPED_MOE_GRAPH, + "lfm_grouped_moe_no_token_clone": LFM_GROUPED_MOE_NO_TOKEN_CLONE_GRAPH, "lstm_cell": LSTM_CELL_GRAPH, "gated_deltanet_decode": DELTANET_DECODE_GRAPH, "gated_deltanet_prefill": DELTANET_PREFILL_GRAPH, @@ -910,6 +921,7 @@ def _index_by_root_op( "attention_cached": _definition("attention_cached", "attention_cached", ATTENTION_CACHED_GRAPH, fusion_fields=("generic", "attention", "cache"), supported_inference_modes=("decode_with_cache",)), "moe_layer_gated": _definition("moe_layer_gated", "moe_layer_gated", MOE_GATED_GRAPH, fusion_fields=("generic", "moe")), "lfm_grouped_moe": _definition("lfm_grouped_moe", "moe_layer_gated", LFM_GROUPED_MOE_GRAPH, fusion_fields=("generic", "moe", "lfm_moe")), + "lfm_grouped_moe_no_token_clone": _definition("lfm_grouped_moe_no_token_clone", "moe_layer_gated", LFM_GROUPED_MOE_NO_TOKEN_CLONE_GRAPH, fusion_fields=("generic", "moe", "lfm_moe")), "lstm_cell": _definition("lstm_cell", "lstm_cell", LSTM_CELL_GRAPH, fusion_fields=("generic", "recurrent", "audio")), "gated_deltanet_decode": _definition("gated_deltanet_decode", "gated_deltanet_decode", DELTANET_DECODE_GRAPH, fusion_fields=("generic", "recurrent", "cache")), "gated_deltanet_prefill": _definition("gated_deltanet_prefill", "gated_deltanet_prefill", DELTANET_PREFILL_GRAPH, fusion_fields=("generic", "recurrent", "cache")), diff --git a/python/cactus/transpiler/ModelProfiles/profiles.py b/python/cactus/transpiler/ModelProfiles/profiles.py index 0bf2ffa83..8be35c27d 100644 --- a/python/cactus/transpiler/ModelProfiles/profiles.py +++ b/python/cactus/transpiler/ModelProfiles/profiles.py @@ -131,7 +131,7 @@ components=TEXT_COMPONENTS, inference_type=QWEN2_5_INFERENCE_PATTERNS, cache_type=("attention_kv", "conv_state"), - cache_policy=(), + cache_policy=("dynamic_cache",), files=("config.json", "generation_config.json", "tokenizer.json", "tokenizer_config.json"), fusion_fields=( "generic", From f97178ca4b09e6e373801b08cc0438842f266d03 Mon Sep 17 00:00:00 2001 From: Parkirat Date: Tue, 28 Jul 2026 20:04:35 -0400 Subject: [PATCH 59/84] Added simplification of IR step (uses fusion matchers to determine fusion) --- python/cactus/transpiler/Converter/convert.py | 14 +- python/cactus/transpiler/IR/models.py | 245 +++++++++++++++++- python/cactus/transpiler/IR/simplify_ir.py | 242 +++++++++++++++++ python/cactus/transpiler/main.py | 2 + 4 files changed, 501 insertions(+), 2 deletions(-) diff --git a/python/cactus/transpiler/Converter/convert.py b/python/cactus/transpiler/Converter/convert.py index be1c42bf9..32a53394f 100644 --- a/python/cactus/transpiler/Converter/convert.py +++ b/python/cactus/transpiler/Converter/convert.py @@ -3,6 +3,7 @@ from . import models from ..ModelProfiles import models as MP_Models from ..ModelProfiles import profiles as MP_Profiles +from ..IR import simplify_ir def export_layer_map( @@ -23,9 +24,20 @@ def convert( output_path: str, custom_profile: MP_Models.ModelProfile | None = None, inference_mode: str = "prefill_no_cache", + simplified_output_path: str | None = None, ) -> models.LayerMap: - layer_map = export_layer_map(model_id=model_id, input_modalities=input_modalities, custom_profile=custom_profile, inference_mode=inference_mode) + profile = custom_profile or MP_Profiles.MODEL_ID_MAP[model_id] + layer_map = export_layer_map(model_id=model_id, input_modalities=input_modalities, custom_profile=profile, inference_mode=inference_mode) Path(output_path).write_text(layer_map.model_dump_json(indent=4), encoding="utf-8") + + if simplified_output_path is not None: + simplify_ir.write_simplified_json( + layer_map, + simplified_output_path, + input_modalities=input_modalities, + fusion_fields=profile.fusion_fields, + ) + return layer_map diff --git a/python/cactus/transpiler/IR/models.py b/python/cactus/transpiler/IR/models.py index f0154d51e..f3ae40063 100644 --- a/python/cactus/transpiler/IR/models.py +++ b/python/cactus/transpiler/IR/models.py @@ -47,12 +47,22 @@ class Graph: outputs: tuple[Node, ...] nodes: tuple[Node, ...] nodes_map: dict[str, Node] + model_name: str = "" + task: str = "" + graph_signature: str = "" + range_constants: str = "" fusions: list["FusionResult"] = field(default_factory=list) @classmethod def from_map(cls, layer_map: CModels.LayerMap) -> "Graph": return generate_graph(cls, layer_map) + def apply_fusions(self, fusion_results: list["FusionResult"] | tuple["FusionResult", ...]) -> "Graph": + return apply_fusions_to_graph(self, tuple(fusion_results)) + + def to_layer_map(self) -> CModels.LayerMap: + return graph_to_layer_map(self) + @dataclass(slots=True) class FusionResult: @@ -222,4 +232,237 @@ def generate_graph(cls: type[Graph], layer_map: CModels.LayerMap) -> Graph: if not outputs: outputs = tuple(node for node in nodes if not node.children) - return cls(source=sources[0] if sources else None, sources=sources, outputs=outputs, nodes=nodes, nodes_map=nodes_map) + return cls( + source=sources[0] if sources else None, + sources=sources, + outputs=outputs, + nodes=nodes, + nodes_map=nodes_map, + model_name=layer_map.model_name, + task=layer_map.task, + graph_signature=layer_map.graph_signature, + range_constants=layer_map.range_constants, + ) + + +def apply_fusions_to_graph(graph: Graph, fusion_results: tuple[FusionResult, ...]) -> Graph: + if not fusion_results: + return graph + + validate_non_overlapping_fusions(fusion_results) + + fused_nodes, replacement_names = build_fused_nodes(graph, fusion_results) + consumed_names = expanded_consumed_names(graph, replacement_names) + kept_nodes = tuple(clone_rewritten_node(node, replacement_names) for node in graph.nodes if node.name not in consumed_names) + rewritten_fused_nodes = tuple(clone_rewritten_node(node, replacement_names) for node in fused_nodes) + return rebuild_graph((*kept_nodes, *rewritten_fused_nodes), graph, fusion_results) + + +def validate_non_overlapping_fusions(fusion_results: tuple[FusionResult, ...]) -> None: + consumed_names: set[str] = set() + + for result in fusion_results: + overlap = consumed_names.intersection(result.consumed_node_names) + + if overlap: + overlap_list = ", ".join(sorted(overlap)) + raise ValueError(f"Fusion {result.fusion_name} overlaps an earlier fusion at: {overlap_list}") + + consumed_names.update(result.consumed_node_names) + + +def build_fused_nodes(graph: Graph, fusion_results: tuple[FusionResult, ...]) -> tuple[tuple[Node, ...], dict[str, str]]: + fused_nodes: list[Node] = [] + replacement_names: dict[str, str] = {} + used_names = set(graph.nodes_map) + + for count, result in enumerate(fusion_results): + fused_name = unique_fused_name(result, count, used_names) + used_names.add(fused_name) + fused_node = fused_node_from_result(result, fused_name) + fused_nodes.append(fused_node) + + for consumed_name in result.consumed_node_names: + replacement_names[consumed_name] = fused_name + + return tuple(fused_nodes), replacement_names + + +def unique_fused_name(result: FusionResult, count: int, used_names: set[str]) -> str: + base = sanitize_node_name(result.cactus_op or result.fusion_name) + candidate = f"{base}_{count}" + suffix = count + + while candidate in used_names: + suffix += 1 + candidate = f"{base}_{suffix}" + + return candidate + + +def sanitize_node_name(name: str) -> str: + return "".join(char if char.isalnum() or char == "_" else "_" for char in name).strip("_") or "fused" + + +def fused_node_from_result(result: FusionResult, fused_name: str) -> Node: + return Node( + index=result.source.index, + name=fused_name, + node_type="call_function", + target=result.target, + args=[{"node": node.name} for node in result.external_inputs], + kwargs=dict(result.attrs), + users=(), + tensor_output_meta=result.source.tensor_output_meta, + module_stack=result.source.module_stack, + value_kind=FModels.ValueKind.ACTIVATION, + attrs=dict(result.attrs), + ) + + +def expanded_consumed_names(graph: Graph, replacement_names: dict[str, str]) -> frozenset[str]: + consumed_names = set(replacement_names) + changed = True + + while changed: + changed = False + + for node in graph.nodes: + if node.name in consumed_names: + continue + + if node.target != "aten._assert_tensor_metadata.default": + continue + + parent_replacements = [replacement_names[parent.name] for parent in node.parents if parent.name in replacement_names] + + if not parent_replacements: + continue + + replacement_names[node.name] = parent_replacements[0] + consumed_names.add(node.name) + changed = True + + return frozenset(consumed_names) + + +def clone_rewritten_node(node: Node, replacement_names: dict[str, str]) -> Node: + return Node( + index=node.index, + name=node.name, + node_type=node.node_type, + target=node.target, + args=rewrite_node_refs(node.args, replacement_names), + kwargs=rewrite_node_refs(node.kwargs, replacement_names), + users=(), + tensor_output_meta=node.tensor_output_meta, + module_stack=node.module_stack, + value_kind=node.value_kind, + attrs=dict(node.attrs), + ) + + +def rewrite_node_refs(value: Any, replacement_names: dict[str, str]) -> Any: + if isinstance(value, dict) and set(value.keys()) == {"node"} and isinstance(value["node"], str): + return {"node": replacement_names.get(value["node"], value["node"])} + + if isinstance(value, list): + return [rewrite_node_refs(item, replacement_names) for item in value] + + if isinstance(value, tuple): + return tuple(rewrite_node_refs(item, replacement_names) for item in value) + + if isinstance(value, dict): + return {key: rewrite_node_refs(item, replacement_names) for key, item in value.items()} + + return value + + +def rebuild_graph(nodes: tuple[Node, ...], original_graph: Graph, fusion_results: tuple[FusionResult, ...] = ()) -> Graph: + nodes_map = {node.name: node for node in nodes} + parent_lists: dict[str, list[Node]] = {node.name: [] for node in nodes} + child_lists: dict[str, list[Node]] = {node.name: [] for node in nodes} + + for node in nodes: + parent_names = extract_node_refs((node.args, node.kwargs)) + + for parent_name in parent_names: + if parent_name not in nodes_map: + raise ValueError(f"{node.name} references missing parent node {parent_name}") + + parent = nodes_map[parent_name] + parent_lists[node.name].append(parent) + + if node not in child_lists[parent.name]: + child_lists[parent.name].append(node) + + for node in nodes: + node.parents = tuple(parent_lists[node.name]) + node.children = tuple(child_lists[node.name]) + node.users = tuple(child.name for child in node.children) + + sources = tuple(node for node in nodes if not node.parents) + outputs = tuple(node for node in nodes if node.is_output) + + if not outputs: + outputs = tuple(node for node in nodes if not node.children) + + return Graph( + source=sources[0] if sources else None, + sources=sources, + outputs=outputs, + nodes=nodes, + nodes_map=nodes_map, + model_name=original_graph.model_name, + task=original_graph.task, + graph_signature=original_graph.graph_signature, + range_constants=original_graph.range_constants, + fusions=list(fusion_results), + ) + + +def graph_to_layer_map(graph: Graph) -> CModels.LayerMap: + records = [ + CModels.LayerRecord( + index=index, + name=node.name, + node_type=node.node_type, + target=node.target, + args=node.args, + kwargs=node.kwargs, + users=[child.name for child in node.children], + tensor_output_meta=node.tensor_output_meta, + module_stack=node.module_stack, + ) + for index, node in enumerate(topological_sort(graph)) + ] + + return CModels.LayerMap( + model_name=graph.model_name, + task=graph.task, + graph_signature=graph.graph_signature, + range_constants=graph.range_constants, + nodes=records, + ) + + +def topological_sort(graph: Graph) -> tuple[Node, ...]: + indegree = {node.name: len({parent.name for parent in node.parents}) for node in graph.nodes} + ready = sorted((node for node in graph.nodes if indegree[node.name] == 0), key=lambda node: (node.index, node.name)) + ordered: list[Node] = [] + + while ready: + node = ready.pop(0) + ordered.append(node) + + for child in sorted(node.children, key=lambda child: (child.index, child.name)): + indegree[child.name] -= 1 + + if indegree[child.name] == 0: + ready.append(child) + ready.sort(key=lambda ready_node: (ready_node.index, ready_node.name)) + + if len(ordered) != len(graph.nodes): + raise ValueError("Cannot topologically sort graph with cycles or missing edges") + + return tuple(ordered) diff --git a/python/cactus/transpiler/IR/simplify_ir.py b/python/cactus/transpiler/IR/simplify_ir.py index e69de29bb..7fdaa3d50 100644 --- a/python/cactus/transpiler/IR/simplify_ir.py +++ b/python/cactus/transpiler/IR/simplify_ir.py @@ -0,0 +1,242 @@ +from collections.abc import Iterable +from pathlib import Path + +from . import match, match_utils, models +from ..Converter import models as CModels +from ..Fusions import fusions as Fusions +from ..Fusions import models as FModels + + +def simplify( + layer_map: CModels.LayerMap, + *, + input_modalities: tuple[str, ...] = (), + fusion_fields: tuple[str, ...] = (), + max_fusions: int | None = None, +) -> CModels.LayerMap: + graph = models.Graph.from_map(layer_map) + simplified_graph = rev_top_sort( + graph, + inference_mode=layer_map.task, + input_modalities=input_modalities, + fusion_fields=fusion_fields, + max_fusions=max_fusions, + ) + return simplified_graph.to_layer_map() + + +def write_simplified_json( + layer_map: CModels.LayerMap, + output_path: str | Path, + *, + input_modalities: tuple[str, ...] = (), + fusion_fields: tuple[str, ...] = (), + max_fusions: int | None = None, +) -> CModels.LayerMap: + simplified = simplify( + layer_map, + input_modalities=input_modalities, + fusion_fields=fusion_fields, + max_fusions=max_fusions, + ) + path = Path(output_path) + path.parent.mkdir(parents=True, exist_ok=True) + path.write_text(simplified.model_dump_json(indent=4), encoding="utf-8") + return simplified + + +def rev_top_sort( + graph: models.Graph, + *, + inference_mode: str | None = None, + input_modalities: tuple[str, ...] = (), + fusion_fields: tuple[str, ...] = (), + max_fusions: int | None = None, +) -> models.Graph: + consumed_names: set[str] = set() + fusion_results: list[models.FusionResult] = [] + + for node in reverse_topological_nodes(graph): + if node.name in consumed_names: + continue + + result = try_match_from_node( + node, + graph, + consumed_names=consumed_names, + inference_mode=inference_mode, + input_modalities=input_modalities, + fusion_fields=fusion_fields, + ) + + if result is None: + continue + + fusion_results.append(result) + consumed_names.update(result.consumed_node_names) + + if max_fusions is not None and len(fusion_results) >= max_fusions: + break + + graph.fusions = fusion_results + return graph.apply_fusions(fusion_results) + + +def try_match_from_node( + node: models.Node, + graph: models.Graph, + *, + consumed_names: set[str] | frozenset[str] = frozenset(), + inference_mode: str | None = None, + input_modalities: tuple[str, ...] = (), + fusion_fields: tuple[str, ...] = (), +) -> models.FusionResult | None: + if not node.is_operation: + return None + + for fusion in candidate_fusions_for_node(node, fusion_fields): + bindings = bind_matching_fusion( + node, + graph, + fusion, + inference_mode=inference_mode, + input_modalities=input_modalities, + fusion_fields=fusion_fields, + ) + + if bindings is None: + continue + + result = fusion_result_from_bindings(fusion, node, bindings) + + if result.consumed_node_names.intersection(consumed_names): + continue + + return result + + return None + + +def bind_matching_fusion( + node: models.Node, + graph: models.Graph, + fusion: FModels.FusionDefinition, + *, + inference_mode: str | None = None, + input_modalities: tuple[str, ...] = (), + fusion_fields: tuple[str, ...] = (), +) -> dict[str, models.Node] | None: + bindings = match_utils.bind_fusion_graph(node, fusion.graph, match.match_nodes) + + if bindings is None: + return None + + if not all( + matcher(node, graph, fusion, bindings, inference_mode, input_modalities, fusion_fields) + for matcher in match.FUSION_DEFINITION_MATCHERS + ): + return None + + return bindings + + +def fusion_result_from_bindings( + fusion: FModels.FusionDefinition, + source: models.Node, + bindings: dict[str, models.Node], +) -> models.FusionResult: + matched_nodes = tuple(sorted(unique_nodes(bindings.values()), key=lambda node: (node.index, node.name))) + external_inputs = collect_external_inputs(fusion.graph, bindings) + attrs = match_utils.collect_definition_attrs(fusion.graph, bindings) + + return models.FusionResult.from_match( + fusion=fusion, + source=source, + matched_nodes=matched_nodes, + bindings=dict(bindings), + external_inputs=external_inputs, + attrs=attrs, + ) + + +def collect_external_inputs( + fusion: FModels.FusionGraph, + bindings: dict[str, models.Node], +) -> tuple[models.Node, ...]: + external_inputs: list[models.Node] = [] + + for input_spec in fusion.inputs: + external_inputs.extend(match_utils.get_fusion_input_nodes(input_spec, bindings)) + + for cache_input in fusion.cache_inputs: + cache_node = match_utils.get_node_ref_parent(cache_input.source, bindings) + + if cache_node is not None and not any(existing is cache_node for existing in external_inputs): + external_inputs.append(cache_node) + + return tuple(external_inputs) + + +def candidate_fusions_for_node( + node: models.Node, + fusion_fields: tuple[str, ...] = (), +) -> tuple[FModels.FusionDefinition, ...]: + candidates: list[FModels.FusionDefinition] = [] + seen: set[int] = set() + + for key in (node.target, node.node_type): + for fusion in Fusions.FUSIONS_BY_ROOT_OP.get(key, ()): + if id(fusion) in seen: + continue + + if not fusion_enabled_for_fields(fusion, fusion_fields): + continue + + seen.add(id(fusion)) + candidates.append(fusion) + + return tuple(sorted(candidates, key=fusion_priority, reverse=True)) + + +def fusion_enabled_for_fields(fusion: FModels.FusionDefinition, fusion_fields: tuple[str, ...]) -> bool: + if not fusion_fields or not fusion.fusion_fields: + return True + + selected_fields = set(fusion_fields) + fusion_specific_fields = set(fusion.fusion_fields) - {"generic"} + + if not fusion_specific_fields: + return "generic" in selected_fields + + if fusion_specific_fields == {"direct"}: + return "generic" in selected_fields or "direct" in selected_fields + + return not fusion_specific_fields.isdisjoint(selected_fields - {"generic"}) + + +def fusion_priority(fusion: FModels.FusionDefinition) -> tuple[int, int, int, int, int]: + required_attrs = fusion.metadata.get("required_attrs", {}) + + return ( + len(fusion.graph.nodes), + len(fusion.graph.edges), + len(fusion.graph.constraints), + len(required_attrs), + int("direct" not in fusion.fusion_fields), + ) + + +def reverse_topological_nodes(graph: models.Graph) -> tuple[models.Node, ...]: + return tuple(reversed(models.topological_sort(graph))) + + +def unique_nodes(nodes: Iterable[models.Node]) -> tuple[models.Node, ...]: + unique: list[models.Node] = [] + + for node in nodes: + if any(existing is node for existing in unique): + continue + + unique.append(node) + + return tuple(unique) diff --git a/python/cactus/transpiler/main.py b/python/cactus/transpiler/main.py index 103615b34..1a3187734 100644 --- a/python/cactus/transpiler/main.py +++ b/python/cactus/transpiler/main.py @@ -17,12 +17,14 @@ def main() -> None: model_id=model_id, input_modalities=input_modalities, output_path=str(output_dir / "output_prefill_with_cache.json"), + simplified_output_path=str(output_dir / "output_prefill_with_cache_simplified.json"), inference_mode=constants.PREFILL_WITH_CACHE_MODE, ) convert( model_id=model_id, input_modalities=input_modalities, output_path=str(output_dir / "output_decode_with_cache.json"), + simplified_output_path=str(output_dir / "output_decode_with_cache_simplified.json"), inference_mode=constants.DECODE_WITH_CACHE_MODE, ) From 3a51ea5bbeffde93d4a290259852913593af3819 Mon Sep 17 00:00:00 2001 From: Parkirat Date: Tue, 28 Jul 2026 20:04:51 -0400 Subject: [PATCH 60/84] Added simplification of IR step (uses fusion matchers to determine fusion) --- python/cactus/transpiler/IR/simplify_ir.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/python/cactus/transpiler/IR/simplify_ir.py b/python/cactus/transpiler/IR/simplify_ir.py index 7fdaa3d50..3953fd48c 100644 --- a/python/cactus/transpiler/IR/simplify_ir.py +++ b/python/cactus/transpiler/IR/simplify_ir.py @@ -7,6 +7,8 @@ from ..Fusions import models as FModels +#TODO: Clean up code + def simplify( layer_map: CModels.LayerMap, *, From 6db22b2c1a54edd5925f11ab623d70a3e35714ff Mon Sep 17 00:00:00 2001 From: Parkirat Date: Wed, 29 Jul 2026 10:03:18 -0400 Subject: [PATCH 61/84] Added cactus ops lowering with generator --- python/cactus/transpiler/Converter/models.py | 71 +- .../cactus/transpiler/Generator/constants.py | 316 +++++ .../cactus/transpiler/Generator/generate.py | 1155 +++++++++++++++++ python/cactus/transpiler/Generator/models.py | 533 ++++++++ python/cactus/transpiler/IR/models.py | 8 + 5 files changed, 2081 insertions(+), 2 deletions(-) create mode 100644 python/cactus/transpiler/Generator/constants.py create mode 100644 python/cactus/transpiler/Generator/generate.py diff --git a/python/cactus/transpiler/Converter/models.py b/python/cactus/transpiler/Converter/models.py index 1393c7764..989a4c2bd 100644 --- a/python/cactus/transpiler/Converter/models.py +++ b/python/cactus/transpiler/Converter/models.py @@ -5,7 +5,7 @@ from typing import Any from huggingface_hub import hf_hub_download import numpy as np -from pydantic import BaseModel +from pydantic import BaseModel, Field import torch from transformers import AutoModel @@ -57,18 +57,36 @@ class LayerRecord(BaseModel): def from_node(cls, num: int, x: torch.fx.Node) -> "LayerRecord": return cls(index=num, name=str(x.name), node_type=str(x.op), target=str(x.target), args=jsonable(x.args), kwargs=jsonable(x.kwargs), users=[user.name for user in x.users], tensor_output_meta=extract_tensor_meta(x), module_stack=extract_module_stack(x)) + +class GraphSpecRecord(BaseModel): + kind: str + arg_name: str | None = None + target: str | None = None + persistent: bool | None = None + + #Serializable top-level export IR container. class LayerMap(BaseModel): model_name: str task: str graph_signature: str range_constants: str + input_specs: list[GraphSpecRecord] = Field(default_factory=list) + output_specs: list[GraphSpecRecord] = Field(default_factory=list) nodes: list[LayerRecord] #Builds a LayerMap from a torch.export ExportedProgram and serialized nodes. @classmethod def from_data(cls, x: torch.export.ExportedProgram, name: str, model_task: str, nodes_list: list[LayerRecord]) -> "LayerMap": - return cls(model_name=name, task=model_task, graph_signature=repr(x.graph_signature), range_constants=repr(x.range_constraints), nodes=nodes_list) + return cls( + model_name=name, + task=model_task, + graph_signature=repr(x.graph_signature), + range_constants=repr(x.range_constraints), + input_specs=extract_graph_signature_specs(x.graph_signature, "input_specs"), + output_specs=extract_graph_signature_specs(x.graph_signature, "output_specs"), + nodes=nodes_list, + ) ###################################################### Model utility helpers!!!!! ########################################################################. @@ -134,6 +152,55 @@ def extract_module_stack(node: torch.fx.Node) -> Any | None: return out +def extract_graph_signature_specs(graph_signature: Any, field_name: str) -> list[GraphSpecRecord]: + specs = getattr(graph_signature, field_name, ()) or () + return [graph_spec_record(spec) for spec in specs] + + +def graph_spec_record(spec: Any) -> GraphSpecRecord: + return GraphSpecRecord( + kind=spec_kind(spec), + arg_name=spec_arg_name(spec), + target=none_or_str(getattr(spec, "target", None)), + persistent=getattr(spec, "persistent", None), + ) + + +def spec_kind(spec: Any) -> str: + kind = getattr(spec, "kind", None) + + if kind is None: + return "unknown" + + name = getattr(kind, "name", None) + + if name is not None: + return str(name).lower() + + return str(kind) + + +def spec_arg_name(spec: Any) -> str | None: + arg = getattr(spec, "arg", None) + + if arg is None: + return None + + name = getattr(arg, "name", None) + + if name is not None: + return str(name) + + return None + + +def none_or_str(value: Any) -> str | None: + if value is None: + return None + + return str(value) + + def load_configs(mp: MP_Models.ModelProfile, model_id: str | None) -> dict[str, dict[str, Any]]: output_dir = constants.CONVERTER_JSON_DIR / mp.model_profiles output_dir.mkdir(parents=True, exist_ok=True) diff --git a/python/cactus/transpiler/Generator/constants.py b/python/cactus/transpiler/Generator/constants.py new file mode 100644 index 000000000..f26f5b3cd --- /dev/null +++ b/python/cactus/transpiler/Generator/constants.py @@ -0,0 +1,316 @@ +CACTUS_TARGET_PREFIX = "cactus." +DEFAULT_COMPONENT_NAME = "model" +DEFAULT_GRAPH_SUFFIX = ".cactus" +DEFAULT_INPUT_PRECISION = "FP16" + +DTYPE_TO_PRECISION: dict[str, str] = { + "torch.float16": "FP16", + "torch.bfloat16": "FP16", + "torch.float32": "FP32", + "torch.float": "FP32", + "torch.int8": "INT8", + "torch.uint8": "INT8", + "torch.int16": "INT8", + "torch.int32": "INT8", + "torch.int64": "INT8", + "torch.long": "INT8", + "torch.bool": "INT8", +} + +WEIGHT_VALUE_KINDS = { + "parameter", + "buffer", +} + +INPUT_VALUE_KINDS = { + "unknown", + "user_input", + "lifted_constant", + "cache_input", + "cache_output", + "cache_state", +} + +PASS_THROUGH_TARGETS = { + "aten._assert_tensor_metadata.default", + "aten.clone.default", + "aten.contiguous.default", + "aten.detach.default", +} + +COPY_TARGETS = { + "aten.copy.default", +} + +CONSTANT_INPUT_TARGETS = { + "aten.arange.default", + "aten.arange.start", + "aten.arange.start_step", + "aten.empty_permuted.default", + "aten.full.default", + "aten.full_like.default", + "aten.scalar_tensor.default", +} + +BINARY_TARGETS: dict[str, str] = { + "cactus.add": "add", + "cactus.subtract": "subtract", + "cactus.multiply": "multiply", + "cactus.divide": "divide", + "cactus.not_equal": "not_equal", + "aten.add.Tensor": "add", + "aten.sub.Tensor": "subtract", + "aten.mul.Tensor": "multiply", + "aten.div.Tensor": "divide", + "aten.ne.Tensor": "not_equal", +} + +SCALAR_TARGETS: dict[str, tuple[str, str]] = { + "cactus.scalar_add": ("scalar_add", "value"), + "cactus.scalar_subtract": ("scalar_subtract", "value"), + "cactus.scalar_multiply": ("scalar_multiply", "value"), + "cactus.scalar_divide": ("scalar_divide", "value"), + "cactus.scalar_floor_divide": ("scalar_floor_divide", "value"), + "cactus.scalar_not_equal": ("scalar_not_equal", "value"), + "aten.add.Scalar": ("scalar_add", "other"), + "aten.sub.Scalar": ("scalar_subtract", "other"), + "aten.mul.Scalar": ("scalar_multiply", "other"), + "aten.div.Scalar": ("scalar_divide", "other"), + "aten.floor_divide.default": ("scalar_floor_divide", "other"), + "aten.ne.Scalar": ("scalar_not_equal", "other"), +} + +UNARY_TARGETS: dict[str, str] = { + "cactus.abs": "abs", + "cactus.scalar_exp": "scalar_exp", + "cactus.scalar_sqrt": "scalar_sqrt", + "cactus.scalar_cos": "scalar_cos", + "cactus.scalar_sin": "scalar_sin", + "cactus.scalar_log": "scalar_log", + "cactus.relu": "relu", + "cactus.silu": "silu", + "cactus.gelu": "gelu", + "cactus.gelu_erf": "gelu_erf", + "cactus.sigmoid": "sigmoid", + "cactus.tanh": "tanh", + "aten.abs.default": "abs", + "aten.exp.default": "scalar_exp", + "aten.sqrt.default": "scalar_sqrt", + "aten.cos.default": "scalar_cos", + "aten.sin.default": "scalar_sin", + "aten.log.default": "scalar_log", + "aten.relu.default": "relu", + "aten.silu.default": "silu", + "aten.gelu.default": "gelu", + "aten.sigmoid.default": "sigmoid", + "aten.tanh.default": "tanh", +} + +POW_TARGETS = { + "cactus.pow", + "aten.pow.Tensor_Scalar", +} + +REDUCE_TARGETS: dict[str, str] = { + "cactus.mean": "mean", + "cactus.sum": "sum", + "cactus.variance": "variance", + "cactus.min": "min", + "cactus.max": "max", + "aten.mean.dim": "mean", + "aten.sum.dim_IntList": "sum", + "aten.var.dim": "variance", + "aten.var.correction": "variance", + "aten.amax.default": "max", + "aten.amin.default": "min", +} + +SHAPE_TARGETS: dict[str, str] = { + "cactus.view": "view", + "cactus.reshape": "reshape", + "aten.view.default": "view", + "aten.reshape.default": "reshape", +} + +EXPAND_TARGETS = { + "aten.expand.default", +} + +UNSQUEEZE_TARGETS = { + "aten.unsqueeze.default", +} + +FLATTEN_TARGETS = { + "aten.flatten.using_ints", +} + +TRANSPOSE_TARGETS = { + "cactus.transpose", + "aten.t.default", + "aten.transpose.int", + "aten.permute.default", +} + +SLICE_TARGETS = { + "cactus.slice", + "aten.slice.Tensor", + "aten.select.int", + "aten.narrow.default", +} + +INDEX_TARGETS = { + "cactus.index", +} + +CAT_TARGETS = { + "cactus.cat", + "aten.cat.default", +} + +MATMUL_TARGETS = { + "cactus.matmul", + "aten.mm.default", + "aten.matmul.default", + "aten.bmm.default", +} + +ADDM_CONST_TARGETS = { + "aten.addmm.default", +} + +SPLIT_TARGETS = { + "aten.split_with_sizes.default", +} + +GETITEM_TARGETS = { + "operator.getitem", + "", +} + +TO_COPY_TARGETS = { + "aten._to_copy.default", +} + +NEG_TARGETS = { + "aten.neg.default", +} + +SOFTMAX_TARGETS = { + "cactus.softmax", + "aten.softmax.int", + "aten._softmax.default", +} + +TOPK_TARGETS = { + "cactus.topk", + "aten.topk.default", +} + +GATHER_TARGETS = { + "cactus.gather", + "aten.gather.default", +} + +EMBEDDING_TARGETS = { + "cactus.embedding_from_tensor", + "aten.embedding.default", +} + +CLAMP_TARGETS = { + "cactus.clamp", + "aten.clamp.default", + "aten.clamp.Tensor", +} + +NORM_TARGETS = { + "cactus.rms_norm", + "cactus.layernorm", + "cactus.layer_norm", + "cactus.groupnorm", + "cactus.group_norm", + "cactus.batchnorm", + "cactus.batch_norm", + "aten.native_layer_norm.default", + "aten.layer_norm.default", + "aten.native_group_norm.default", + "aten.group_norm.default", + "aten.native_batch_norm.default", + "aten.batch_norm.default", +} + +CONV_TARGETS = { + "cactus.conv1d", + "cactus.conv1d_causal", + "cactus.conv1d_k3", + "cactus.conv1d_k7s3", + "cactus.conv1d_same_depthwise_k9", + "cactus.conv1d_pointwise", + "cactus.conv2d", + "cactus.conv2d_depthwise_k3s2p1", + "cactus.conv2d_k3s1p1", + "cactus.conv2d_k3s2p1", + "cactus.conv2d_pointwise_1x1", + "aten.convolution.default", + "aten.conv1d.default", + "aten.conv2d.default", +} + +ATTENTION_TARGETS = { + "cactus.attention", + "cactus.attention_cached", + "aten.scaled_dot_product_attention.default", +} + +CACHE_TARGETS = { + "cactus.kv_cache_append", + "cactus.kv_cache_state", + "cactus.conv_cache_append", + "cactus.conv_cache_initialize", + "cactus.conv_cache_state", + "cactus.recurrent_cache_state", + "cactus.recurrent_cache_write", +} + +MOE_TARGETS = { + "cactus.moe_layer_gated", + "cactus.moe_layer_ungated", + "cactus.dense_mlp_tq_fused", +} + +SPECIAL_CACTUS_TARGETS = { + "cactus.rope", + "cactus.rope_gptj", + "cactus.glu", + "cactus.lstm_cell", + "cactus.gated_deltanet_decode", + "cactus.gated_deltanet_prefill", + "cactus.rel_pos_bias", + "cactus.sample", + "cactus.scatter_topk", + "cactus.gaussian_topk", + "cactus.altup_predict", + "cactus.altup_correct", + "cactus.stft", + "cactus.rfft", + "cactus.irfft", + "cactus.spectrogram", + "cactus.image_preprocess", + "cactus.bilinear_interpolation", +} + +UNSUPPORTED_SEMANTIC_TARGETS = { + "aten.any.dim", + "aten.bitwise_and.Tensor", + "aten.bitwise_or.Tensor", + "aten.constant_pad_nd.default", + "aten.eq.Scalar", + "aten.eq.Tensor", + "aten.fmod.Scalar", + "aten.index.Tensor", + "aten.le.Tensor", + "aten.logical_not.default", + "aten.slice_scatter.default", + "aten.where.self", +} + +OPEN_SLICE_END = 9_223_372_036_854_775_807 diff --git a/python/cactus/transpiler/Generator/generate.py b/python/cactus/transpiler/Generator/generate.py new file mode 100644 index 000000000..039c7392d --- /dev/null +++ b/python/cactus/transpiler/Generator/generate.py @@ -0,0 +1,1155 @@ +from __future__ import annotations + +import json +from collections.abc import Mapping +from pathlib import Path +from typing import Any + +from . import constants +from . import models +from ..Converter import models as CModels +from ..IR import models as IRModels + + +class UnsupportedLoweringError(NotImplementedError): + pass + + +GraphInput = CModels.LayerMap | IRModels.Graph +ComponentInput = GraphInput | Mapping[str, GraphInput] + + +def generate( + ir_output: ComponentInput, + output_dir: str | Path | None = None, + *, + model_profile: Any | None = None, + component_name: str | None = None, + weights_dir: str | Path | None = None, + weights_manifest_path: str | Path | None = None, + config: models.GeneratorConfig | None = None, + lowerings: Mapping[str, models.LoweringRule | models.LoweringFn] | None = None, + strict: bool = True, + allow_unsupported_ops: bool = False, +) -> models.GenerationResult: + if config is None and output_dir is None: + raise ValueError("generate requires output_dir unless a GeneratorConfig is provided") + + generator_config = config or models.GeneratorConfig( + output_dir=Path(output_dir), + weights_dir=Path(weights_dir) if weights_dir is not None else None, + weights_manifest_path=Path(weights_manifest_path) if weights_manifest_path is not None else None, + graph_suffix=constants.DEFAULT_GRAPH_SUFFIX, + strict=strict, + allow_unsupported_ops=allow_unsupported_ops, + ) + components = component_graphs_from_input(ir_output, generator_config, component_name, model_profile) + lowering_rules = build_lowering_rules(lowerings) + + for component in components: + lower_component(component, generator_config, lowering_rules) + + return models.GenerationResult.from_components(components) + + +def generate_from_json( + input_path: str | Path, + output_dir: str | Path, + *, + model_profile: Any | None = None, + component_name: str | None = None, + weights_dir: str | Path | None = None, + weights_manifest_path: str | Path | None = None, + strict: bool = True, + allow_unsupported_ops: bool = False, +) -> models.GenerationResult: + layer_map = read_layer_map(input_path) + return generate( + layer_map, + output_dir, + model_profile=model_profile, + component_name=component_name, + weights_dir=weights_dir, + weights_manifest_path=weights_manifest_path, + strict=strict, + allow_unsupported_ops=allow_unsupported_ops, + ) + + +################################################# Generator Utils!!!!!!! ################################################# + + +def read_layer_map(input_path: str | Path) -> CModels.LayerMap: + path = Path(input_path) + return CModels.LayerMap.model_validate_json(path.read_text(encoding="utf-8")) + + +def component_graphs_from_input( + ir_output: ComponentInput, + config: models.GeneratorConfig, + component_name: str | None, + model_profile: Any | None, +) -> tuple[models.ComponentGraph, ...]: + if isinstance(ir_output, Mapping): + return tuple( + models.ComponentGraph.from_ir(name, graph_from_input(component_input), config) + for name, component_input in ir_output.items() + ) + + graph = graph_from_input(ir_output) + name = component_name or default_component_name(graph, model_profile) + return (models.ComponentGraph.from_ir(name, graph, config),) + + +def graph_from_input(ir_output: GraphInput) -> IRModels.Graph: + if isinstance(ir_output, IRModels.Graph): + return ir_output + + if isinstance(ir_output, CModels.LayerMap): + return IRModels.Graph.from_map(ir_output) + + raise TypeError(f"Unsupported generator input: {type(ir_output).__name__}") + + +def default_component_name(graph: IRModels.Graph, model_profile: Any | None) -> str: + profile_name = getattr(model_profile, "model_profiles", None) + model_name = profile_name or graph.model_name or constants.DEFAULT_COMPONENT_NAME + task = graph.task + + if task: + return f"{model_name}_{task}" + + return model_name + + +def create_cactus_graph() -> Any: + try: + from cactus import Graph + except ModuleNotFoundError: + from bindings.cactus import Graph + + return Graph() + + +def component_weight_resolver( + component: models.ComponentGraph, + config: models.GeneratorConfig, +) -> models.WeightResolver | None: + if config.weights_dir is None: + return None + + resolver = models.WeightResolver.from_graph( + component.ir_graph, + config.weights_dir, + config.weights_manifest_path, + ) + + if not resolver.records: + component.warn(f"{component.name}: no converted weight records found in {config.weights_dir}") + + return resolver + + +def build_lowering_rules( + extra_lowerings: Mapping[str, models.LoweringRule | models.LoweringFn] | None = None, +) -> dict[str, models.LoweringRule]: + rules: dict[str, models.LoweringRule] = {} + + add_rules(rules, constants.BINARY_TARGETS, lower_binary) + add_rules(rules, constants.SCALAR_TARGETS, lower_scalar) + add_rules(rules, constants.UNARY_TARGETS, lower_unary) + add_rules(rules, constants.POW_TARGETS, lower_pow) + add_rules(rules, constants.REDUCE_TARGETS, lower_reduce) + add_rules(rules, constants.SHAPE_TARGETS, lower_shape) + add_rules(rules, constants.EXPAND_TARGETS, lower_expand) + add_rules(rules, constants.UNSQUEEZE_TARGETS, lower_unsqueeze) + add_rules(rules, constants.FLATTEN_TARGETS, lower_flatten) + add_rules(rules, constants.TRANSPOSE_TARGETS, lower_transpose) + add_rules(rules, constants.SLICE_TARGETS, lower_slice) + add_rules(rules, constants.INDEX_TARGETS, lower_index) + add_rules(rules, constants.CAT_TARGETS, lower_cat) + add_rules(rules, constants.MATMUL_TARGETS, lower_matmul) + add_rules(rules, constants.ADDM_CONST_TARGETS, lower_addmm) + add_rules(rules, constants.SPLIT_TARGETS, lower_split) + add_rules(rules, constants.GETITEM_TARGETS, lower_getitem) + add_rules(rules, constants.TO_COPY_TARGETS, lower_to_copy) + add_rules(rules, constants.PASS_THROUGH_TARGETS, lower_pass_through) + add_rules(rules, constants.COPY_TARGETS, lower_copy) + add_rules(rules, constants.CONSTANT_INPUT_TARGETS, lower_constant_input) + add_rules(rules, constants.NEG_TARGETS, lower_neg) + add_rules(rules, constants.SOFTMAX_TARGETS, lower_softmax) + add_rules(rules, constants.TOPK_TARGETS, lower_topk) + add_rules(rules, constants.GATHER_TARGETS, lower_gather) + add_rules(rules, constants.EMBEDDING_TARGETS, lower_embedding) + add_rules(rules, constants.CLAMP_TARGETS, lower_clamp) + add_rules(rules, constants.NORM_TARGETS, lower_norm) + add_rules(rules, constants.CONV_TARGETS, lower_conv) + add_rules(rules, constants.ATTENTION_TARGETS, lower_attention) + add_rules(rules, constants.CACHE_TARGETS, lower_cache) + add_rules(rules, constants.MOE_TARGETS, lower_moe) + add_rules(rules, constants.SPECIAL_CACTUS_TARGETS, lower_special_cactus) + add_rules(rules, constants.UNSUPPORTED_SEMANTIC_TARGETS, lower_unsupported_semantic) + + if extra_lowerings: + for target, rule in extra_lowerings.items(): + if isinstance(rule, models.LoweringRule): + rules[target] = rule + else: + rules[target] = models.LoweringRule(target, rule) + + return rules + + +def add_rules(rules: dict[str, models.LoweringRule], targets: Any, lower: models.LoweringFn) -> None: + target_names = targets.keys() if isinstance(targets, dict) else targets + + for target in target_names: + rules[target] = models.LoweringRule(target, lower) + + +def lower_component( + component: models.ComponentGraph, + config: models.GeneratorConfig, + lowering_rules: dict[str, models.LoweringRule], +) -> models.ComponentGraph: + component.graph = create_cactus_graph() + component.weight_resolver = component_weight_resolver(component, config) + context = models.GenerationContext(component=component, config=config, lowerings=lowering_rules) + + for node in IRModels.topological_sort(component.ir_graph): + try: + lower_node(context, node) + except UnsupportedLoweringError as e: + component.mark_unsupported(node) + component.warn(str(e)) + + if config.strict and not config.allow_unsupported_ops: + raise + except KeyError as e: + component.mark_unsupported(node) + component.warn(f"{node.name}: missing lowered input {e}") + + if config.strict and not config.allow_unsupported_ops: + raise UnsupportedLoweringError(f"{node.name}: missing lowered input {e}") from e + + component.save() + return component + + +def lower_node(context: models.GenerationContext, node: IRModels.Node) -> None: + if node.is_placeholder: + context.bind(node, lower_placeholder(context, node)) + return + + if node.is_output: + context.bind(node, lower_output(context, node)) + return + + rule = context.lowering_for(node) + + if rule is None: + raise UnsupportedLoweringError(f"{node.name}: no lowering rule for {node.target}") + + result = rule.lower(context, node) + + if result is not None: + context.bind(node, result) + + +def lower_placeholder(context: models.GenerationContext, node: IRModels.Node) -> Any: + shape, dynamic_dims = graph_input_shape(node) + dtype = cactus_precision(context.graph, tensor_dtype(node)) + tensor = context.graph.input(shape, dtype=dtype, dynamic_dims=dynamic_dims if any(dynamic_dims) else None) + + if node.value_kind in constants.WEIGHT_VALUE_KINDS: + bind_weight_placeholder(context, node, tensor) + return tensor + + context.component.add_runtime_input(tensor) + return tensor + + +def bind_weight_placeholder(context: models.GenerationContext, node: IRModels.Node, tensor: Any) -> None: + resolver = context.component.weight_resolver + + if resolver is None: + context.component.warn(f"{node.name}: weight lowered as graph input because no weights_dir was provided") + return + + source_target = resolver.source_target_for(node.name) or node.target + record = resolver.resolve(node.name) + + if record is None or record.output_name is None: + message = f"{node.name}: could not resolve converted weight for source target {source_target}" + + if context.config.strict and not context.config.allow_unsupported_ops: + raise UnsupportedLoweringError(message) + + context.component.warn(message) + context.component.add_runtime_input(tensor) + return + + node_id = models.tensor_node_id(tensor) + + if node_id is None: + message = f"{node.name}: lowered weight tensor does not expose a Cactus node id" + + if context.config.strict and not context.config.allow_unsupported_ops: + raise UnsupportedLoweringError(message) + + context.component.warn(message) + return + + context.component.add_weight_binding( + models.WeightBinding( + placeholder=node.name, + source_target=source_target, + node_id=node_id, + path=record.output_name, + output_name=record.output_name, + precision=record.precision, + component=record.component, + ) + ) + + +def lower_output(context: models.GenerationContext, node: IRModels.Node) -> Any: + refs = IRModels.extract_node_refs((node.args, node.kwargs)) + values = tuple(context.require(ref) for ref in refs) + + for value in values: + if isinstance(value, (tuple, list)): + for item in value: + context.component.add_output(item) + else: + context.component.add_output(value) + + if len(values) == 1: + return values[0] + + return values + + +def lower_binary(context: models.GenerationContext, node: IRModels.Node) -> Any: + inputs = context.inputs_for(node) + method = constants.BINARY_TARGETS[node.target] + + if len(inputs) == 2: + return getattr(context.graph, method)(inputs[0], inputs[1]) + + scalar = scalar_attr(node, "other") + + if len(inputs) == 1 and scalar is not None: + scalar_method = binary_method_to_scalar_method(method) + return getattr(context.graph, scalar_method)(inputs[0], scalar) + + raise unsupported_arity(node, len(inputs), "two tensor inputs or one tensor plus scalar attr") + + +def lower_scalar(context: models.GenerationContext, node: IRModels.Node) -> Any: + inputs = require_input_count(context, node, 1) + method, attr_name = constants.SCALAR_TARGETS[node.target] + value = scalar_attr(node, attr_name) + + if value is None: + raise UnsupportedLoweringError(f"{node.name}: scalar lowering {node.target} missing attr {attr_name}") + + return getattr(context.graph, method)(inputs[0], value) + + +def lower_unary(context: models.GenerationContext, node: IRModels.Node) -> Any: + inputs = require_input_count(context, node, 1) + method = constants.UNARY_TARGETS[node.target] + return getattr(context.graph, method)(inputs[0]) + + +def lower_pow(context: models.GenerationContext, node: IRModels.Node) -> Any: + inputs = require_input_count(context, node, 1) + exponent = scalar_attr(node, "exponent") + + if exponent is None: + raise UnsupportedLoweringError(f"{node.name}: pow lowering missing exponent attr") + + return context.graph.pow(inputs[0], exponent) + + +def lower_reduce(context: models.GenerationContext, node: IRModels.Node) -> Any: + inputs = require_input_count(context, node, 1) + method = constants.REDUCE_TARGETS[node.target] + axis = axis_attr(node) + + if axis is None: + raise UnsupportedLoweringError(f"{node.name}: reduce lowering missing axis/dim attr") + + return getattr(context.graph, method)(inputs[0], axis) + + +def lower_shape(context: models.GenerationContext, node: IRModels.Node) -> Any: + inputs = require_input_count(context, node, 1) + method = constants.SHAPE_TARGETS[node.target] + shape = shape_attr(node) + + return getattr(context.graph, method)(inputs[0], shape) + + +def lower_expand(context: models.GenerationContext, node: IRModels.Node) -> Any: + inputs = require_input_count(context, node, 1) + return context.graph.expand(inputs[0], shape_attr(node)) + + +def lower_unsqueeze(context: models.GenerationContext, node: IRModels.Node) -> Any: + inputs = require_input_count(context, node, 1) + return context.graph.reshape(inputs[0], output_shape(node)) + + +def lower_flatten(context: models.GenerationContext, node: IRModels.Node) -> Any: + inputs = require_input_count(context, node, 1) + return context.graph.flatten( + inputs[0], + start_dim=int(node.attrs.get("start_dim", 0)), + end_dim=int(node.attrs.get("end_dim", -1)), + ) + + +def lower_transpose(context: models.GenerationContext, node: IRModels.Node) -> Any: + inputs = require_input_count(context, node, 1) + permutation = node.attrs.get("permutation") + + if permutation is not None: + return context.graph.permute(inputs[0], tuple_ints(permutation)) + + if node.target == "aten.t.default": + return context.graph.transpose(inputs[0]) + + dim0 = node.attrs.get("dim0") + dim1 = node.attrs.get("dim1") + + if dim0 is None or dim1 is None: + raise UnsupportedLoweringError(f"{node.name}: transpose lowering missing permutation or dim attrs") + + return context.graph.permute(inputs[0], swap_permutation(parent_rank(node), int(dim0), int(dim1))) + + +def lower_slice(context: models.GenerationContext, node: IRModels.Node) -> Any: + inputs = require_input_count(context, node, 1) + axis = int(node.attrs.get("axis", node.attrs.get("dim", 0))) + + if node.target == "aten.select.int" or "index_value" in node.attrs: + return context.graph.index(inputs[0], int(node.attrs.get("index_value", node.attrs.get("index", 0))), axis=axis) + + start = int(node.attrs.get("start", 0) or 0) + step = node.attrs.get("step", 1) + + if step is not None and int(step) != 1: + raise UnsupportedLoweringError(f"{node.name}: Cactus slice lowering only supports step 1") + + length = slice_length(node, start) + return context.graph.slice(inputs[0], axis, start, length=length) + + +def lower_index(context: models.GenerationContext, node: IRModels.Node) -> Any: + inputs = require_input_count(context, node, 1) + index_value = node.attrs.get("index_value", node.attrs.get("index")) + + if index_value is None: + raise UnsupportedLoweringError(f"{node.name}: index lowering missing index_value/index attr") + + return context.graph.index(inputs[0], int(index_value), axis=int(node.attrs.get("axis", node.attrs.get("dim", 0)))) + + +def lower_cat(context: models.GenerationContext, node: IRModels.Node) -> Any: + inputs = context.inputs_for(node) + + if not inputs: + raise unsupported_arity(node, 0, "at least one input") + + return context.graph.cat(inputs, axis=int(node.attrs.get("axis", node.attrs.get("dim", 0)))) + + +def lower_matmul(context: models.GenerationContext, node: IRModels.Node) -> Any: + inputs = require_input_count(context, node, 2) + return context.graph.matmul( + inputs[0], + inputs[1], + pretransposed_rhs=bool(node.attrs.get("pretransposed_rhs", False)), + ) + + +def lower_addmm(context: models.GenerationContext, node: IRModels.Node) -> Any: + inputs = require_input_count(context, node, 3) + product = context.graph.matmul(inputs[1], inputs[2]) + return context.graph.add(inputs[0], product) + + +def lower_split(context: models.GenerationContext, node: IRModels.Node) -> Any: + inputs = require_input_count(context, node, 1) + split_sizes = node.attrs.get("split_sizes") + + if not isinstance(split_sizes, list): + raise UnsupportedLoweringError(f"{node.name}: split lowering requires split_sizes list") + + axis = int(node.attrs.get("axis", node.attrs.get("dim", 0))) + offset = 0 + outputs = [] + + for size in split_sizes: + size_int = int(size) + outputs.append(context.graph.slice(inputs[0], axis, offset, length=size_int)) + offset += size_int + + return tuple(outputs) + + +def lower_getitem(context: models.GenerationContext, node: IRModels.Node) -> Any: + inputs = require_input_count(context, node, 1) + index = int(node.attrs.get("index", 0)) + source = inputs[0] + + if isinstance(source, (tuple, list)): + return source[index] + + if index == 0: + return source + + raise UnsupportedLoweringError(f"{node.name}: getitem index {index} requires tuple-producing parent") + + +def lower_to_copy(context: models.GenerationContext, node: IRModels.Node) -> Any: + inputs = require_input_count(context, node, 1) + dtype = node.attrs.get("dtype") + + if dtype is None or dtype == "torch.bool": + return inputs[0] + + precision_name = constants.DTYPE_TO_PRECISION.get(str(dtype)) + + if precision_name is None: + return inputs[0] + + return context.graph.precision_cast(inputs[0], cactus_precision(context.graph, str(dtype))) + + +def lower_pass_through(context: models.GenerationContext, node: IRModels.Node) -> Any: + inputs = require_at_least_one_input(context, node) + return inputs[0] + + +def lower_copy(context: models.GenerationContext, node: IRModels.Node) -> Any: + inputs = require_at_least_one_input(context, node) + return inputs[-1] + + +def lower_constant_input(context: models.GenerationContext, node: IRModels.Node) -> Any: + context.component.warn(f"{node.name}: constant-producing op {node.target} lowered as graph input") + shape, dynamic_dims = graph_input_shape(node) + dtype = cactus_precision(context.graph, tensor_dtype(node)) + tensor = context.graph.input(shape, dtype=dtype, dynamic_dims=dynamic_dims if any(dynamic_dims) else None) + context.component.add_runtime_input(tensor) + return tensor + + +def lower_neg(context: models.GenerationContext, node: IRModels.Node) -> Any: + inputs = require_input_count(context, node, 1) + return context.graph.scalar_multiply(inputs[0], -1.0) + + +def lower_softmax(context: models.GenerationContext, node: IRModels.Node) -> Any: + inputs = require_input_count(context, node, 1) + return context.graph.softmax(inputs[0], axis=axis_attr(node, default=-1)) + + +def lower_topk(context: models.GenerationContext, node: IRModels.Node) -> Any: + inputs = require_input_count(context, node, 1) + k = node.attrs.get("k") + + if k is None: + raise UnsupportedLoweringError(f"{node.name}: topk lowering missing k attr") + + return context.graph.topk(inputs[0], int(k)) + + +def lower_gather(context: models.GenerationContext, node: IRModels.Node) -> Any: + inputs = require_input_count(context, node, 2) + return context.graph.gather(inputs[0], inputs[1]) + + +def lower_embedding(context: models.GenerationContext, node: IRModels.Node) -> Any: + inputs = require_input_count(context, node, 2) + return context.graph.embedding_from_tensor(inputs[0], inputs[1]) + + +def lower_clamp(context: models.GenerationContext, node: IRModels.Node) -> Any: + inputs = require_input_count(context, node, 1) + lo = node.attrs.get("lo", node.attrs.get("min")) + hi = node.attrs.get("hi", node.attrs.get("max")) + + if lo is None or hi is None: + raise UnsupportedLoweringError(f"{node.name}: clamp lowering requires min/lo and max/hi attrs") + + return context.graph.clamp(inputs[0], lo, hi) + + +def lower_norm(context: models.GenerationContext, node: IRModels.Node) -> Any: + inputs = context.inputs_for(node) + target = node.target + + if target == "cactus.rms_norm": + require_len(node, inputs, 2) + return context.graph.rms_norm(inputs[0], inputs[1], eps=epsilon_attr(node)) + + if target in {"cactus.layernorm", "cactus.layer_norm", "aten.native_layer_norm.default", "aten.layer_norm.default"}: + require_len(node, inputs, 2) + bias = inputs[2] if len(inputs) > 2 else None + return context.graph.layernorm(inputs[0], inputs[1], bias=bias, eps=epsilon_attr(node)) + + if target in {"cactus.groupnorm", "cactus.group_norm", "aten.native_group_norm.default", "aten.group_norm.default"}: + require_len(node, inputs, 3) + return context.graph.groupnorm(inputs[0], inputs[1], inputs[2], node.attrs.get("num_groups", 1), eps=epsilon_attr(node)) + + if target in {"cactus.batchnorm", "cactus.batch_norm", "aten.native_batch_norm.default", "aten.batch_norm.default"}: + require_len(node, inputs, 5) + return context.graph.batchnorm(inputs[0], inputs[1], inputs[2], inputs[3], inputs[4], axis=int(node.attrs.get("axis", 1)), eps=epsilon_attr(node)) + + raise UnsupportedLoweringError(f"{node.name}: unsupported norm target {target}") + + +def lower_conv(context: models.GenerationContext, node: IRModels.Node) -> Any: + inputs = context.inputs_for(node) + target = node.target + + if target.startswith(constants.CACTUS_TARGET_PREFIX): + method = target.removeprefix(constants.CACTUS_TARGET_PREFIX) + return lower_named_conv(context, node, method, inputs) + + require_len(node, inputs, 2) + bias = inputs[2] if len(inputs) > 2 else None + rank = output_rank(node) + + if target == "aten.conv1d.default" or rank == 3: + return context.graph.conv1d(inputs[0], inputs[1], bias=bias, stride=first_int(node.attrs.get("stride"), 1)) + + if target == "aten.conv2d.default" or rank == 4: + return context.graph.conv2d( + inputs[0], + inputs[1], + bias=bias, + stride=first_int(node.attrs.get("stride"), 1), + padding=first_int(node.attrs.get("padding"), 0), + dilation=first_int(node.attrs.get("dilation"), 1), + groups=int(node.attrs.get("groups", 1)), + ) + + raise UnsupportedLoweringError(f"{node.name}: cannot infer conv rank for {target}") + + +def lower_named_conv(context: models.GenerationContext, node: IRModels.Node, method: str, inputs: tuple[Any, ...]) -> Any: + require_len(node, inputs, 2) + bias = inputs[2] if len(inputs) > 2 else None + + if method == "conv1d": + return context.graph.conv1d(inputs[0], inputs[1], bias=bias, stride=first_int(node.attrs.get("stride"), 1)) + + if method == "conv1d_causal": + return context.graph.conv1d_causal(inputs[0], inputs[1], kernel_size=int(node.attrs.get("kernel_size", 0)), dilation=first_int(node.attrs.get("dilation"), 1)) + + if method == "conv1d_k3": + return context.graph.conv1d_k3(inputs[0], inputs[1], stride=first_int(node.attrs.get("stride"), 1)) + + if method == "conv1d_k7s3": + require_len(node, inputs, 3) + return context.graph.conv1d_k7s3(inputs[0], inputs[1], inputs[2]) + + if method in {"conv1d_same_depthwise_k9", "conv1d_pointwise", "conv2d_k3s2p1", "conv2d_depthwise_k3s2p1", "conv2d_pointwise_1x1", "conv2d_k3s1p1"}: + return getattr(context.graph, method)(inputs[0], inputs[1], bias=bias) + + if method == "conv2d": + return context.graph.conv2d( + inputs[0], + inputs[1], + bias=bias, + stride=first_int(node.attrs.get("stride"), 1), + padding=first_int(node.attrs.get("padding"), 0), + dilation=first_int(node.attrs.get("dilation"), 1), + groups=int(node.attrs.get("groups", 1)), + ) + + raise UnsupportedLoweringError(f"{node.name}: unsupported Cactus conv method {method}") + + +def lower_attention(context: models.GenerationContext, node: IRModels.Node) -> Any: + inputs = context.inputs_for(node) + + if node.target == "cactus.attention" or node.target == "aten.scaled_dot_product_attention.default": + require_len(node, inputs, 3) + mask = inputs[3] if len(inputs) > 3 else None + return context.graph.attention( + inputs[0], + inputs[1], + inputs[2], + scale=float(node.attrs.get("scale", 1.0)), + is_causal=bool(node.attrs.get("is_causal", True)), + position_offset=int(node.attrs.get("position_offset", 0)), + window_size=int(node.attrs.get("window_size", 0)), + mask=mask, + additive_mask=bool(node.attrs.get("additive_mask", False)), + ) + + if node.target == "cactus.attention_cached": + require_len(node, inputs, 5) + return context.graph.attention_cached( + inputs[0], + inputs[1], + inputs[2], + inputs[3], + inputs[4], + scale=float(node.attrs.get("scale", 1.0)), + position_offset=int(node.attrs.get("position_offset", 0)), + window_size=int(node.attrs.get("window_size", 0)), + v_head_dim=int(node.attrs.get("v_head_dim", 0)), + ) + + raise UnsupportedLoweringError(f"{node.name}: unsupported attention target {node.target}") + + +def lower_cache(context: models.GenerationContext, node: IRModels.Node) -> Any: + inputs = context.inputs_for(node) + target = node.target + + if target == "cactus.kv_cache_append": + require_len(node, inputs, 2) + return context.graph.kv_cache_append(inputs[0], inputs[1], window_size=int(node.attrs.get("window_size", 0)), sink_size=int(node.attrs.get("sink_size", 0))) + + if target == "cactus.conv_cache_append": + require_len(node, inputs, 2) + return context.graph.conv_cache_append(inputs[0], inputs[1]) + + if target == "cactus.conv_cache_initialize": + require_len(node, inputs, 2) + return context.graph.conv_cache_initialize(inputs[0], inputs[1]) + + if target == "cactus.recurrent_cache_write": + require_len(node, inputs, 2) + return context.graph.recurrent_cache_write(inputs[0], inputs[1]) + + if target == "cactus.kv_cache_state": + return context.graph.kv_cache_state( + required_int_attr(node, "max_seq_len"), + required_int_attr(node, "num_kv_heads"), + required_int_attr(node, "head_dim"), + window_size=int(node.attrs.get("window_size", 0)), + sink_size=int(node.attrs.get("sink_size", 0)), + num_slots=int(node.attrs.get("num_slots", 1)), + ) + + if target == "cactus.conv_cache_state": + return context.graph.conv_cache_state(required_int_attr(node, "window_size"), required_int_attr(node, "hidden_dim")) + + if target == "cactus.recurrent_cache_state": + return context.graph.recurrent_cache_state(shape_attr(node), dtype=cactus_precision(context.graph, tensor_dtype(node))) + + raise UnsupportedLoweringError(f"{node.name}: unsupported cache target {target}") + + +def lower_moe(context: models.GenerationContext, node: IRModels.Node) -> Any: + inputs = context.inputs_for(node) + + if node.target == "cactus.dense_mlp_tq_fused": + require_len(node, inputs, 4) + return context.graph.dense_mlp_tq_fused(inputs[0], inputs[1], inputs[2], inputs[3], product_scale=float(node.attrs.get("product_scale", 1.0))) + + if node.target == "cactus.moe_layer_gated": + if len(inputs) < 6: + raise UnsupportedLoweringError( + f"{node.name}: cactus.moe_layer_gated requires hidden, routing_probs, topk_indices, w1, w3, and w2 inputs; " + "packed/grouped MoE lowering still needs a weight-unpacking policy" + ) + + return context.graph.moe_layer_gated( + inputs[0], + inputs[1], + inputs[2], + ensure_tensor_sequence(inputs[3]), + ensure_tensor_sequence(inputs[4]), + ensure_tensor_sequence(inputs[5]), + required_int_attr(node, "num_experts"), + required_int_attr(node, "num_experts_per_tok"), + normalize_routing=bool(node.attrs.get("normalize_routing", True)), + epsilon=float(node.attrs.get("epsilon", 1e-6)), + routed_scaling_factor=float(node.attrs.get("routed_scaling_factor", 1.0)), + ) + + if node.target == "cactus.moe_layer_ungated": + if len(inputs) < 5: + raise unsupported_arity(node, len(inputs), "hidden, routing_probs, topk_indices, w1, and w2") + + return context.graph.moe_layer_ungated( + inputs[0], + inputs[1], + inputs[2], + ensure_tensor_sequence(inputs[3]), + ensure_tensor_sequence(inputs[4]), + required_int_attr(node, "num_experts"), + required_int_attr(node, "num_experts_per_tok"), + normalize_routing=bool(node.attrs.get("normalize_routing", True)), + epsilon=float(node.attrs.get("epsilon", 1e-6)), + routed_scaling_factor=float(node.attrs.get("routed_scaling_factor", 1.0)), + ) + + raise UnsupportedLoweringError(f"{node.name}: unsupported MoE target {node.target}") + + +def lower_special_cactus(context: models.GenerationContext, node: IRModels.Node) -> Any: + inputs = context.inputs_for(node) + target = node.target + + if target == "cactus.rope": + require_len(node, inputs, 1) + return context.graph.rope(inputs[0], theta=float(node.attrs.get("theta", 10_000.0)), position_offset=int(node.attrs.get("position_offset", 0))) + + if target == "cactus.rope_gptj": + require_len(node, inputs, 1) + return context.graph.rope_gptj(inputs[0], theta=float(node.attrs.get("theta", 10_000.0)), position_offset=int(node.attrs.get("position_offset", 0)), rot_dim=int(node.attrs.get("rot_dim", 0))) + + if target == "cactus.glu": + require_len(node, inputs, 1) + return context.graph.glu(inputs[0], axis=axis_attr(node, default=-1)) + + if target == "cactus.lstm_cell": + require_len(node, inputs, 7) + return context.graph.lstm_cell(*inputs[:7]) + + if target == "cactus.gated_deltanet_decode": + require_len(node, inputs, 6) + return context.graph.gated_deltanet_decode(*inputs[:6], scale=float(node.attrs.get("scale", 1.0))) + + if target == "cactus.gated_deltanet_prefill": + require_len(node, inputs, 6) + return context.graph.gated_deltanet_prefill(*inputs[:6], chunk_size=int(node.attrs.get("chunk_size", 1)), scale=float(node.attrs.get("scale", 1.0))) + + if target == "cactus.rel_pos_bias": + require_len(node, inputs, 2) + return context.graph.rel_pos_bias(inputs[0], inputs[1], scale=float(node.attrs.get("scale", 1.0))) + + if target == "cactus.sample": + require_len(node, inputs, 1) + return context.graph.sample(inputs[0], temperature=float(node.attrs.get("temperature", 0.6)), top_p=float(node.attrs.get("top_p", 0.95)), top_k=int(node.attrs.get("top_k", 20))) + + if target == "cactus.scatter_topk": + require_len(node, inputs, 2) + return context.graph.scatter_topk(inputs[0], inputs[1], num_classes=required_int_attr(node, "num_classes")) + + if target == "cactus.gaussian_topk": + require_len(node, inputs, 1) + return context.graph.gaussian_topk(inputs[0], ppf=float(node.attrs.get("ppf", 0.0))) + + if target == "cactus.altup_predict": + require_len(node, inputs, 2) + return context.graph.altup_predict(inputs[0], list(inputs[1:])) + + if target == "cactus.altup_correct": + require_len(node, inputs, 3) + return context.graph.altup_correct(inputs[0], inputs[1], list(inputs[2:])) + + if target == "cactus.stft": + require_len(node, inputs, 2) + return context.graph.stft(inputs[0], inputs[1], stride=required_int_attr(node, "stride"), num_fft_bins=required_int_attr(node, "num_fft_bins")) + + if target == "cactus.rfft": + require_len(node, inputs, 1) + return context.graph.rfft(inputs[0]) + + if target == "cactus.irfft": + require_len(node, inputs, 1) + return context.graph.irfft(inputs[0], output_length=required_int_attr(node, "output_length")) + + if target == "cactus.spectrogram": + require_len(node, inputs, 2) + return context.graph.spectrogram( + inputs[0], + inputs[1], + frame_length=required_int_attr(node, "frame_length"), + hop_length=required_int_attr(node, "hop_length"), + fft_length=required_int_attr(node, "fft_length"), + power=float(node.attrs.get("power", 2.0)), + ) + + if target == "cactus.image_preprocess": + require_len(node, inputs, 1) + return context.graph.image_preprocess( + inputs[0], + src_width=required_int_attr(node, "src_width"), + src_height=required_int_attr(node, "src_height"), + target_width=required_int_attr(node, "target_width"), + target_height=required_int_attr(node, "target_height"), + patch_size=required_int_attr(node, "patch_size"), + channels=required_int_attr(node, "channels"), + rescale_factor=float(node.attrs.get("rescale_factor", 1.0)), + mean=node.attrs.get("mean", (0.0, 0.0, 0.0)), + std_dev=node.attrs.get("std_dev", (1.0, 1.0, 1.0)), + ) + + if target == "cactus.bilinear_interpolation": + require_len(node, inputs, 1) + return context.graph.bilinear_interpolation(inputs[0], required_int_attr(node, "dst_height"), required_int_attr(node, "dst_width")) + + raise UnsupportedLoweringError(f"{node.name}: unsupported special Cactus target {target}") + + +def lower_unsupported_semantic(context: models.GenerationContext, node: IRModels.Node) -> Any: + raise UnsupportedLoweringError(f"{node.name}: {node.target} has no safe Cactus lowering yet") + + +def require_input_count(context: models.GenerationContext, node: IRModels.Node, min_count: int) -> tuple[Any, ...]: + inputs = context.inputs_for(node) + require_len(node, inputs, min_count) + return inputs + + +def require_at_least_one_input(context: models.GenerationContext, node: IRModels.Node) -> tuple[Any, ...]: + return require_input_count(context, node, 1) + + +def require_len(node: IRModels.Node, inputs: tuple[Any, ...], min_count: int) -> None: + if len(inputs) < min_count: + raise unsupported_arity(node, len(inputs), f"at least {min_count} inputs") + + +def unsupported_arity(node: IRModels.Node, actual: int, expected: str) -> UnsupportedLoweringError: + return UnsupportedLoweringError(f"{node.name}: {node.target} got {actual} lowered inputs; expected {expected}") + + +def tensor_dtype(node: IRModels.Node) -> str | None: + if isinstance(node.tensor_output_meta, dict): + dtype = node.tensor_output_meta.get("dtype") + + if dtype is not None: + return str(dtype) + + return None + + +def cactus_precision(graph: Any, dtype: str | None) -> int: + precision_name = constants.DTYPE_TO_PRECISION.get(str(dtype), constants.DEFAULT_INPUT_PRECISION) + return int(getattr(graph, precision_name)) + + +def graph_input_shape(node: IRModels.Node) -> tuple[tuple[int, ...], tuple[bool, ...]]: + shape = meta_shape(node) + + if not shape: + raw_shape = node.tensor_output_meta.get("shape") if isinstance(node.tensor_output_meta, dict) else None + if raw_shape == [] or raw_shape == (): + return (1,), (False,) + + return (), () + + dims: list[int] = [] + dynamic_dims: list[bool] = [] + + for dim in shape: + if isinstance(dim, int) and dim >= 0: + dims.append(dim) + dynamic_dims.append(False) + continue + + if isinstance(dim, str) and dim.isdigit(): + dims.append(int(dim)) + dynamic_dims.append(False) + continue + + dims.append(1) + dynamic_dims.append(True) + + return tuple(dims), tuple(dynamic_dims) + + +def meta_shape(node: IRModels.Node) -> tuple[Any, ...]: + if isinstance(node.tensor_output_meta, dict): + shape = node.tensor_output_meta.get("shape") + + if isinstance(shape, list): + return tuple(shape) + + if isinstance(shape, tuple): + return shape + + return () + + +def output_shape(node: IRModels.Node) -> tuple[int, ...]: + shape, _ = graph_input_shape(node) + + if shape: + return shape + + raise UnsupportedLoweringError(f"{node.name}: missing concrete output shape") + + +def shape_attr(node: IRModels.Node) -> tuple[int, ...]: + raw_shape = node.attrs.get("shape") + + if raw_shape is None: + return output_shape(node) + + resolved_output_shape = output_shape(node) + dims: list[int] = [] + + for index, dim in enumerate(raw_shape): + if isinstance(dim, int) and dim >= 0: + dims.append(dim) + elif isinstance(dim, str) and dim.isdigit(): + dims.append(int(dim)) + elif index < len(resolved_output_shape): + dims.append(resolved_output_shape[index]) + else: + raise UnsupportedLoweringError(f"{node.name}: cannot resolve shape dim {dim!r}") + + return tuple(dims) + + +def axis_attr(node: IRModels.Node, default: int | None = None) -> int | None: + axis = node.attrs.get("axis", node.attrs.get("dim", default)) + + if isinstance(axis, list): + if len(axis) != 1: + raise UnsupportedLoweringError(f"{node.name}: Cactus reduction lowering only supports one axis") + return int(axis[0]) + + if axis is None: + return None + + return int(axis) + + +def scalar_attr(node: IRModels.Node, name: str) -> Any | None: + value = node.attrs.get(name) + + if isinstance(value, dict) and "node" in value: + return None + + if isinstance(value, list) and len(value) == 1: + return value[0] + + return value + + +def epsilon_attr(node: IRModels.Node) -> float: + return float(node.attrs.get("epsilon", node.attrs.get("eps", 1e-5))) + + +def required_int_attr(node: IRModels.Node, name: str) -> int: + value = node.attrs.get(name) + + if value is None: + raise UnsupportedLoweringError(f"{node.name}: missing required attr {name}") + + return int(value) + + +def first_int(value: Any, default: int) -> int: + if value is None: + return default + + if isinstance(value, (list, tuple)): + if not value: + return default + return int(value[0]) + + return int(value) + + +def tuple_ints(value: Any) -> tuple[int, ...]: + if not isinstance(value, (list, tuple)): + raise TypeError(f"Expected list/tuple of ints, got {type(value).__name__}") + + return tuple(int(item) for item in value) + + +def parent_rank(node: IRModels.Node) -> int: + if not node.parents: + raise UnsupportedLoweringError(f"{node.name}: cannot infer parent rank without parents") + + rank = len(meta_shape(node.parents[0])) + + if rank == 0: + raise UnsupportedLoweringError(f"{node.name}: cannot infer parent rank without parent shape metadata") + + return rank + + +def output_rank(node: IRModels.Node) -> int: + return len(meta_shape(node)) + + +def swap_permutation(rank: int, dim0: int, dim1: int) -> tuple[int, ...]: + dim0 = normalize_dim(dim0, rank) + dim1 = normalize_dim(dim1, rank) + permutation = list(range(rank)) + permutation[dim0], permutation[dim1] = permutation[dim1], permutation[dim0] + return tuple(permutation) + + +def normalize_dim(dim: int, rank: int) -> int: + if dim < 0: + dim += rank + + if dim < 0 or dim >= rank: + raise ValueError(f"Dimension {dim} is outside rank {rank}") + + return dim + + +def slice_length(node: IRModels.Node, start: int) -> int: + if "length" in node.attrs: + return int(node.attrs["length"]) + + end = node.attrs.get("end") + + if end is None: + return 0 + + end = int(end) + + if end >= constants.OPEN_SLICE_END: + return 0 + + return max(end - start, 0) + + +def binary_method_to_scalar_method(method: str) -> str: + return { + "add": "scalar_add", + "subtract": "scalar_subtract", + "multiply": "scalar_multiply", + "divide": "scalar_divide", + "not_equal": "scalar_not_equal", + }[method] + + +def ensure_tensor_sequence(value: Any) -> list[Any]: + if isinstance(value, list): + return value + + if isinstance(value, tuple): + return list(value) + + return [value] + + +def dump_result_manifest(result: models.GenerationResult, output_path: str | Path) -> Path: + path = Path(output_path) + path.parent.mkdir(parents=True, exist_ok=True) + path.write_text( + json.dumps( + { + "component_paths": {name: str(path_) for name, path_ in result.component_paths.items()}, + "component_manifest_paths": {name: str(path_) for name, path_ in result.component_manifest_paths.items()}, + "unsupported_nodes": list(result.unsupported_nodes), + "warnings": list(result.warnings), + "ok": result.ok, + }, + indent=4, + ), + encoding="utf-8", + ) + return path diff --git a/python/cactus/transpiler/Generator/models.py b/python/cactus/transpiler/Generator/models.py index e69de29bb..2728e428f 100644 --- a/python/cactus/transpiler/Generator/models.py +++ b/python/cactus/transpiler/Generator/models.py @@ -0,0 +1,533 @@ +from __future__ import annotations + +import json +import re +from dataclasses import dataclass, field +from pathlib import Path +from typing import Any, Callable, Mapping + +from ..IR import models as IRModels + + +GraphTensor = Any +CactusGraph = Any +LoweringFn = Callable[["GenerationContext", IRModels.Node], GraphTensor | tuple[GraphTensor, ...] | None] + + +@dataclass(slots=True, frozen=True) +class TensorSpec: + name: str + shape: tuple[Any, ...] = () + dtype: str | None = None + dynamic_dims: tuple[bool, ...] = () + + @classmethod + def from_node(cls, node: IRModels.Node, *, name: str | None = None) -> "TensorSpec": + return tensor_spec_from_node(cls, node, name=name) + + +@dataclass(slots=True, frozen=True) +class LoweringRule: + target: str + lower: LoweringFn + + +@dataclass(slots=True, frozen=True) +class GeneratorConfig: + output_dir: Path + weights_dir: Path | None = None + weights_manifest_path: Path | None = None + graph_suffix: str = ".cactus" + strict: bool = True + allow_unsupported_ops: bool = False + + def component_path(self, component_name: str) -> Path: + return generator_component_path(self, component_name) + + +@dataclass(slots=True) +class ComponentGraph: + name: str + ir_graph: IRModels.Graph + output_path: Path + manifest_path: Path + graph: CactusGraph | None = None + weight_resolver: "WeightResolver | None" = None + weight_bindings: list["WeightBinding"] = field(default_factory=list) + runtime_input_ids: list[int] = field(default_factory=list) + output_node_ids: list[int] = field(default_factory=list) + unsupported_nodes: list[str] = field(default_factory=list) + warnings: list[str] = field(default_factory=list) + + @classmethod + def from_ir(cls, name: str, ir_graph: IRModels.Graph, config: GeneratorConfig) -> "ComponentGraph": + return component_graph_from_ir(cls, name, ir_graph, config) + + def save(self) -> Path: + return save_component_graph(self) + + def save_manifest(self) -> Path: + return save_component_manifest(self) + + def mark_unsupported(self, node: IRModels.Node) -> None: + mark_component_node_unsupported(self, node) + + def warn(self, message: str) -> None: + add_component_warning(self, message) + + def add_weight_binding(self, binding: "WeightBinding") -> None: + add_component_weight_binding(self, binding) + + def add_runtime_input(self, tensor: GraphTensor) -> None: + add_component_runtime_input(self, tensor) + + def add_output(self, tensor: GraphTensor) -> None: + add_component_output(self, tensor) + + +@dataclass(slots=True) +class GenerationContext: + component: ComponentGraph + config: GeneratorConfig + lowerings: dict[str, LoweringRule] = field(default_factory=dict) + values: dict[str, GraphTensor] = field(default_factory=dict) + + @property + def graph(self) -> CactusGraph: + return active_cactus_graph(self) + + def bind(self, node: IRModels.Node, value: GraphTensor) -> GraphTensor: + return bind_context_value(self, node, value) + + def lookup(self, node_name: str) -> GraphTensor | None: + return lookup_context_value(self, node_name) + + def require(self, node_name: str) -> GraphTensor: + return require_context_value(self, node_name) + + def inputs_for(self, node: IRModels.Node) -> tuple[GraphTensor, ...]: + return context_inputs_for_node(self, node) + + def lowering_for(self, node: IRModels.Node) -> LoweringRule | None: + return context_lowering_for_node(self, node) + + def mark_unsupported(self, node: IRModels.Node) -> None: + self.component.mark_unsupported(node) + + +@dataclass(slots=True, frozen=True) +class WeightRecord: + source_name: str | None = None + hf_name: str | None = None + adapter_name: str | None = None + output_name: str | None = None + shape: tuple[Any, ...] = () + precision: str | None = None + status: str | None = None + component: str | None = None + scale_factor: float = 1.0 + adapter_family: str | None = None + source_names: tuple[str, ...] = () + transform: str = "none" + qdq_restore: str = "hf_key" + + @classmethod + def from_manifest_row(cls, row: Mapping[str, Any]) -> "WeightRecord": + return weight_record_from_manifest_row(cls, row) + + @property + def aliases(self) -> tuple[str, ...]: + return weight_record_aliases(self) + + +@dataclass(slots=True, frozen=True) +class WeightBinding: + placeholder: str + source_target: str + node_id: int + path: str + output_name: str + precision: str | None = None + component: str | None = None + binding_kind: str = "mmap_weight" + + +@dataclass(slots=True) +class WeightResolver: + weights_dir: Path + records: tuple[WeightRecord, ...] = () + records_by_name: dict[str, WeightRecord] = field(default_factory=dict) + placeholder_targets: dict[str, str] = field(default_factory=dict) + + @classmethod + def from_graph(cls, graph: IRModels.Graph, weights_dir: Path, manifest_path: Path | None = None) -> "WeightResolver": + return weight_resolver_from_graph(cls, graph, weights_dir, manifest_path) + + def resolve(self, placeholder_name: str) -> WeightRecord | None: + return resolve_weight_record(self, placeholder_name) + + def source_target_for(self, placeholder_name: str) -> str | None: + return self.placeholder_targets.get(placeholder_name) + + +@dataclass(slots=True, frozen=True) +class ComponentGraphManifest: + component: str + graph_path: str + weight_bindings: tuple[WeightBinding, ...] = () + runtime_input_node_ids: tuple[int, ...] = () + output_node_ids: tuple[int, ...] = () + unsupported_nodes: tuple[str, ...] = () + warnings: tuple[str, ...] = () + + @classmethod + def from_component(cls, component: ComponentGraph) -> "ComponentGraphManifest": + return component_manifest_from_component(cls, component) + + def to_dict(self) -> dict[str, Any]: + return component_manifest_to_dict(self) + + +@dataclass(slots=True, frozen=True) +class GenerationResult: + component_paths: dict[str, Path] = field(default_factory=dict) + component_manifest_paths: dict[str, Path] = field(default_factory=dict) + unsupported_nodes: tuple[str, ...] = () + warnings: tuple[str, ...] = () + + @classmethod + def from_components(cls, components: tuple[ComponentGraph, ...] | list[ComponentGraph]) -> "GenerationResult": + return generation_result_from_components(cls, components) + + @property + def ok(self) -> bool: + return generation_result_ok(self) + + +################################################# Model Utils!!!!!!! ################################################# + + +def tensor_spec_from_node(cls: type[TensorSpec], node: IRModels.Node, *, name: str | None = None) -> TensorSpec: + meta = node.tensor_output_meta if isinstance(node.tensor_output_meta, dict) else {} + shape = meta.get("shape", ()) + + return cls( + name=name or node.name, + shape=tuple(shape or ()), + dtype=meta.get("dtype"), + ) + + +def generator_component_path(config: GeneratorConfig, component_name: str) -> Path: + safe_name = sanitize_component_name(component_name) + return config.output_dir / f"{safe_name}{config.graph_suffix}" + + +def generator_component_manifest_path(config: GeneratorConfig, component_name: str) -> Path: + safe_name = sanitize_component_name(component_name) + return config.output_dir / f"{safe_name}.graph_manifest.json" + + +def component_graph_from_ir(cls: type[ComponentGraph], name: str, ir_graph: IRModels.Graph, config: GeneratorConfig) -> ComponentGraph: + return cls( + name=name, + ir_graph=ir_graph, + output_path=config.component_path(name), + manifest_path=generator_component_manifest_path(config, name), + ) + + +def save_component_graph(component: ComponentGraph) -> Path: + if component.graph is None: + raise ValueError(f"Component {component.name} has no Cactus graph to save") + + component.output_path.parent.mkdir(parents=True, exist_ok=True) + component.graph.save(component.output_path) + component.save_manifest() + return component.output_path + + +def save_component_manifest(component: ComponentGraph) -> Path: + component.manifest_path.parent.mkdir(parents=True, exist_ok=True) + manifest = ComponentGraphManifest.from_component(component) + component.manifest_path.write_text(json.dumps(manifest.to_dict(), indent=4), encoding="utf-8") + return component.manifest_path + + +def mark_component_node_unsupported(component: ComponentGraph, node: IRModels.Node) -> None: + if node.name not in component.unsupported_nodes: + component.unsupported_nodes.append(node.name) + + +def add_component_warning(component: ComponentGraph, message: str) -> None: + if message not in component.warnings: + component.warnings.append(message) + + +def add_component_weight_binding(component: ComponentGraph, binding: WeightBinding) -> None: + if all(existing.node_id != binding.node_id for existing in component.weight_bindings): + component.weight_bindings.append(binding) + + +def add_component_runtime_input(component: ComponentGraph, tensor: GraphTensor) -> None: + node_id = tensor_node_id(tensor) + + if node_id is not None and node_id not in component.runtime_input_ids: + component.runtime_input_ids.append(node_id) + + +def add_component_output(component: ComponentGraph, tensor: GraphTensor) -> None: + node_id = tensor_node_id(tensor) + + if node_id is not None and node_id not in component.output_node_ids: + component.output_node_ids.append(node_id) + + +def tensor_node_id(tensor: GraphTensor) -> int | None: + node_id = getattr(tensor, "id", None) + + if node_id is None: + return None + + return int(node_id) + + +def active_cactus_graph(context: GenerationContext) -> CactusGraph: + if context.component.graph is None: + raise ValueError(f"Component {context.component.name} has no active Cactus graph") + + return context.component.graph + + +def bind_context_value(context: GenerationContext, node: IRModels.Node, value: GraphTensor) -> GraphTensor: + context.values[node.name] = value + return value + + +def lookup_context_value(context: GenerationContext, node_name: str) -> GraphTensor | None: + return context.values.get(node_name) + + +def require_context_value(context: GenerationContext, node_name: str) -> GraphTensor: + value = context.lookup(node_name) + + if value is None: + raise KeyError(f"Node {node_name} has not been lowered in component {context.component.name}") + + return value + + +def context_inputs_for_node(context: GenerationContext, node: IRModels.Node) -> tuple[GraphTensor, ...]: + return tuple(context.require(parent.name) for parent in node.parents) + + +def context_lowering_for_node(context: GenerationContext, node: IRModels.Node) -> LoweringRule | None: + return context.lowerings.get(node.target) + + +def generation_result_from_components(cls: type[GenerationResult], components: tuple[ComponentGraph, ...] | list[ComponentGraph]) -> GenerationResult: + return cls( + component_paths={component.name: component.output_path for component in components}, + component_manifest_paths={component.name: component.manifest_path for component in components}, + unsupported_nodes=tuple( + unsupported + for component in components + for unsupported in component.unsupported_nodes + ), + warnings=tuple( + warning + for component in components + for warning in component.warnings + ), + ) + + +def generation_result_ok(result: GenerationResult) -> bool: + return not result.unsupported_nodes + + +def sanitize_component_name(name: str) -> str: + safe = "".join(char if char.isalnum() or char in {"_", "-"} else "_" for char in name) + return safe.strip("_-") or "component" + + +def weight_record_from_manifest_row(cls: type[WeightRecord], row: Mapping[str, Any]) -> WeightRecord: + return cls( + source_name=none_or_str(row.get("source_name")), + hf_name=none_or_str(row.get("hf_name")), + adapter_name=none_or_str(row.get("adapter_name")), + output_name=none_or_str(row.get("output_name") or row.get("output_file")), + shape=tuple(row.get("shape") or ()), + precision=none_or_str(row.get("precision")), + status=none_or_str(row.get("status")), + component=none_or_str(row.get("component")), + scale_factor=float(row.get("scale_factor", 1.0) or 1.0), + adapter_family=none_or_str(row.get("adapter_family")), + source_names=tuple(str(name) for name in row.get("source_names") or ()), + transform=str(row.get("transform", "none") or "none"), + qdq_restore=str(row.get("qdq_restore", "hf_key") or "hf_key"), + ) + + +def weight_record_aliases(record: WeightRecord) -> tuple[str, ...]: + aliases = ( + record.source_name, + record.hf_name, + record.adapter_name, + *record.source_names, + ) + return tuple(alias for alias in unique_strings(aliases) if alias) + + +def weight_resolver_from_graph( + cls: type[WeightResolver], + graph: IRModels.Graph, + weights_dir: Path, + manifest_path: Path | None = None, +) -> WeightResolver: + records = load_weight_records(manifest_path or weights_dir / "weights_manifest.json") + records_by_name = index_weight_records(records) + placeholder_targets = input_spec_placeholder_targets(graph) + + return cls( + weights_dir=weights_dir, + records=records, + records_by_name=records_by_name, + placeholder_targets=placeholder_targets, + ) + + +def load_weight_records(manifest_path: Path) -> tuple[WeightRecord, ...]: + if not manifest_path.exists(): + return () + + data = json.loads(manifest_path.read_text(encoding="utf-8")) + rows = data if isinstance(data, list) else data.get("weights", ()) + return tuple(WeightRecord.from_manifest_row(row) for row in rows if isinstance(row, Mapping)) + + +def index_weight_records(records: tuple[WeightRecord, ...]) -> dict[str, WeightRecord]: + index: dict[str, WeightRecord] = {} + + for record in records: + if not record.output_name: + continue + + for alias in record.aliases: + for normalized in weight_name_variants(alias): + index.setdefault(normalized, record) + + return index + + +def input_spec_placeholder_targets(graph: IRModels.Graph) -> dict[str, str]: + targets: dict[str, str] = {} + + for spec in graph.input_specs: + arg_name = getattr(spec, "arg_name", None) + target = getattr(spec, "target", None) + + if arg_name and target: + targets[str(arg_name)] = str(target) + + if targets: + return targets + + return parse_graph_signature_placeholder_targets(graph.graph_signature) + + +def parse_graph_signature_placeholder_targets(graph_signature: str) -> dict[str, str]: + targets: dict[str, str] = {} + pattern = re.compile( + r"InputSpec\(kind=]+>, " + r"arg=TensorArgument\(name='([^']+)'\), target='([^']+)'" + ) + + for match in pattern.finditer(graph_signature): + targets[match.group(1)] = match.group(2) + + return targets + + +def resolve_weight_record(resolver: WeightResolver, placeholder_name: str) -> WeightRecord | None: + source_target = resolver.placeholder_targets.get(placeholder_name) + + if source_target is None: + return None + + for variant in weight_name_variants(source_target): + record = resolver.records_by_name.get(variant) + + if record is not None: + return record + + return None + + +def weight_name_variants(name: str) -> tuple[str, ...]: + variants = [name] + + current = name + while current.startswith("model."): + current = current[len("model."):] + variants.append(current) + + if name.startswith("_orig_mod."): + variants.append(name[len("_orig_mod."):]) + + return tuple(unique_strings(variants)) + + +def component_manifest_from_component(cls: type[ComponentGraphManifest], component: ComponentGraph) -> ComponentGraphManifest: + return cls( + component=component.name, + graph_path=component.output_path.name, + weight_bindings=tuple(component.weight_bindings), + runtime_input_node_ids=tuple(component.runtime_input_ids), + output_node_ids=tuple(component.output_node_ids), + unsupported_nodes=tuple(component.unsupported_nodes), + warnings=tuple(component.warnings), + ) + + +def component_manifest_to_dict(manifest: ComponentGraphManifest) -> dict[str, Any]: + return { + "component": manifest.component, + "graph_path": manifest.graph_path, + "weight_bindings": [ + { + "placeholder": binding.placeholder, + "source_target": binding.source_target, + "node_id": binding.node_id, + "path": binding.path, + "output_name": binding.output_name, + "precision": binding.precision, + "component": binding.component, + "binding_kind": binding.binding_kind, + } + for binding in manifest.weight_bindings + ], + "runtime_input_node_ids": list(manifest.runtime_input_node_ids), + "output_node_ids": list(manifest.output_node_ids), + "unsupported_nodes": list(manifest.unsupported_nodes), + "warnings": list(manifest.warnings), + } + + +def none_or_str(value: Any) -> str | None: + if value is None: + return None + + return str(value) + + +def unique_strings(values: tuple[str | None, ...] | list[str | None]) -> tuple[str, ...]: + unique: list[str] = [] + + for value in values: + if value is None or value in unique: + continue + + unique.append(value) + + return tuple(unique) diff --git a/python/cactus/transpiler/IR/models.py b/python/cactus/transpiler/IR/models.py index f3ae40063..74084fcd2 100644 --- a/python/cactus/transpiler/IR/models.py +++ b/python/cactus/transpiler/IR/models.py @@ -51,6 +51,8 @@ class Graph: task: str = "" graph_signature: str = "" range_constants: str = "" + input_specs: tuple[CModels.GraphSpecRecord, ...] = field(default_factory=tuple) + output_specs: tuple[CModels.GraphSpecRecord, ...] = field(default_factory=tuple) fusions: list["FusionResult"] = field(default_factory=list) @classmethod @@ -242,6 +244,8 @@ def generate_graph(cls: type[Graph], layer_map: CModels.LayerMap) -> Graph: task=layer_map.task, graph_signature=layer_map.graph_signature, range_constants=layer_map.range_constants, + input_specs=tuple(layer_map.input_specs), + output_specs=tuple(layer_map.output_specs), ) @@ -417,6 +421,8 @@ def rebuild_graph(nodes: tuple[Node, ...], original_graph: Graph, fusion_results task=original_graph.task, graph_signature=original_graph.graph_signature, range_constants=original_graph.range_constants, + input_specs=original_graph.input_specs, + output_specs=original_graph.output_specs, fusions=list(fusion_results), ) @@ -442,6 +448,8 @@ def graph_to_layer_map(graph: Graph) -> CModels.LayerMap: task=graph.task, graph_signature=graph.graph_signature, range_constants=graph.range_constants, + input_specs=list(graph.input_specs), + output_specs=list(graph.output_specs), nodes=records, ) From ddcbcc868c22af4c8816332d48314dedabfc836e Mon Sep 17 00:00:00 2001 From: Parkirat Date: Wed, 29 Jul 2026 10:49:36 -0400 Subject: [PATCH 62/84] Further LFM support --- python/cactus/transpiler/Fusions/edges.py | 88 +++++++ python/cactus/transpiler/Fusions/fusions.py | 141 ++++++++++- python/cactus/transpiler/Fusions/nodes.py | 85 ++++++- .../cactus/transpiler/Generator/constants.py | 5 +- .../cactus/transpiler/Generator/generate.py | 236 ++++++++++++++++-- python/cactus/transpiler/IR/extra_matchers.py | 82 ++++++ python/cactus/transpiler/IR/match_utils.py | 10 +- python/cactus/transpiler/IR/models.py | 29 ++- python/cactus/transpiler/IR/simplify_ir.py | 3 + 9 files changed, 641 insertions(+), 38 deletions(-) diff --git a/python/cactus/transpiler/Fusions/edges.py b/python/cactus/transpiler/Fusions/edges.py index e8ed2d806..a32dbcd63 100644 --- a/python/cactus/transpiler/Fusions/edges.py +++ b/python/cactus/transpiler/Fusions/edges.py @@ -91,6 +91,30 @@ def _index_edges_by_dest(edges: dict[str, M.FusionEdge]) -> dict[str, tuple[str, "attn_mask_add_to_softmax": _edge("attn_mask_add", "attn_softmax", 0), "attn_softmax_to_value": _edge("attn_softmax", "attn_value", 0), "v_view_to_value": _edge("v_view", "attn_value", 1), + "lfm_attn_query_expand_to_flat": _edge("lfm_attn_query_expand", "lfm_attn_query_flat", 0), + "lfm_attn_key_expand_to_flat": _edge("lfm_attn_key_expand", "lfm_attn_key_flat", 0), + "lfm_attn_value_expand_to_flat": _edge("lfm_attn_value_expand", "lfm_attn_value_flat", 0), + "lfm_attn_query_flat_to_qk": _edge("lfm_attn_query_flat", "lfm_attn_qk", 0), + "lfm_attn_key_flat_to_qk": _edge("lfm_attn_key_flat", "lfm_attn_qk", 1), + "lfm_attn_qk_to_logits_view": _edge("lfm_attn_qk", "lfm_attn_logits_view", 0), + "lfm_attn_logits_view_to_mask_add": _edge("lfm_attn_logits_view", "lfm_attn_mask_add", 0), + "lfm_attn_mask_where_to_add": _edge("lfm_attn_mask_where", "lfm_attn_mask_add", 1), + "lfm_attn_mask_zero_to_where": _edge("lfm_attn_mask_zero", "lfm_attn_mask_where", 1), + "lfm_attn_mask_fill_to_where": _edge("lfm_attn_mask_fill", "lfm_attn_mask_where", 2), + "lfm_attn_mask_add_to_softmax": _edge("lfm_attn_mask_add", "lfm_attn_softmax", 0), + "lfm_attn_mask_add_to_null_eq": _edge("lfm_attn_mask_add", "lfm_attn_null_eq", 0), + "lfm_attn_null_eq_to_valid_not": _edge("lfm_attn_null_eq", "lfm_attn_valid_not", 0), + "lfm_attn_valid_not_to_any": _edge("lfm_attn_valid_not", "lfm_attn_any_valid", 0), + "lfm_attn_any_to_invalid_not": _edge("lfm_attn_any_valid", "lfm_attn_invalid_not", 0), + "lfm_attn_invalid_not_to_probs_where": _edge("lfm_attn_invalid_not", "lfm_attn_probs_where", 0), + "lfm_attn_softmax_to_zero_like": _edge("lfm_attn_softmax", "lfm_attn_zero_like", 0), + "lfm_attn_zero_like_to_probs_where": _edge("lfm_attn_zero_like", "lfm_attn_probs_where", 1), + "lfm_attn_softmax_to_probs_where": _edge("lfm_attn_softmax", "lfm_attn_probs_where", 2), + "lfm_attn_probs_where_to_expand": _edge("lfm_attn_probs_where", "lfm_attn_probs_expand", 0), + "lfm_attn_probs_expand_to_flat": _edge("lfm_attn_probs_expand", "lfm_attn_probs_flat", 0), + "lfm_attn_probs_flat_to_output_bmm": _edge("lfm_attn_probs_flat", "lfm_attn_output_bmm", 0), + "lfm_attn_value_flat_to_output_bmm": _edge("lfm_attn_value_flat", "lfm_attn_output_bmm", 1), + "lfm_attn_output_bmm_to_view": _edge("lfm_attn_output_bmm", "lfm_attn_output_view", 0), "rope_slice_odd_to_neg": _edge("rope_slice_odd", "rope_neg", 0), "rope_neg_to_rotate_cat": _edge("rope_neg", "rope_rotate_cat", 0), "rope_slice_even_to_rotate_cat": _edge("rope_slice_even", "rope_rotate_cat", 1), @@ -105,6 +129,24 @@ def _index_edges_by_dest(edges: dict[str, M.FusionEdge]) -> dict[str, tuple[str, "delta_q_to_decode": _edge("delta_q", "delta_gate", 0), "delta_k_to_decode": _edge("delta_k", "delta_gate", 1), "delta_v_to_decode": _edge("delta_v", "delta_gate", 2), + "lfm_conv_cache_clone_to_output_copy": _edge("lfm_conv_cache_clone", "lfm_conv_cache_output_copy", 0), + "lfm_conv_cache_scatter_dim0_to_output_copy": _edge("lfm_conv_cache_scatter_dim0", "lfm_conv_cache_output_copy", 1), + "lfm_conv_cache_clone_to_gather": _edge("lfm_conv_cache_clone", "lfm_conv_cache_gather", 0), + "lfm_conv_cache_roll_mod_to_gather": _edge("lfm_conv_cache_roll_mod", "lfm_conv_cache_gather", 1), + "lfm_conv_cache_roll_add_to_mod": _edge("lfm_conv_cache_roll_add", "lfm_conv_cache_roll_mod", 0), + "lfm_conv_cache_roll_arange_to_add": _edge("lfm_conv_cache_roll_arange", "lfm_conv_cache_roll_add", 0), + "lfm_conv_cache_gather_to_scatter_dim0": _edge("lfm_conv_cache_gather", "lfm_conv_cache_scatter_dim0", 0), + "lfm_conv_cache_scatter_dim1_to_scatter_dim0": _edge("lfm_conv_cache_scatter_dim1", "lfm_conv_cache_scatter_dim0", 1), + "lfm_conv_cache_scatter_slice_dim0_to_scatter_dim1": _edge("lfm_conv_cache_scatter_slice_dim0", "lfm_conv_cache_scatter_dim1", 0), + "lfm_conv_cache_scatter_dim2_to_scatter_dim1": _edge("lfm_conv_cache_scatter_dim2", "lfm_conv_cache_scatter_dim1", 1), + "lfm_conv_cache_scatter_slice_dim1_to_scatter_dim2": _edge("lfm_conv_cache_scatter_slice_dim1", "lfm_conv_cache_scatter_dim2", 0), + "lfm_conv_cache_value_copy_to_scatter_dim2": _edge("lfm_conv_cache_value_copy", "lfm_conv_cache_scatter_dim2", 1), + "lfm_conv_cache_gather_to_scatter_slice_dim0": _edge("lfm_conv_cache_gather", "lfm_conv_cache_scatter_slice_dim0", 0), + "lfm_conv_cache_scatter_slice_dim0_to_scatter_slice_dim1": _edge("lfm_conv_cache_scatter_slice_dim0", "lfm_conv_cache_scatter_slice_dim1", 0), + "lfm_conv_cache_gather_to_value_slice_dim0": _edge("lfm_conv_cache_gather", "lfm_conv_cache_value_slice_dim0", 0), + "lfm_conv_cache_value_slice_dim0_to_value_slice_dim1": _edge("lfm_conv_cache_value_slice_dim0", "lfm_conv_cache_value_slice_dim1", 0), + "lfm_conv_cache_value_slice_dim1_to_value_slice_last": _edge("lfm_conv_cache_value_slice_dim1", "lfm_conv_cache_value_slice_last", 0), + "lfm_conv_cache_value_slice_last_to_value_copy": _edge("lfm_conv_cache_value_slice_last", "lfm_conv_cache_value_copy", 0), "router_logits_to_probs": _edge("router_logits", "routing_probs", 0), "routing_probs_to_topk": _edge("routing_probs", "topk", 0), "moe_hidden_to_router": _edge("moe_hidden_view", "moe_router_logits", 0), @@ -227,6 +269,32 @@ def _index_edges_by_dest(edges: dict[str, M.FusionEdge]) -> dict[str, tuple[str, "attn_softmax_to_value", "v_view_to_value", ), + "lfm_bmm_masked_attention": ( + "lfm_attn_query_expand_to_flat", + "lfm_attn_key_expand_to_flat", + "lfm_attn_value_expand_to_flat", + "lfm_attn_query_flat_to_qk", + "lfm_attn_key_flat_to_qk", + "lfm_attn_qk_to_logits_view", + "lfm_attn_logits_view_to_mask_add", + "lfm_attn_mask_where_to_add", + "lfm_attn_mask_zero_to_where", + "lfm_attn_mask_fill_to_where", + "lfm_attn_mask_add_to_softmax", + "lfm_attn_mask_add_to_null_eq", + "lfm_attn_null_eq_to_valid_not", + "lfm_attn_valid_not_to_any", + "lfm_attn_any_to_invalid_not", + "lfm_attn_invalid_not_to_probs_where", + "lfm_attn_softmax_to_zero_like", + "lfm_attn_zero_like_to_probs_where", + "lfm_attn_softmax_to_probs_where", + "lfm_attn_probs_where_to_expand", + "lfm_attn_probs_expand_to_flat", + "lfm_attn_probs_flat_to_output_bmm", + "lfm_attn_value_flat_to_output_bmm", + "lfm_attn_output_bmm_to_view", + ), "rope": ( "rope_slice_odd_to_neg", "rope_neg_to_rotate_cat", @@ -250,6 +318,26 @@ def _index_edges_by_dest(edges: dict[str, M.FusionEdge]) -> dict[str, tuple[str, "expert_down_to_weighted", "expert_weighted_to_combine", ), + "lfm_conv_cache_roll_append": ( + "lfm_conv_cache_clone_to_output_copy", + "lfm_conv_cache_scatter_dim0_to_output_copy", + "lfm_conv_cache_clone_to_gather", + "lfm_conv_cache_roll_mod_to_gather", + "lfm_conv_cache_roll_add_to_mod", + "lfm_conv_cache_roll_arange_to_add", + "lfm_conv_cache_gather_to_scatter_dim0", + "lfm_conv_cache_scatter_dim1_to_scatter_dim0", + "lfm_conv_cache_scatter_slice_dim0_to_scatter_dim1", + "lfm_conv_cache_scatter_dim2_to_scatter_dim1", + "lfm_conv_cache_scatter_slice_dim1_to_scatter_dim2", + "lfm_conv_cache_value_copy_to_scatter_dim2", + "lfm_conv_cache_gather_to_scatter_slice_dim0", + "lfm_conv_cache_scatter_slice_dim0_to_scatter_slice_dim1", + "lfm_conv_cache_gather_to_value_slice_dim0", + "lfm_conv_cache_value_slice_dim0_to_value_slice_dim1", + "lfm_conv_cache_value_slice_dim1_to_value_slice_last", + "lfm_conv_cache_value_slice_last_to_value_copy", + ), "lfm_grouped_moe": ( "moe_hidden_to_router", "moe_router_weight_to_router", diff --git a/python/cactus/transpiler/Fusions/fusions.py b/python/cactus/transpiler/Fusions/fusions.py index d2705823d..d094a1ae8 100644 --- a/python/cactus/transpiler/Fusions/fusions.py +++ b/python/cactus/transpiler/Fusions/fusions.py @@ -532,6 +532,59 @@ def _index_by_root_op( }, ) +LFM_BMM_MASKED_ATTENTION_GRAPH = _graph( + "lfm_bmm_masked_attention", + "lfm_attn_output_view", + ( + "lfm_attn_query_expand", + "lfm_attn_key_expand", + "lfm_attn_value_expand", + "lfm_attn_query_flat", + "lfm_attn_key_flat", + "lfm_attn_value_flat", + "lfm_attn_qk", + "lfm_attn_logits_view", + "lfm_attn_mask_where", + "lfm_attn_mask_zero", + "lfm_attn_mask_fill", + "lfm_attn_mask_add", + "lfm_attn_softmax", + "lfm_attn_null_eq", + "lfm_attn_valid_not", + "lfm_attn_any_valid", + "lfm_attn_invalid_not", + "lfm_attn_zero_like", + "lfm_attn_probs_where", + "lfm_attn_probs_expand", + "lfm_attn_probs_flat", + "lfm_attn_output_bmm", + "lfm_attn_output_view", + ), + edge_names=E.EDGE_GROUPS["lfm_bmm_masked_attention"], + inputs=( + _input("query", "lfm_attn_query_expand", 0), + _input("key", "lfm_attn_key_expand", 0), + _input("value", "lfm_attn_value_expand", 0), + _input("mask_predicate", "lfm_attn_mask_where", 0, metadata={"drop_after_fusion": True}), + ), + attr_captures=( + M.AttrCapture("scale", default=1.0, required=False), + M.AttrCapture("is_causal", default=True, required=False), + M.AttrCapture("position_offset", default=0, required=False), + M.AttrCapture("window_size", default=0, required=False), + M.AttrCapture("input_layout", default="bhqd_bhds_bhsd", required=False), + M.AttrCapture("output_layout", default="bhqd", required=False), + M.AttrCapture("dropped_mask_builder", default=True, required=False), + ), + constraints={ + "node_attr_equals": ( + {"node": "lfm_attn_softmax", "attr": "axis", "value": -1, "allow_missing": True}, + {"node": "lfm_attn_null_eq", "attr": "other", "value": None}, + {"node": "lfm_attn_any_valid", "attr": "arg_1", "value": -1}, + ), + }, +) + ROPE_GRAPH = _graph( "rope", "rope_add", @@ -569,7 +622,10 @@ def _index_by_root_op( "kv_cache_append", "kv_cache_append", ("kv_cache_append",), - inputs=(_input("new_kv", "kv_cache_append", 0), _input("cache_state", "kv_cache_append", 1)), + inputs=( + _input("new_kv", "kv_cache_append", 0), + _input("cache_state", "kv_cache_append", 1, allowed_value_kinds=(M.ValueKind.CACHE_INPUT, M.ValueKind.CACHE_STATE, M.ValueKind.CACHE_OUTPUT)), + ), cache_inputs=(_cache_input("cache_state", "kv_cache_append", 1, tensor_role=M.CacheTensorRole.STATE),), cache_outputs=(_cache_output("updated_cache_state", "kv_cache_append", tensor_role=M.CacheTensorRole.STATE),), cache_mutations=(_cache_mutation("append_kv", read_roles=("cache_state",), write_roles=("updated_cache_state",)),), @@ -739,7 +795,10 @@ def _index_by_root_op( "rel_pos_bias", "rel_pos_bias", ("rel_pos_bias",), - inputs=(_input("query", "rel_pos_bias", 0), _input("relative_key", "rel_pos_bias", 1)), + inputs=( + _input("query", "rel_pos_bias", 0, tensor_constraints=(M.TensorConstraint(rank=4),)), + _input("relative_key", "rel_pos_bias", 1, tensor_constraints=(M.TensorConstraint(rank=4),)), + ), attr_captures=(M.AttrCapture("scale", "rel_pos_bias", "scale", required=False),), ) @@ -763,17 +822,81 @@ def _index_by_root_op( "conv_cache_append", "conv_cache_append", ("conv_cache_append",), - inputs=(_input("new_data", "conv_cache_append", 0), _input("cache_state", "conv_cache_append", 1)), + inputs=( + _input("new_data", "conv_cache_append", 0), + _input("cache_state", "conv_cache_append", 1, allowed_value_kinds=(M.ValueKind.CACHE_INPUT, M.ValueKind.CACHE_STATE, M.ValueKind.CACHE_OUTPUT)), + ), cache_inputs=(_cache_input("cache_state", "conv_cache_append", 1, cache_kind=M.CacheKind.CONV, tensor_role=M.CacheTensorRole.STATE),), cache_outputs=(_cache_output("updated_cache_state", "conv_cache_append", cache_kind=M.CacheKind.CONV, tensor_role=M.CacheTensorRole.STATE),), cache_mutations=(_cache_mutation("append_conv_state", cache_kind=M.CacheKind.CONV, read_roles=("cache_state",), write_roles=("updated_cache_state",)),), ) +LFM_CONV_CACHE_ROLL_APPEND_GRAPH = _graph( + "lfm_conv_cache_roll_append", + "lfm_conv_cache_output_copy", + ( + "lfm_conv_cache_clone", + "lfm_conv_cache_roll_arange", + "lfm_conv_cache_roll_add", + "lfm_conv_cache_roll_mod", + "lfm_conv_cache_gather", + "lfm_conv_cache_scatter_slice_dim0", + "lfm_conv_cache_scatter_slice_dim1", + "lfm_conv_cache_value_slice_dim0", + "lfm_conv_cache_value_slice_dim1", + "lfm_conv_cache_value_slice_last", + "lfm_conv_cache_value_copy", + "lfm_conv_cache_scatter_dim2", + "lfm_conv_cache_scatter_dim1", + "lfm_conv_cache_scatter_dim0", + "lfm_conv_cache_output_copy", + ), + edge_names=E.EDGE_GROUPS["lfm_conv_cache_roll_append"], + inputs=( + _input("new_data", "lfm_conv_cache_value_copy", 1), + _input("cache_state", "lfm_conv_cache_clone", 0, allowed_value_kinds=(M.ValueKind.CACHE_INPUT, M.ValueKind.CACHE_STATE, M.ValueKind.CACHE_OUTPUT)), + ), + cache_inputs=( + _cache_input("cache_state", "lfm_conv_cache_clone", 0, cache_kind=M.CacheKind.CONV, tensor_role=M.CacheTensorRole.STATE), + ), + cache_outputs=( + _cache_output("updated_cache_state", "lfm_conv_cache_output_copy", cache_kind=M.CacheKind.CONV, tensor_role=M.CacheTensorRole.STATE), + ), + cache_mutations=( + _cache_mutation("roll_append_conv_state", cache_kind=M.CacheKind.CONV, read_roles=("cache_state",), write_roles=("updated_cache_state",)), + ), + attr_captures=( + M.AttrCapture("window_size", default=3, required=False), + ), + constraints={ + "cache_roll_append_structure": { + "cache_role": "cache_state", + "new_data_role": "new_data", + "clone_node": "lfm_conv_cache_clone", + "gather_node": "lfm_conv_cache_gather", + "arange_node": "lfm_conv_cache_roll_arange", + "add_node": "lfm_conv_cache_roll_add", + "mod_node": "lfm_conv_cache_roll_mod", + "scatter_slice_dim0_node": "lfm_conv_cache_scatter_slice_dim0", + "scatter_slice_dim1_node": "lfm_conv_cache_scatter_slice_dim1", + "value_slice_dim0_node": "lfm_conv_cache_value_slice_dim0", + "value_slice_dim1_node": "lfm_conv_cache_value_slice_dim1", + "value_slice_last_node": "lfm_conv_cache_value_slice_last", + "scatter_dim2_node": "lfm_conv_cache_scatter_dim2", + "scatter_dim1_node": "lfm_conv_cache_scatter_dim1", + "scatter_dim0_node": "lfm_conv_cache_scatter_dim0", + }, + }, +) + CONV_CACHE_INITIALIZE_GRAPH = _graph( "conv_cache_initialize", "conv_cache_initialize", ("conv_cache_initialize",), - inputs=(_input("rows", "conv_cache_initialize", 0), _input("cache_state", "conv_cache_initialize", 1)), + inputs=( + _input("rows", "conv_cache_initialize", 0), + _input("cache_state", "conv_cache_initialize", 1, allowed_value_kinds=(M.ValueKind.CACHE_INPUT, M.ValueKind.CACHE_STATE, M.ValueKind.CACHE_OUTPUT)), + ), cache_inputs=(_cache_input("cache_state", "conv_cache_initialize", 1, cache_kind=M.CacheKind.CONV, tensor_role=M.CacheTensorRole.STATE),), cache_outputs=(_cache_output("initialized_cache_state", "conv_cache_initialize", cache_kind=M.CacheKind.CONV, tensor_role=M.CacheTensorRole.STATE),), ) @@ -858,6 +981,7 @@ def _index_by_root_op( "scaled_dot_product_attention": ATTENTION_DIRECT_GRAPH, "attention_core": ATTENTION_CORE_GRAPH, "attention_masked": ATTENTION_MASKED_GRAPH, + "lfm_bmm_masked_attention": LFM_BMM_MASKED_ATTENTION_GRAPH, "rope": ROPE_GRAPH, "conv": CONV_GRAPH, "conv_bias": CONV_BIAS_GRAPH, @@ -873,6 +997,7 @@ def _index_by_root_op( "kv_cache_state": KV_CACHE_STATE_GRAPH, "conv_cache_state": CONV_CACHE_STATE_GRAPH, "conv_cache_append": CONV_CACHE_APPEND_GRAPH, + "lfm_conv_cache_roll_append": LFM_CONV_CACHE_ROLL_APPEND_GRAPH, "conv_cache_initialize": CONV_CACHE_INITIALIZE_GRAPH, "recurrent_cache_state": RECURRENT_CACHE_STATE_GRAPH, "recurrent_cache_write": RECURRENT_CACHE_WRITE_GRAPH, @@ -906,10 +1031,11 @@ def _index_by_root_op( "scaled_dot_product_attention": _definition("scaled_dot_product_attention", "attention", ATTENTION_DIRECT_GRAPH, fusion_fields=("generic", "attention")), "attention_core": _definition("attention_core", "attention", ATTENTION_CORE_GRAPH, fusion_fields=("generic", "attention", "gemma4_attention", "qwen2_5_attention", "lfm_attention")), "attention_masked": _definition("attention_masked", "attention", ATTENTION_MASKED_GRAPH, fusion_fields=("generic", "attention", "gemma4_attention")), + "lfm_bmm_masked_attention": _definition("lfm_bmm_masked_attention", "attention", LFM_BMM_MASKED_ATTENTION_GRAPH, fusion_fields=("generic", "attention", "lfm_attention", "lfm_moe")), "rope": _definition("rope", "rope", ROPE_GRAPH, fusion_fields=("generic", "rope", "gemma4_rope", "qwen2_5_rope")), - "conv": _definition("conv", "conv1d", CONV_GRAPH, fusion_fields=("generic", "conv", "audio", "vision")), + "conv": _definition("conv", "conv1d", CONV_GRAPH, fusion_fields=("generic", "conv", "audio", "vision"), metadata=_required_attrs(padding=0, dilation=1, groups=1)), "conv_bias": _definition("conv_bias", "conv1d", CONV_BIAS_GRAPH, fusion_fields=("generic", "conv", "audio", "vision")), - "conv1d_k3": _definition("conv1d_k3", "conv1d_k3", CONV_GRAPH, fusion_fields=("generic", "conv", "audio"), metadata=_required_attrs(kernel_size=3)), + "conv1d_k3": _definition("conv1d_k3", "conv1d_k3", CONV_GRAPH, fusion_fields=("generic", "conv", "audio"), metadata=_required_attrs(kernel_size=3, stride=1, padding=0, dilation=1, groups=1)), "conv1d_k7s3": _definition("conv1d_k7s3", "conv1d_k7s3", CONV_BIAS_GRAPH, fusion_fields=("generic", "conv", "audio"), metadata=_required_attrs(kernel_size=7, stride=3)), "conv1d_causal": _definition("conv1d_causal", "conv1d_causal", CONV_GRAPH, fusion_fields=("generic", "conv", "cache"), metadata=_required_attrs(causal=True)), "conv1d_pointwise": _definition("conv1d_pointwise", "conv1d_pointwise", CONV_GRAPH, fusion_fields=("generic", "conv", "audio"), metadata=_required_attrs(kernel_size=1)), @@ -929,6 +1055,7 @@ def _index_by_root_op( "kv_cache_state": _definition("kv_cache_state", "kv_cache_state", KV_CACHE_STATE_GRAPH, fusion_fields=("generic", "cache"), supported_inference_modes=("prefill_with_cache", "decode_with_cache")), "conv_cache_state": _definition("conv_cache_state", "conv_cache_state", CONV_CACHE_STATE_GRAPH, fusion_fields=("generic", "cache", "conv")), "conv_cache_append": _definition("conv_cache_append", "conv_cache_append", CONV_CACHE_APPEND_GRAPH, fusion_fields=("generic", "cache", "conv")), + "lfm_conv_cache_roll_append": _definition("lfm_conv_cache_roll_append", "conv_cache_append", LFM_CONV_CACHE_ROLL_APPEND_GRAPH, fusion_fields=("generic", "cache", "conv", "lfm_moe"), supported_inference_modes=("decode_with_cache",)), "conv_cache_initialize": _definition("conv_cache_initialize", "conv_cache_initialize", CONV_CACHE_INITIALIZE_GRAPH, fusion_fields=("generic", "cache", "conv")), "recurrent_cache_state": _definition("recurrent_cache_state", "recurrent_cache_state", RECURRENT_CACHE_STATE_GRAPH, fusion_fields=("generic", "cache", "recurrent")), "recurrent_cache_write": _definition("recurrent_cache_write", "recurrent_cache_write", RECURRENT_CACHE_WRITE_GRAPH, fusion_fields=("generic", "cache", "recurrent")), @@ -941,8 +1068,6 @@ def _index_by_root_op( "rfft": _definition("rfft", "rfft", DSP_GRAPHS["rfft"], fusion_fields=("generic", "audio", "dsp")), "irfft": _definition("irfft", "irfft", DSP_GRAPHS["irfft"], fusion_fields=("generic", "audio", "dsp")), "spectrogram": _definition("spectrogram", "spectrogram", DSP_GRAPHS["spectrogram"], fusion_fields=("generic", "audio", "dsp")), - "image_preprocess": _definition("image_preprocess", "image_preprocess", DSP_GRAPHS["image_preprocess"], fusion_fields=("generic", "vision", "preprocess")), - "bilinear_interpolation": _definition("bilinear_interpolation", "bilinear_interpolation", DSP_GRAPHS["bilinear_interpolation"], fusion_fields=("generic", "vision")), } diff --git a/python/cactus/transpiler/Fusions/nodes.py b/python/cactus/transpiler/Fusions/nodes.py index 9aaf81bfe..8a81c67d2 100644 --- a/python/cactus/transpiler/Fusions/nodes.py +++ b/python/cactus/transpiler/Fusions/nodes.py @@ -12,12 +12,13 @@ "placeholder": ("placeholder",), "getitem": ("operator.getitem", ""), "embedding": ("aten.embedding.default", "aten.index_select.default", "aten.gather.default"), - "matmul": ("aten.mm.default", "aten.matmul.default", "aten.bmm.default", "aten.addmm.default", "aten.linear.default"), - "add": ("aten.add.Tensor", "aten.add.Scalar", "operator.add", ""), - "sub": ("aten.sub.Tensor", "aten.sub.Scalar", "operator.sub", ""), - "mul": ("aten.mul.Tensor", "aten.mul.Scalar", "operator.mul", ""), - "div": ("aten.div.Tensor", "aten.div.Scalar", "operator.truediv", ""), + "matmul": ("aten.mm.default", "aten.matmul.default", "aten.bmm.default", "aten.addmm.default", "aten.linear.default", "cactus.matmul"), + "add": ("aten.add.Tensor", "aten.add.Scalar", "operator.add", "", "cactus.add", "cactus.scalar_add"), + "sub": ("aten.sub.Tensor", "aten.sub.Scalar", "operator.sub", "", "cactus.subtract", "cactus.scalar_subtract"), + "mul": ("aten.mul.Tensor", "aten.mul.Scalar", "operator.mul", "", "cactus.multiply", "cactus.scalar_multiply"), + "div": ("aten.div.Tensor", "aten.div.Scalar", "operator.truediv", "", "cactus.divide", "cactus.scalar_divide"), "floor_divide": ("aten.floor_divide.default", "aten.div.Tensor_mode"), + "fmod": ("aten.fmod.Scalar",), "abs": ("aten.abs.default",), "pow": ("aten.pow.Tensor_Scalar", "aten.pow.Tensor_Tensor", "aten.pow.Scalar"), "sqrt": ("aten.sqrt.default", "aten.rsqrt.default"), @@ -30,18 +31,19 @@ "var": ("aten.var.dim", "aten.var.correction"), "max": ("aten.max.dim", "aten.amax.default"), "min": ("aten.min.dim", "aten.amin.default"), - "softmax": ("aten.softmax.int", "aten._softmax.default"), + "softmax": ("aten.softmax.int", "aten._softmax.default", "cactus.softmax"), "topk": ("aten.topk.default",), "sort": ("aten.sort.default",), - "view": ("aten.view.default", "aten.reshape.default", "aten.flatten.using_ints"), + "view": ("aten.view.default", "aten.reshape.default", "aten.flatten.using_ints", "cactus.view", "cactus.reshape"), "unsqueeze": ("aten.unsqueeze.default",), "expand": ("aten.expand.default",), - "transpose": ("aten.t.default", "aten.transpose.int", "aten.permute.default"), + "transpose": ("aten.t.default", "aten.transpose.int", "aten.permute.default", "cactus.transpose"), "split": ("aten.split_with_sizes.default", "aten.chunk.default"), - "slice": ("aten.slice.Tensor", "aten.select.int", "aten.narrow.default"), - "index": ("aten.select.int", "aten.index.Tensor"), - "index_put": ("aten.index_put.default",), - "cat": ("aten.cat.default",), + "slice": ("aten.slice.Tensor", "aten.narrow.default", "cactus.slice"), + "index": ("aten.select.int", "aten.index.Tensor", "cactus.index"), + "gather": ("aten.gather.default", "aten.index_select.default", "cactus.gather"), + "index_put": ("aten.index_put.default", "cactus.index_put"), + "cat": ("aten.cat.default", "cactus.cat"), "scatter": ("aten.scatter.src", "aten.scatter.value", "aten.scatter_reduce.two"), "histc": ("aten.histc.default",), "cumsum": ("aten.cumsum.default",), @@ -50,6 +52,8 @@ "neg": ("aten.neg.default", "operator.neg", ""), "where": ("aten.where.self",), "eq": ("aten.eq.Tensor", "aten.eq.Scalar"), + "logical_not": ("aten.logical_not.default",), + "any": ("aten.any.dim",), "ne": ("aten.ne.Tensor", "aten.ne.Scalar"), "lt": ("aten.lt.Tensor", "aten.lt.Scalar"), "gt": ("aten.gt.Tensor", "aten.gt.Scalar"), @@ -57,6 +61,8 @@ "arange": ("aten.arange.default", "aten.arange.start", "aten.arange.start_step"), "scalar_tensor": ("aten.scalar_tensor.default",), "copy": ("aten.clone.default", "aten._to_copy.default", "aten.contiguous.default"), + "copy_write": ("aten.copy.default",), + "slice_scatter": ("aten.slice_scatter.default",), "relu": ("aten.relu.default",), "silu": ("aten.silu.default", "aten.sigmoid.default"), "gelu": ("aten.gelu.default",), @@ -156,7 +162,7 @@ def _expert_node(name: str, op_key: str) -> M.FusionNode: "index": _node("index", "index"), "index_put": _node("index_put", "index_put"), "cat": _node("cat", "cat"), - "gather": _node("gather", "embedding"), + "gather": _node("gather", "gather"), "histc": _node("histc", "histc"), "cumsum": _node("cumsum", "cumsum"), "empty": _node("empty", "empty"), @@ -214,6 +220,29 @@ def _expert_node(name: str, op_key: str) -> M.FusionNode: "attn_softmax": _node("attn_softmax", "softmax"), "attn_value": _node("attn_value", "matmul"), "attn_sdp": _node("attn_sdp", "scaled_dot_product_attention"), + "lfm_attn_query_expand": _node("lfm_attn_query_expand", "expand"), + "lfm_attn_key_expand": _node("lfm_attn_key_expand", "expand"), + "lfm_attn_value_expand": _node("lfm_attn_value_expand", "expand"), + "lfm_attn_query_flat": _node("lfm_attn_query_flat", "view"), + "lfm_attn_key_flat": _node("lfm_attn_key_flat", "view"), + "lfm_attn_value_flat": _node("lfm_attn_value_flat", "view"), + "lfm_attn_qk": _node("lfm_attn_qk", "matmul"), + "lfm_attn_logits_view": _node("lfm_attn_logits_view", "view"), + "lfm_attn_mask_where": _node("lfm_attn_mask_where", "where"), + "lfm_attn_mask_zero": _node("lfm_attn_mask_zero", "scalar_tensor"), + "lfm_attn_mask_fill": _node("lfm_attn_mask_fill", "scalar_tensor"), + "lfm_attn_mask_add": _node("lfm_attn_mask_add", "add"), + "lfm_attn_softmax": _node("lfm_attn_softmax", "softmax"), + "lfm_attn_null_eq": _node("lfm_attn_null_eq", "eq"), + "lfm_attn_valid_not": _node("lfm_attn_valid_not", "logical_not"), + "lfm_attn_any_valid": _node("lfm_attn_any_valid", "any"), + "lfm_attn_invalid_not": _node("lfm_attn_invalid_not", "logical_not"), + "lfm_attn_zero_like": _node("lfm_attn_zero_like", "full"), + "lfm_attn_probs_where": _node("lfm_attn_probs_where", "where"), + "lfm_attn_probs_expand": _node("lfm_attn_probs_expand", "expand"), + "lfm_attn_probs_flat": _node("lfm_attn_probs_flat", "view"), + "lfm_attn_output_bmm": _node("lfm_attn_output_bmm", "matmul"), + "lfm_attn_output_view": _node("lfm_attn_output_view", "view"), "rel_pos_bias": _node("rel_pos_bias", "matmul"), "rope_slice_even": _node("rope_slice_even", "slice"), "rope_slice_odd": _node("rope_slice_odd", "slice"), @@ -226,6 +255,21 @@ def _expert_node(name: str, op_key: str) -> M.FusionNode: "conv_bias_add": _node("conv_bias_add", "add"), "conv_cache_append": _node("conv_cache_append", "cat"), "conv_cache_initialize": _node("conv_cache_initialize", "copy"), + "lfm_conv_cache_clone": _node("lfm_conv_cache_clone", "copy"), + "lfm_conv_cache_roll_arange": _node("lfm_conv_cache_roll_arange", "arange"), + "lfm_conv_cache_roll_add": _node("lfm_conv_cache_roll_add", "add"), + "lfm_conv_cache_roll_mod": _node("lfm_conv_cache_roll_mod", "fmod"), + "lfm_conv_cache_gather": _node("lfm_conv_cache_gather", "gather"), + "lfm_conv_cache_scatter_slice_dim0": _node("lfm_conv_cache_scatter_slice_dim0", "slice"), + "lfm_conv_cache_scatter_slice_dim1": _node("lfm_conv_cache_scatter_slice_dim1", "slice"), + "lfm_conv_cache_value_slice_dim0": _node("lfm_conv_cache_value_slice_dim0", "slice"), + "lfm_conv_cache_value_slice_dim1": _node("lfm_conv_cache_value_slice_dim1", "slice"), + "lfm_conv_cache_value_slice_last": _node("lfm_conv_cache_value_slice_last", "slice"), + "lfm_conv_cache_value_copy": _node("lfm_conv_cache_value_copy", "copy_write"), + "lfm_conv_cache_scatter_dim2": _node("lfm_conv_cache_scatter_dim2", "slice_scatter"), + "lfm_conv_cache_scatter_dim1": _node("lfm_conv_cache_scatter_dim1", "slice_scatter"), + "lfm_conv_cache_scatter_dim0": _node("lfm_conv_cache_scatter_dim0", "slice_scatter"), + "lfm_conv_cache_output_copy": _node("lfm_conv_cache_output_copy", "copy_write"), "conv_cache_state": _node("conv_cache_state", "placeholder", allowed_value_kinds=CACHE_STATE_VALUE_KINDS), "kv_cache_append": _node("kv_cache_append", "cat"), "kv_cache_state": _node("kv_cache_state", "placeholder", allowed_value_kinds=CACHE_STATE_VALUE_KINDS), @@ -407,6 +451,21 @@ def _expert_node(name: str, op_key: str) -> M.FusionNode: "conv_cache_state", "conv_cache_append", "conv_cache_initialize", + "lfm_conv_cache_clone", + "lfm_conv_cache_roll_arange", + "lfm_conv_cache_roll_add", + "lfm_conv_cache_roll_mod", + "lfm_conv_cache_gather", + "lfm_conv_cache_scatter_slice_dim0", + "lfm_conv_cache_scatter_slice_dim1", + "lfm_conv_cache_value_slice_dim0", + "lfm_conv_cache_value_slice_dim1", + "lfm_conv_cache_value_slice_last", + "lfm_conv_cache_value_copy", + "lfm_conv_cache_scatter_dim2", + "lfm_conv_cache_scatter_dim1", + "lfm_conv_cache_scatter_dim0", + "lfm_conv_cache_output_copy", "recurrent_cache_state", "recurrent_cache_write", ), diff --git a/python/cactus/transpiler/Generator/constants.py b/python/cactus/transpiler/Generator/constants.py index f26f5b3cd..5ad7b2030 100644 --- a/python/cactus/transpiler/Generator/constants.py +++ b/python/cactus/transpiler/Generator/constants.py @@ -42,6 +42,10 @@ "aten.copy.default", } +PAD_TARGETS = { + "aten.constant_pad_nd.default", +} + CONSTANT_INPUT_TARGETS = { "aten.arange.default", "aten.arange.start", @@ -302,7 +306,6 @@ "aten.any.dim", "aten.bitwise_and.Tensor", "aten.bitwise_or.Tensor", - "aten.constant_pad_nd.default", "aten.eq.Scalar", "aten.eq.Tensor", "aten.fmod.Scalar", diff --git a/python/cactus/transpiler/Generator/generate.py b/python/cactus/transpiler/Generator/generate.py index 039c7392d..f620ebb33 100644 --- a/python/cactus/transpiler/Generator/generate.py +++ b/python/cactus/transpiler/Generator/generate.py @@ -175,6 +175,7 @@ def build_lowering_rules( add_rules(rules, constants.TO_COPY_TARGETS, lower_to_copy) add_rules(rules, constants.PASS_THROUGH_TARGETS, lower_pass_through) add_rules(rules, constants.COPY_TARGETS, lower_copy) + add_rules(rules, constants.PAD_TARGETS, lower_constant_pad_nd) add_rules(rules, constants.CONSTANT_INPUT_TARGETS, lower_constant_input) add_rules(rules, constants.NEG_TARGETS, lower_neg) add_rules(rules, constants.SOFTMAX_TARGETS, lower_softmax) @@ -335,7 +336,8 @@ def lower_binary(context: models.GenerationContext, node: IRModels.Node) -> Any: method = constants.BINARY_TARGETS[node.target] if len(inputs) == 2: - return getattr(context.graph, method)(inputs[0], inputs[1]) + left, right = align_binary_inputs(context, node, inputs) + return getattr(context.graph, method)(left, right) scalar = scalar_attr(node, "other") @@ -357,6 +359,25 @@ def lower_scalar(context: models.GenerationContext, node: IRModels.Node) -> Any: return getattr(context.graph, method)(inputs[0], value) +def align_binary_inputs(context: models.GenerationContext, node: IRModels.Node, inputs: tuple[Any, ...]) -> tuple[Any, Any]: + target_shape = output_shape(node) + left = align_input_to_shape(context, inputs[0], node.parents[0], target_shape) + right = align_input_to_shape(context, inputs[1], node.parents[1], target_shape) + return left, right + + +def align_input_to_shape(context: models.GenerationContext, value: Any, source_node: IRModels.Node, target_shape: tuple[Any, ...]) -> Any: + source_shape = meta_shape(source_node) + + if shape_matches_tensor(source_node, target_shape): + return value + + if element_count(source_shape) == element_count(target_shape): + return context.graph.reshape(value, target_shape) + + return value + + def lower_unary(context: models.GenerationContext, node: IRModels.Node) -> Any: inputs = require_input_count(context, node, 1) method = constants.UNARY_TARGETS[node.target] @@ -394,7 +415,21 @@ def lower_shape(context: models.GenerationContext, node: IRModels.Node) -> Any: def lower_expand(context: models.GenerationContext, node: IRModels.Node) -> Any: inputs = require_input_count(context, node, 1) - return context.graph.expand(inputs[0], shape_attr(node)) + target_shape = shape_attr(node) + + if shape_matches_tensor(node.parents[0], target_shape): + return inputs[0] + + if element_count(meta_shape(node.parents[0])) == element_count(target_shape): + return context.graph.reshape(inputs[0], target_shape) + + try: + return context.graph.expand(inputs[0], target_shape) + except AttributeError as e: + raise UnsupportedLoweringError( + f"{node.name}: native Cactus expand is unavailable for broadcast shape " + f"{meta_shape(node.parents[0])} -> {target_shape}" + ) from e def lower_unsqueeze(context: models.GenerationContext, node: IRModels.Node) -> Any: @@ -432,7 +467,18 @@ def lower_transpose(context: models.GenerationContext, node: IRModels.Node) -> A def lower_slice(context: models.GenerationContext, node: IRModels.Node) -> Any: inputs = require_input_count(context, node, 1) - axis = int(node.attrs.get("axis", node.attrs.get("dim", 0))) + axis_value = node.attrs.get("axis") + + if axis_value is None: + axis_value = node.attrs.get("dim") + + if axis_value is None: + if shape_matches_tensor(node.parents[0], output_shape(node)): + return inputs[0] + + raise UnsupportedLoweringError(f"{node.name}: slice lowering missing axis/dim attr") + + axis = normalize_dim(int(axis_value), len(meta_shape(node.parents[0]))) if node.target == "aten.select.int" or "index_value" in node.attrs: return context.graph.index(inputs[0], int(node.attrs.get("index_value", node.attrs.get("index", 0))), axis=axis) @@ -468,6 +514,10 @@ def lower_cat(context: models.GenerationContext, node: IRModels.Node) -> Any: def lower_matmul(context: models.GenerationContext, node: IRModels.Node) -> Any: inputs = require_input_count(context, node, 2) + + if node.target == "aten.bmm.default": + return lower_bmm(context, node, inputs) + return context.graph.matmul( inputs[0], inputs[1], @@ -475,6 +525,22 @@ def lower_matmul(context: models.GenerationContext, node: IRModels.Node) -> Any: ) +def lower_bmm(context: models.GenerationContext, node: IRModels.Node, inputs: tuple[Any, ...]) -> Any: + left_shape = meta_shape(node.parents[0]) + right_shape = meta_shape(node.parents[1]) + + if len(left_shape) != 3 or len(right_shape) != 3: + raise UnsupportedLoweringError(f"{node.name}: bmm lowering requires rank-3 inputs") + + if left_shape[0] != 1 or right_shape[0] != 1: + raise UnsupportedLoweringError(f"{node.name}: bmm lowering currently supports only batch size 1") + + left_2d = context.graph.reshape(inputs[0], (left_shape[1], left_shape[2])) + right_2d = context.graph.reshape(inputs[1], (right_shape[1], right_shape[2])) + product = context.graph.matmul(left_2d, right_2d) + return context.graph.reshape(product, output_shape(node)) + + def lower_addmm(context: models.GenerationContext, node: IRModels.Node) -> Any: inputs = require_input_count(context, node, 3) product = context.graph.matmul(inputs[1], inputs[2]) @@ -488,7 +554,7 @@ def lower_split(context: models.GenerationContext, node: IRModels.Node) -> Any: if not isinstance(split_sizes, list): raise UnsupportedLoweringError(f"{node.name}: split lowering requires split_sizes list") - axis = int(node.attrs.get("axis", node.attrs.get("dim", 0))) + axis = normalize_dim(int(node.attrs.get("axis", node.attrs.get("dim", 0))), len(meta_shape(node.parents[0]))) offset = 0 outputs = [] @@ -539,6 +605,49 @@ def lower_copy(context: models.GenerationContext, node: IRModels.Node) -> Any: return inputs[-1] +def lower_constant_pad_nd(context: models.GenerationContext, node: IRModels.Node) -> Any: + inputs = require_input_count(context, node, 1) + pads = node.attrs.get("pad", node.attrs.get("arg_1")) + + if not isinstance(pads, (list, tuple)) or len(pads) % 2 != 0: + raise UnsupportedLoweringError(f"{node.name}: constant_pad_nd lowering requires an even pad list") + + if any(int(pad) > 0 for pad in pads): + raise UnsupportedLoweringError(f"{node.name}: positive constant_pad_nd requires a native pad op") + + result = inputs[0] + parent_shape = meta_shape(node.parents[0]) + rank = len(parent_shape) + + for pair_index in range(0, len(pads), 2): + left_pad = int(pads[pair_index]) + right_pad = int(pads[pair_index + 1]) + + if left_pad == 0 and right_pad == 0: + continue + + axis = rank - 1 - (pair_index // 2) + + if axis < 0: + raise UnsupportedLoweringError(f"{node.name}: pad rank exceeds input rank") + + axis_size = parent_shape[axis] + + if not isinstance(axis_size, int): + raise UnsupportedLoweringError(f"{node.name}: cannot crop symbolic pad dimension {axis_size}") + + start = max(-left_pad, 0) + end_crop = max(-right_pad, 0) + length = axis_size - start - end_crop + + if length < 0: + raise UnsupportedLoweringError(f"{node.name}: pad crop removes more than the full axis") + + result = context.graph.slice(result, axis, start, length=length) + + return result + + def lower_constant_input(context: models.GenerationContext, node: IRModels.Node) -> Any: context.component.warn(f"{node.name}: constant-producing op {node.target} lowered as graph input") shape, dynamic_dims = graph_input_shape(node) @@ -626,7 +735,7 @@ def lower_conv(context: models.GenerationContext, node: IRModels.Node) -> Any: rank = output_rank(node) if target == "aten.conv1d.default" or rank == 3: - return context.graph.conv1d(inputs[0], inputs[1], bias=bias, stride=first_int(node.attrs.get("stride"), 1)) + return lower_generic_conv1d(context, node, inputs, bias) if target == "aten.conv2d.default" or rank == 4: return context.graph.conv2d( @@ -647,12 +756,13 @@ def lower_named_conv(context: models.GenerationContext, node: IRModels.Node, met bias = inputs[2] if len(inputs) > 2 else None if method == "conv1d": - return context.graph.conv1d(inputs[0], inputs[1], bias=bias, stride=first_int(node.attrs.get("stride"), 1)) + return lower_generic_conv1d(context, node, inputs, bias) if method == "conv1d_causal": return context.graph.conv1d_causal(inputs[0], inputs[1], kernel_size=int(node.attrs.get("kernel_size", 0)), dilation=first_int(node.attrs.get("dilation"), 1)) if method == "conv1d_k3": + require_plain_conv1d_attrs(node, method) return context.graph.conv1d_k3(inputs[0], inputs[1], stride=first_int(node.attrs.get("stride"), 1)) if method == "conv1d_k7s3": @@ -676,16 +786,33 @@ def lower_named_conv(context: models.GenerationContext, node: IRModels.Node, met raise UnsupportedLoweringError(f"{node.name}: unsupported Cactus conv method {method}") +def lower_generic_conv1d(context: models.GenerationContext, node: IRModels.Node, inputs: tuple[Any, ...], bias: Any | None) -> Any: + require_plain_conv1d_attrs(node, "conv1d") + return context.graph.conv1d(inputs[0], inputs[1], bias=bias, stride=first_int(node.attrs.get("stride"), 1)) + + +def require_plain_conv1d_attrs(node: IRModels.Node, method: str) -> None: + padding = first_int(node.attrs.get("padding"), 0) + dilation = first_int(node.attrs.get("dilation"), 1) + groups = first_int(node.attrs.get("groups"), 1) + + if padding != 0 or dilation != 1 or groups != 1: + raise UnsupportedLoweringError( + f"{node.name}: {method} lowering only supports padding=0, dilation=1, groups=1; " + f"got padding={padding}, dilation={dilation}, groups={groups}" + ) + + def lower_attention(context: models.GenerationContext, node: IRModels.Node) -> Any: inputs = context.inputs_for(node) if node.target == "cactus.attention" or node.target == "aten.scaled_dot_product_attention.default": require_len(node, inputs, 3) - mask = inputs[3] if len(inputs) > 3 else None - return context.graph.attention( - inputs[0], - inputs[1], - inputs[2], + query, key, value, mask = attention_inputs_for_layout(context, node, inputs) + output = context.graph.attention( + query, + key, + value, scale=float(node.attrs.get("scale", 1.0)), is_causal=bool(node.attrs.get("is_causal", True)), position_offset=int(node.attrs.get("position_offset", 0)), @@ -693,6 +820,7 @@ def lower_attention(context: models.GenerationContext, node: IRModels.Node) -> A mask=mask, additive_mask=bool(node.attrs.get("additive_mask", False)), ) + return attention_output_for_layout(context, node, output) if node.target == "cactus.attention_cached": require_len(node, inputs, 5) @@ -711,16 +839,39 @@ def lower_attention(context: models.GenerationContext, node: IRModels.Node) -> A raise UnsupportedLoweringError(f"{node.name}: unsupported attention target {node.target}") +def attention_inputs_for_layout(context: models.GenerationContext, node: IRModels.Node, inputs: tuple[Any, ...]) -> tuple[Any, Any, Any, Any | None]: + mask = inputs[3] if len(inputs) > 3 else None + + if node.attrs.get("input_layout") == "bhqd_bhds_bhsd": + return ( + context.graph.permute(inputs[0], (0, 2, 1, 3)), + context.graph.permute(inputs[1], (0, 3, 1, 2)), + context.graph.permute(inputs[2], (0, 2, 1, 3)), + mask, + ) + + return inputs[0], inputs[1], inputs[2], mask + + +def attention_output_for_layout(context: models.GenerationContext, node: IRModels.Node, output: Any) -> Any: + if node.attrs.get("output_layout") == "bhqd": + return context.graph.permute(output, (0, 2, 1, 3)) + + return output + + def lower_cache(context: models.GenerationContext, node: IRModels.Node) -> Any: inputs = context.inputs_for(node) target = node.target if target == "cactus.kv_cache_append": require_len(node, inputs, 2) + require_cache_state_parent(node, 1, "cactus.kv_cache_state") return context.graph.kv_cache_append(inputs[0], inputs[1], window_size=int(node.attrs.get("window_size", 0)), sink_size=int(node.attrs.get("sink_size", 0))) if target == "cactus.conv_cache_append": require_len(node, inputs, 2) + require_cache_state_parent(node, 1, "cactus.conv_cache_state") return context.graph.conv_cache_append(inputs[0], inputs[1]) if target == "cactus.conv_cache_initialize": @@ -750,6 +901,15 @@ def lower_cache(context: models.GenerationContext, node: IRModels.Node) -> Any: raise UnsupportedLoweringError(f"{node.name}: unsupported cache target {target}") +def require_cache_state_parent(node: IRModels.Node, parent_index: int, expected_target: str) -> None: + if parent_index >= len(node.parents) or node.parents[parent_index].target != expected_target: + actual = node.parents[parent_index].target if parent_index < len(node.parents) else "" + raise UnsupportedLoweringError( + f"{node.name}: {node.target} requires parent {parent_index} to be {expected_target}; got {actual}. " + "Raw HF cache tensors need a cache-state bridge before native cache append lowering." + ) + + def lower_moe(context: models.GenerationContext, node: IRModels.Node) -> Any: inputs = context.inputs_for(node) @@ -828,7 +988,7 @@ def lower_special_cactus(context: models.GenerationContext, node: IRModels.Node) if target == "cactus.rel_pos_bias": require_len(node, inputs, 2) - return context.graph.rel_pos_bias(inputs[0], inputs[1], scale=float(node.attrs.get("scale", 1.0))) + return context.graph.rel_pos_bias(inputs[0], inputs[1], scale=float(node.attrs.get("scale") or 1.0)) if target == "cactus.sample": require_len(node, inputs, 1) @@ -875,6 +1035,9 @@ def lower_special_cactus(context: models.GenerationContext, node: IRModels.Node) if target == "cactus.image_preprocess": require_len(node, inputs, 1) + if not has_attrs(node, ("src_width", "src_height", "target_width", "target_height", "patch_size", "channels")): + return context.graph.reshape(inputs[0], output_shape(node)) + return context.graph.image_preprocess( inputs[0], src_width=required_int_attr(node, "src_width"), @@ -985,6 +1148,26 @@ def output_shape(node: IRModels.Node) -> tuple[int, ...]: raise UnsupportedLoweringError(f"{node.name}: missing concrete output shape") +def shape_matches_tensor(node: IRModels.Node, expected_shape: tuple[int, ...]) -> bool: + actual_shape = meta_shape(node) + return tuple(actual_shape) == tuple(expected_shape) + + +def element_count(shape: tuple[Any, ...]) -> int | None: + if not shape: + return None + + count = 1 + + for dim in shape: + if not isinstance(dim, int) or dim < 0: + return None + + count *= dim + + return count + + def shape_attr(node: IRModels.Node) -> tuple[int, ...]: raw_shape = node.attrs.get("shape") @@ -1046,6 +1229,10 @@ def required_int_attr(node: IRModels.Node, name: str) -> int: return int(value) +def has_attrs(node: IRModels.Node, names: tuple[str, ...]) -> bool: + return all(node.attrs.get(name) is not None for name in names) + + def first_int(value: Any, default: int) -> int: if value is None: return default @@ -1106,16 +1293,37 @@ def slice_length(node: IRModels.Node, start: int) -> int: end = node.attrs.get("end") if end is None: - return 0 + return open_slice_length(node, start) end = int(end) if end >= constants.OPEN_SLICE_END: - return 0 + return open_slice_length(node, start) return max(end - start, 0) +def open_slice_length(node: IRModels.Node, start: int) -> int: + if not node.parents: + return 0 + + axis = int(node.attrs.get("axis", node.attrs.get("dim", 0))) + parent_shape = meta_shape(node.parents[0]) + + if axis < 0: + axis += len(parent_shape) + + if axis < 0 or axis >= len(parent_shape): + return 0 + + axis_size = parent_shape[axis] + + if not isinstance(axis_size, int): + return 0 + + return max(axis_size - start, 0) + + def binary_method_to_scalar_method(method: str) -> str: return { "add": "scalar_add", diff --git a/python/cactus/transpiler/IR/extra_matchers.py b/python/cactus/transpiler/IR/extra_matchers.py index ae9a7b524..9aa4a1511 100644 --- a/python/cactus/transpiler/IR/extra_matchers.py +++ b/python/cactus/transpiler/IR/extra_matchers.py @@ -1108,6 +1108,87 @@ def match_grouped_moe_structure(source: models.Node, graph: models.Graph, fusion return True +def match_cache_roll_append_structure(source: models.Node, graph: models.Graph, fusion: FModels.FusionGraph, bindings: dict[str, models.Node], spec: dict[str, Any]) -> bool: + cache_node = match_utils.get_first_input_by_role(fusion, bindings, spec["cache_role"]) + new_data_node = match_utils.get_first_input_by_role(fusion, bindings, spec["new_data_role"]) + + if cache_node is None or new_data_node is None: + return False + + cache_shape = match_utils.get_tensor_shape(cache_node) + new_data_shape = match_utils.get_tensor_shape(new_data_node) + output_shape = match_utils.get_tensor_shape(source) + + if not cache_shape or not new_data_shape or not output_shape: + return False + + if not match_utils.values_equal(cache_shape, output_shape): + return False + + if len(cache_shape) != len(new_data_shape): + return False + + window_size = get_known_int(cache_shape[-1]) + + if window_size is None or window_size <= 0: + return False + + if not cache_append_shapes_match(cache_shape, new_data_shape): + return False + + arange_node = bindings.get(spec["arange_node"]) + add_node = bindings.get(spec["add_node"]) + mod_node = bindings.get(spec["mod_node"]) + + if arange_node is None or add_node is None or mod_node is None: + return False + + if not match_utils.values_equal(get_attr(arange_node, "start", "arg_0"), 0): + return False + + if not match_utils.values_equal(get_attr(arange_node, "end", "arg_1"), window_size): + return False + + if not match_utils.values_equal(get_attr(add_node, "other", "value", "arg_1"), 1): + return False + + if not match_utils.values_equal(get_attr(mod_node, "other", "value", "arg_1"), window_size): + return False + + return ( + slice_node_matches(bindings.get(spec["scatter_slice_dim0_node"]), axis=0, start=0) + and slice_node_matches(bindings.get(spec["scatter_slice_dim1_node"]), axis=1, start=0) + and slice_node_matches(bindings.get(spec["value_slice_dim0_node"]), axis=0, start=0) + and slice_node_matches(bindings.get(spec["value_slice_dim1_node"]), axis=1, start=0) + and slice_node_matches(bindings.get(spec["value_slice_last_node"]), axis=2, start=-1) + and scatter_node_matches(bindings.get(spec["scatter_dim2_node"]), axis=2, start=-1) + and scatter_node_matches(bindings.get(spec["scatter_dim1_node"]), axis=1, start=0) + and scatter_node_matches(bindings.get(spec["scatter_dim0_node"]), axis=0, start=0) + ) + + +def cache_append_shapes_match(cache_shape: list[Any], new_data_shape: list[Any]) -> bool: + for cache_dim, new_dim in zip(cache_shape[:-1], new_data_shape[:-1]): + if not match_utils.values_equal(cache_dim, new_dim): + return False + + return match_utils.values_equal(new_data_shape[-1], 1) + + +def slice_node_matches(node: models.Node | None, *, axis: int, start: int) -> bool: + if node is None: + return False + + return match_utils.values_equal(get_attr(node, "axis", "dim", "arg_1"), axis) and match_utils.values_equal(get_attr(node, "start", "arg_2"), start) + + +def scatter_node_matches(node: models.Node | None, *, axis: int, start: int) -> bool: + if node is None: + return False + + return match_utils.values_equal(get_attr(node, "dim", "axis", "arg_2"), axis) and match_utils.values_equal(get_attr(node, "start", "arg_3"), start) + + def match_node_attr_equals(source: models.Node, graph: models.Graph, fusion: FModels.FusionGraph, bindings: dict[str, models.Node], spec: dict[str, Any]) -> bool: """ Checks that one matched node has a specific normalized attr value. @@ -1254,6 +1335,7 @@ def match_same_layer(source: models.Node, graph: models.Graph, fusion: FModels.F "moe_expert_branch_routing": match_moe_expert_branch_routing, "moe_routing_weights_combine": match_moe_routing_weights_combine, "grouped_moe_structure": match_grouped_moe_structure, + "cache_roll_append_structure": match_cache_roll_append_structure, "node_attr_equals": match_node_attr_equals, "node_attrs_equal": match_node_attrs_equal, "input_value_kind": match_input_value_kind, diff --git a/python/cactus/transpiler/IR/match_utils.py b/python/cactus/transpiler/IR/match_utils.py index 08f28ee05..947c4c346 100644 --- a/python/cactus/transpiler/IR/match_utils.py +++ b/python/cactus/transpiler/IR/match_utils.py @@ -350,7 +350,15 @@ def get_attr_capture_value(capture: FModels.AttrCapture, bindings: dict[str, mod if source_node is None: return MISSING if capture.required else capture.default - return source_node.attrs.get(capture.source_attr, MISSING if capture.required else capture.default) + value = source_node.attrs.get(capture.source_attr, MISSING) + + if value is MISSING: + value = source_node.attrs.get(capture.name, MISSING) + + if value is MISSING: + return MISSING if capture.required else capture.default + + return value def infer_definition_attrs(fusion: FModels.FusionGraph, bindings: dict[str, models.Node]) -> dict[str, Any]: diff --git a/python/cactus/transpiler/IR/models.py b/python/cactus/transpiler/IR/models.py index 74084fcd2..9cf8a206a 100644 --- a/python/cactus/transpiler/IR/models.py +++ b/python/cactus/transpiler/IR/models.py @@ -259,7 +259,8 @@ def apply_fusions_to_graph(graph: Graph, fusion_results: tuple[FusionResult, ... consumed_names = expanded_consumed_names(graph, replacement_names) kept_nodes = tuple(clone_rewritten_node(node, replacement_names) for node in graph.nodes if node.name not in consumed_names) rewritten_fused_nodes = tuple(clone_rewritten_node(node, replacement_names) for node in fused_nodes) - return rebuild_graph((*kept_nodes, *rewritten_fused_nodes), graph, fusion_results) + rebuilt_graph = rebuild_graph((*kept_nodes, *rewritten_fused_nodes), graph, fusion_results) + return prune_dead_nodes(rebuilt_graph) def validate_non_overlapping_fusions(fusion_results: tuple[FusionResult, ...]) -> None: @@ -427,6 +428,32 @@ def rebuild_graph(nodes: tuple[Node, ...], original_graph: Graph, fusion_results ) +def prune_dead_nodes(graph: Graph) -> Graph: + live_names = output_ancestor_names(graph) + + if len(live_names) == len(graph.nodes): + return graph + + live_nodes = tuple(node for node in graph.nodes if node.name in live_names) + return rebuild_graph(live_nodes, graph, tuple(graph.fusions)) + + +def output_ancestor_names(graph: Graph) -> frozenset[str]: + live_names: set[str] = set() + stack = list(graph.outputs) + + while stack: + node = stack.pop() + + if node.name in live_names: + continue + + live_names.add(node.name) + stack.extend(node.parents) + + return frozenset(live_names) + + def graph_to_layer_map(graph: Graph) -> CModels.LayerMap: records = [ CModels.LayerRecord( diff --git a/python/cactus/transpiler/IR/simplify_ir.py b/python/cactus/transpiler/IR/simplify_ir.py index 3953fd48c..aa9589005 100644 --- a/python/cactus/transpiler/IR/simplify_ir.py +++ b/python/cactus/transpiler/IR/simplify_ir.py @@ -168,6 +168,9 @@ def collect_external_inputs( external_inputs: list[models.Node] = [] for input_spec in fusion.inputs: + if input_spec.metadata.get("drop_after_fusion"): + continue + external_inputs.extend(match_utils.get_fusion_input_nodes(input_spec, bindings)) for cache_input in fusion.cache_inputs: From fb5963e3dd21004e5f0b1e666633b7ac16147004 Mon Sep 17 00:00:00 2001 From: Parkirat Date: Wed, 29 Jul 2026 11:30:47 -0400 Subject: [PATCH 63/84] Additional conv-state matching logic --- python/cactus/transpiler/IR/extra_matchers.py | 216 +++++++++++++ python/cactus/transpiler/IR/match.py | 2 +- python/cactus/transpiler/IR/match_utils.py | 19 ++ python/cactus/transpiler/IR/models.py | 305 +++++++++++++++++- python/cactus/transpiler/IR/simplify_ir.py | 41 ++- 5 files changed, 573 insertions(+), 10 deletions(-) diff --git a/python/cactus/transpiler/IR/extra_matchers.py b/python/cactus/transpiler/IR/extra_matchers.py index 9aa4a1511..f0f4d2d50 100644 --- a/python/cactus/transpiler/IR/extra_matchers.py +++ b/python/cactus/transpiler/IR/extra_matchers.py @@ -1167,6 +1167,204 @@ def match_cache_roll_append_structure(source: models.Node, graph: models.Graph, ) +def match_empty_cache_initializer(source: models.Node, graph: models.Graph, fusion: FModels.FusionGraph, bindings: dict[str, models.Node], spec: dict[str, Any]) -> bool: + empty_node = match_utils.get_first_input_by_role(fusion, bindings, spec["empty_role"]) + new_value_node = match_utils.get_first_input_by_role(fusion, bindings, spec["new_value_role"]) + output_node = bindings.get(spec["output_node"]) + + if empty_node is None or new_value_node is None or output_node is None: + return False + + empty_shape = match_utils.get_tensor_shape(empty_node) + new_shape = match_utils.get_tensor_shape(new_value_node) + output_shape = match_utils.get_tensor_shape(output_node) + + if not empty_shape or not new_shape or not output_shape: + return False + + if element_count(empty_shape) != 0: + return False + + return match_utils.values_equal(new_shape, output_shape) + + +def match_conv_cache_initialize_structure(source: models.Node, graph: models.Graph, fusion: FModels.FusionGraph, bindings: dict[str, models.Node], spec: dict[str, Any]) -> bool: + rows_node = match_utils.get_first_input_by_role(fusion, bindings, spec["rows_role"]) + empty_node = match_utils.get_first_input_by_role(fusion, bindings, spec["empty_role"]) + output_node = bindings.get(spec["output_node"]) + + if rows_node is None or empty_node is None or output_node is None: + return False + + rows_shape = match_utils.get_tensor_shape(rows_node) + empty_shape = match_utils.get_tensor_shape(empty_node) + output_shape = match_utils.get_tensor_shape(output_node) + + if len(rows_shape) != 3 or len(empty_shape) != 3 or len(output_shape) != 3: + return False + + if not match_utils.values_equal(rows_shape, output_shape): + return False + + if not match_utils.values_equal(empty_shape, output_shape): + return False + + if empty_node.target != "aten.full_like.default": + return False + + if not match_utils.values_equal(get_attr(empty_node, "fill_value", "arg_1"), 0): + return False + + return get_known_int(rows_shape[1]) is not None and get_known_int(rows_shape[2]) is not None + + +def match_short_conv_prefill_structure(source: models.Node, graph: models.Graph, fusion: FModels.FusionGraph, bindings: dict[str, models.Node], spec: dict[str, Any]) -> bool: + conv_node = bindings.get(spec["conv_node"]) + slice_node = bindings.get(spec["slice_node"]) + x_node = match_utils.get_first_input_by_role(fusion, bindings, spec["x_role"]) + weight_node = match_utils.get_first_input_by_role(fusion, bindings, spec["weight_role"]) + + if conv_node is None or slice_node is None or x_node is None or weight_node is None: + return False + + x_shape = match_utils.get_tensor_shape(x_node) + weight_shape = match_utils.get_tensor_shape(weight_node) + conv_shape = match_utils.get_tensor_shape(conv_node) + output_shape = match_utils.get_tensor_shape(slice_node) + + if len(x_shape) != 3 or len(weight_shape) != 3 or len(conv_shape) != 3 or len(output_shape) != 3: + return False + + hidden_dim = get_known_int(x_shape[1]) + kernel_size = get_known_int(weight_shape[2]) + input_length = get_known_int(x_shape[2]) + + if hidden_dim is None or kernel_size is None or input_length is None: + return False + + if not match_utils.values_equal(weight_shape[0], hidden_dim): + return False + + if not match_utils.values_equal(weight_shape[1], 1): + return False + + if first_int_attr(conv_node, "stride", 1) != 1: + return False + + if first_int_attr(conv_node, "dilation", 1) != 1: + return False + + if first_int_attr(conv_node, "padding", 0) != kernel_size - 1: + return False + + if first_int_attr(conv_node, "groups", 1) != hidden_dim: + return False + + slice_axis = first_int_attr(slice_node, "axis", first_int_attr(slice_node, "dim", -1)) + if normalize_dim(slice_axis, 3) != 2: + return False + + if first_int_attr(slice_node, "start", 0) != 0: + return False + + return match_utils.values_equal(output_shape, x_shape) and get_shape_dim(conv_node, 2) == input_length + kernel_size - 1 + + +def match_short_conv_decode_structure(source: models.Node, graph: models.Graph, fusion: FModels.FusionGraph, bindings: dict[str, models.Node], spec: dict[str, Any]) -> bool: + cache_window_node = match_utils.get_first_input_by_role(fusion, bindings, spec["cache_window_role"]) + multiply_node = bindings.get(spec["multiply_node"]) + sum_node = bindings.get(spec["sum_node"]) + weight_node = match_utils.get_first_input_by_role(fusion, bindings, spec["weight_role"]) + + if cache_window_node is None or multiply_node is None or sum_node is None or weight_node is None: + return False + + if cache_window_node.target != "cactus.conv_cache_append" and not ( + cache_window_node.cache is not None and cache_window_node.cache.kind == FModels.CacheKind.CONV + ): + return False + + cache_window_shape = match_utils.get_tensor_shape(cache_window_node) + weight_shape = match_utils.get_tensor_shape(weight_node) + multiply_shape = match_utils.get_tensor_shape(multiply_node) + output_shape = match_utils.get_tensor_shape(sum_node) + + if len(cache_window_shape) != 3 or len(multiply_shape) != 3 or len(output_shape) != 2: + return False + + if not match_utils.values_equal(cache_window_shape, multiply_shape): + return False + + batch_size = get_known_int(cache_window_shape[0]) + hidden_dim = get_known_int(cache_window_shape[1]) + window_size = get_known_int(cache_window_shape[2]) + + if batch_size is None or hidden_dim is None or window_size is None: + return False + + if not match_utils.values_equal(output_shape, [batch_size, hidden_dim]): + return False + + if not short_conv_decode_weight_shape_matches(weight_shape, hidden_dim, window_size): + return False + + sum_axis = first_int_attr(sum_node, "axis", first_int_attr(sum_node, "dim", -1)) + if normalize_dim(sum_axis, 3) != 2: + return False + + keepdim = get_attr(sum_node, "keepdim") + return keepdim is match_utils.MISSING or keepdim is False + + +def match_cache_output_consumers(source: models.Node, graph: models.Graph, fusion: FModels.FusionGraph, bindings: dict[str, models.Node], spec: dict[str, Any]) -> bool: + node = bindings.get(spec["node"]) + + if node is None: + return False + + allow_output_child = bool(spec.get("allow_output_child", True)) + + for child in node.children: + if child.target in match_utils.METADATA_ONLY_TARGETS: + continue + + if allow_output_child and child.is_output: + continue + + return False + + return True + + +def element_count(shape: list[Any]) -> int | None: + count = 1 + + for dim in shape: + known_dim = get_known_int(dim) + + if known_dim is None: + return None + + count *= known_dim + + return count + + +def first_int_attr(node: models.Node, name: str, default: int) -> int: + value = get_attr(node, name) + + if value is match_utils.MISSING: + return default + + if isinstance(value, (list, tuple)): + if not value: + return default + + return int(value[0]) + + return int(value) + + def cache_append_shapes_match(cache_shape: list[Any], new_data_shape: list[Any]) -> bool: for cache_dim, new_dim in zip(cache_shape[:-1], new_data_shape[:-1]): if not match_utils.values_equal(cache_dim, new_dim): @@ -1175,6 +1373,19 @@ def cache_append_shapes_match(cache_shape: list[Any], new_data_shape: list[Any]) return match_utils.values_equal(new_data_shape[-1], 1) +def short_conv_decode_weight_shape_matches(weight_shape: list[Any], hidden_dim: int, window_size: int) -> bool: + if len(weight_shape) == 2: + return match_utils.values_equal(weight_shape[0], hidden_dim) and match_utils.values_equal(weight_shape[1], window_size) + + if len(weight_shape) != 3: + return False + + return ( + match_utils.values_equal(weight_shape, [1, hidden_dim, window_size]) + or match_utils.values_equal(weight_shape, [hidden_dim, 1, window_size]) + ) + + def slice_node_matches(node: models.Node | None, *, axis: int, start: int) -> bool: if node is None: return False @@ -1336,6 +1547,11 @@ def match_same_layer(source: models.Node, graph: models.Graph, fusion: FModels.F "moe_routing_weights_combine": match_moe_routing_weights_combine, "grouped_moe_structure": match_grouped_moe_structure, "cache_roll_append_structure": match_cache_roll_append_structure, + "empty_cache_initializer": match_empty_cache_initializer, + "conv_cache_initialize_structure": match_conv_cache_initialize_structure, + "short_conv_prefill_structure": match_short_conv_prefill_structure, + "short_conv_decode_structure": match_short_conv_decode_structure, + "cache_output_consumers": match_cache_output_consumers, "node_attr_equals": match_node_attr_equals, "node_attrs_equal": match_node_attrs_equal, "input_value_kind": match_input_value_kind, diff --git a/python/cactus/transpiler/IR/match.py b/python/cactus/transpiler/IR/match.py index de925d276..eb233c6f2 100644 --- a/python/cactus/transpiler/IR/match.py +++ b/python/cactus/transpiler/IR/match.py @@ -154,7 +154,7 @@ def match_cache_inputs(node: models.Node, graph: models.Graph, fusion: FModels.F return False - if not match_utils.match_boundary_value(cache_node, ANY, cache_input.tensor_constraints): + if not match_utils.match_cache_boundary_value(cache_node, cache_input): return False return True diff --git a/python/cactus/transpiler/IR/match_utils.py b/python/cactus/transpiler/IR/match_utils.py index 947c4c346..8f0f5ac0d 100644 --- a/python/cactus/transpiler/IR/match_utils.py +++ b/python/cactus/transpiler/IR/match_utils.py @@ -274,6 +274,25 @@ def match_boundary_value(node: models.Node, allowed_value_kinds: tuple[str, ...] return all(match_tensor_constraint(constraint, node, {}) for constraint in tensor_constraints) +def match_cache_boundary_value(node: models.Node, cache_input: FModels.CacheInput) -> bool: + if not match_boundary_value(node, (), cache_input.tensor_constraints): + return False + + if node.cache is None: + return cache_input.optional + + if node.cache.kind != cache_input.cache_kind: + return False + + if cache_input.tensor_role is not None and node.cache.role != cache_input.tensor_role: + return False + + if cache_input.layer_index is not None and node.cache.layer_index != cache_input.layer_index: + return False + + return True + + def all_external_parents_declared(bindings: dict[str, models.Node], declared_input_ids: set[int]) -> bool: internal_ids = {id(node) for node in bindings.values()} diff --git a/python/cactus/transpiler/IR/models.py b/python/cactus/transpiler/IR/models.py index 9cf8a206a..fe4487187 100644 --- a/python/cactus/transpiler/IR/models.py +++ b/python/cactus/transpiler/IR/models.py @@ -7,6 +7,22 @@ from ..Fusions import models as FModels +@dataclass(slots=True) +class CacheAnnotation: + kind: str + role: str + tensor_index: int | None = None + layer_index: int | None = None + shape: tuple[Any, ...] = () + layout: str | None = None + sequence_length: int | None = None + window_size: int | None = None + hidden_dim: int | None = None + num_kv_heads: int | None = None + head_dim: int | None = None + source: str = "inferred" + + @dataclass(slots=True) class Node: index: int @@ -20,6 +36,8 @@ class Node: module_stack: Any | None value_kind: str attrs: dict[str, Any] = field(default_factory=dict) + ir_metadata: dict[str, Any] = field(default_factory=dict) + cache: CacheAnnotation | None = None parents: tuple["Node", ...] = field(default_factory=tuple, repr=False) children: tuple["Node", ...] = field(default_factory=tuple, repr=False) @@ -53,6 +71,7 @@ class Graph: range_constants: str = "" input_specs: tuple[CModels.GraphSpecRecord, ...] = field(default_factory=tuple) output_specs: tuple[CModels.GraphSpecRecord, ...] = field(default_factory=tuple) + cache_annotations: tuple[CacheAnnotation, ...] = field(default_factory=tuple) fusions: list["FusionResult"] = field(default_factory=list) @classmethod @@ -201,6 +220,8 @@ def generate_node(record: CModels.LayerRecord) -> Node: module_stack=record.module_stack, value_kind=infer_value_kind(record), attrs=extract_attrs(record), + ir_metadata=dict(record.ir_metadata), + cache=cache_annotation_from_metadata(record.ir_metadata), ) @@ -228,6 +249,7 @@ def generate_graph(cls: type[Graph], layer_map: CModels.LayerMap) -> Graph: node.parents = tuple(parent_lists[node.name]) node.children = tuple(child_lists[node.name]) + annotate_cache_nodes(nodes) sources = tuple(node for node in nodes if not node.parents) outputs = tuple(node for node in nodes if node.is_output) @@ -246,6 +268,7 @@ def generate_graph(cls: type[Graph], layer_map: CModels.LayerMap) -> Graph: range_constants=layer_map.range_constants, input_specs=tuple(layer_map.input_specs), output_specs=tuple(layer_map.output_specs), + cache_annotations=cache_annotations_from_nodes(nodes), ) @@ -310,6 +333,9 @@ def sanitize_node_name(name: str) -> str: def fused_node_from_result(result: FusionResult, fused_name: str) -> Node: + cache = cache_annotation_from_fusion_result(result) + ir_metadata = fusion_result_metadata(result, cache) + return Node( index=result.source.index, name=fused_name, @@ -320,8 +346,10 @@ def fused_node_from_result(result: FusionResult, fused_name: str) -> Node: users=(), tensor_output_meta=result.source.tensor_output_meta, module_stack=result.source.module_stack, - value_kind=FModels.ValueKind.ACTIVATION, + value_kind=FModels.ValueKind.CACHE_OUTPUT if cache is not None else FModels.ValueKind.ACTIVATION, attrs=dict(result.attrs), + ir_metadata=ir_metadata, + cache=cache, ) @@ -364,6 +392,8 @@ def clone_rewritten_node(node: Node, replacement_names: dict[str, str]) -> Node: module_stack=node.module_stack, value_kind=node.value_kind, attrs=dict(node.attrs), + ir_metadata=dict(node.ir_metadata), + cache=node.cache, ) @@ -424,6 +454,7 @@ def rebuild_graph(nodes: tuple[Node, ...], original_graph: Graph, fusion_results range_constants=original_graph.range_constants, input_specs=original_graph.input_specs, output_specs=original_graph.output_specs, + cache_annotations=cache_annotations_from_nodes(nodes), fusions=list(fusion_results), ) @@ -466,6 +497,7 @@ def graph_to_layer_map(graph: Graph) -> CModels.LayerMap: users=[child.name for child in node.children], tensor_output_meta=node.tensor_output_meta, module_stack=node.module_stack, + ir_metadata=node_ir_metadata(node), ) for index, node in enumerate(topological_sort(graph)) ] @@ -501,3 +533,274 @@ def topological_sort(graph: Graph) -> tuple[Node, ...]: raise ValueError("Cannot topologically sort graph with cycles or missing edges") return tuple(ordered) + + +def annotate_cache_nodes(nodes: tuple[Node, ...]) -> None: + kv_count = 0 + conv_count = 0 + + for node in nodes: + if node.cache is not None: + continue + + tensor_index = past_key_value_index(node) + + if tensor_index is None: + if node.target == "cache_position": + node.cache = CacheAnnotation(kind="position", role=FModels.CacheTensorRole.POSITION, source="placeholder_name") + node.ir_metadata["cache"] = cache_annotation_to_dict(node.cache) + + continue + + shape = tensor_shape(node) + + if len(shape) == 4: + role = FModels.CacheTensorRole.KEY if kv_count % 2 == 0 else FModels.CacheTensorRole.VALUE + node.cache = CacheAnnotation( + kind=FModels.CacheKind.KV, + role=role, + tensor_index=tensor_index, + layer_index=cache_layer_index_from_node(node, kv_count // 2), + shape=tuple(shape), + layout="batch_heads_sequence_head_dim", + num_kv_heads=known_int(shape[1]), + sequence_length=known_int(shape[2]), + head_dim=known_int(shape[3]), + source="past_key_values_placeholder", + ) + kv_count += 1 + elif len(shape) == 3: + node.cache = CacheAnnotation( + kind=FModels.CacheKind.CONV, + role=FModels.CacheTensorRole.STATE, + tensor_index=tensor_index, + layer_index=cache_layer_index_from_node(node, conv_count), + shape=tuple(shape), + layout="batch_hidden_window", + hidden_dim=known_int(shape[1]), + window_size=known_int(shape[2]), + source="past_key_values_placeholder", + ) + conv_count += 1 + + if node.cache is not None: + node.ir_metadata["cache"] = cache_annotation_to_dict(node.cache) + + +def cache_annotation_from_metadata(metadata: dict[str, Any]) -> CacheAnnotation | None: + raw_cache = metadata.get("cache") + + if not isinstance(raw_cache, dict): + return None + + return CacheAnnotation( + kind=str(raw_cache.get("kind")), + role=str(raw_cache.get("role")), + tensor_index=optional_int(raw_cache.get("tensor_index")), + layer_index=optional_int(raw_cache.get("layer_index")), + shape=tuple(raw_cache.get("shape") or ()), + layout=optional_str(raw_cache.get("layout")), + sequence_length=optional_int(raw_cache.get("sequence_length")), + window_size=optional_int(raw_cache.get("window_size")), + hidden_dim=optional_int(raw_cache.get("hidden_dim")), + num_kv_heads=optional_int(raw_cache.get("num_kv_heads")), + head_dim=optional_int(raw_cache.get("head_dim")), + source=str(raw_cache.get("source", "metadata")), + ) + + +def cache_annotation_to_dict(annotation: CacheAnnotation) -> dict[str, Any]: + return { + "kind": annotation.kind, + "role": annotation.role, + "tensor_index": annotation.tensor_index, + "layer_index": annotation.layer_index, + "shape": list(annotation.shape), + "layout": annotation.layout, + "sequence_length": annotation.sequence_length, + "window_size": annotation.window_size, + "hidden_dim": annotation.hidden_dim, + "num_kv_heads": annotation.num_kv_heads, + "head_dim": annotation.head_dim, + "source": annotation.source, + } + + +def cache_annotation_from_fusion_result(result: FusionResult) -> CacheAnnotation | None: + if not result.fusion.graph.cache_outputs: + return None + + cache_output = result.fusion.graph.cache_outputs[0] + base = first_cache_input_annotation(result) + shape = tensor_shape(result.source) + kind = cache_output.cache_kind or (base.kind if base is not None else FModels.CacheKind.KV) + role = cache_output.tensor_role or (base.role if base is not None else FModels.CacheTensorRole.STATE) + + return CacheAnnotation( + kind=kind, + role=role, + tensor_index=base.tensor_index if base is not None else None, + layer_index=base.layer_index if base is not None else infer_layer_index(result), + shape=tuple(shape), + layout=base.layout if base is not None else inferred_cache_layout(kind, shape), + sequence_length=known_int(shape[2]) if len(shape) == 4 else None, + window_size=cache_window_size(result, shape, base), + hidden_dim=cache_hidden_dim(result, shape, base), + num_kv_heads=known_int(shape[1]) if len(shape) == 4 else None, + head_dim=known_int(shape[3]) if len(shape) == 4 else None, + source=f"fusion:{result.fusion_name}", + ) + + +def first_cache_input_annotation(result: FusionResult) -> CacheAnnotation | None: + for node in result.external_inputs: + if node.cache is not None: + return node.cache + + return None + + +def infer_layer_index(result: FusionResult) -> int | None: + for node in result.matched_nodes: + layer_index = layer_index_from_text(f"{node.name} {node.target} {node.module_stack!r}") + + if layer_index is not None: + return layer_index + + return None + + +def cache_layer_index_from_node(node: Node, default: int | None = None) -> int | None: + for candidate in (node, *node.children, *node.parents): + layer_index = layer_index_from_text(f"{candidate.name} {candidate.target} {candidate.module_stack!r}") + + if layer_index is not None: + return layer_index + + return default + + +def layer_index_from_text(value: str) -> int | None: + parts = value.replace("/", ".").replace("_", ".").split(".") + + for index, part in enumerate(parts[:-1]): + if part in {"layers", "layer", "blocks", "block", "h"} and parts[index + 1].isdigit(): + return int(parts[index + 1]) + + return None + + +def cache_window_size(result: FusionResult, shape: tuple[Any, ...], base: CacheAnnotation | None) -> int | None: + if "window_size" in result.attrs and result.attrs["window_size"] is not None: + return int(result.attrs["window_size"]) + + if base is not None and base.window_size is not None: + return base.window_size + + if len(shape) == 3: + return known_int(shape[2]) + + return None + + +def cache_hidden_dim(result: FusionResult, shape: tuple[Any, ...], base: CacheAnnotation | None) -> int | None: + if "hidden_dim" in result.attrs and result.attrs["hidden_dim"] is not None: + return int(result.attrs["hidden_dim"]) + + if base is not None and base.hidden_dim is not None: + return base.hidden_dim + + if len(shape) == 3: + return known_int(shape[1]) + + return None + + +def inferred_cache_layout(kind: str, shape: tuple[Any, ...]) -> str | None: + if kind == FModels.CacheKind.KV and len(shape) == 4: + return "batch_heads_sequence_head_dim" + + if kind == FModels.CacheKind.CONV and len(shape) == 3: + return "batch_hidden_window" + + return None + + +def fusion_result_metadata(result: FusionResult, cache: CacheAnnotation | None) -> dict[str, Any]: + metadata = { + "fusion": { + "name": result.fusion_name, + "cactus_op": result.cactus_op, + "matched_nodes": [node.name for node in result.matched_nodes], + } + } + + if cache is not None: + metadata["cache"] = cache_annotation_to_dict(cache) + + return metadata + + +def cache_annotations_from_nodes(nodes: tuple[Node, ...]) -> tuple[CacheAnnotation, ...]: + return tuple(node.cache for node in nodes if node.cache is not None) + + +def node_ir_metadata(node: Node) -> dict[str, Any]: + metadata = dict(node.ir_metadata) + + if node.cache is not None: + metadata["cache"] = cache_annotation_to_dict(node.cache) + + return metadata + + +def past_key_value_index(node: Node) -> int | None: + for value in (node.target, node.name): + prefix = "past_key_values_" + + if isinstance(value, str) and value.startswith(prefix): + suffix = value[len(prefix):] + + if suffix.isdigit(): + return int(suffix) + + return None + + +def tensor_shape(node: Node) -> tuple[Any, ...]: + if not isinstance(node.tensor_output_meta, dict): + return () + + shape = node.tensor_output_meta.get("shape") + + if isinstance(shape, list): + return tuple(shape) + + if isinstance(shape, tuple): + return shape + + return () + + +def known_int(value: Any) -> int | None: + if isinstance(value, int) and not isinstance(value, bool): + return value + + if isinstance(value, str) and value.isdigit(): + return int(value) + + return None + + +def optional_int(value: Any) -> int | None: + if value is None: + return None + + return known_int(value) + + +def optional_str(value: Any) -> str | None: + if value is None: + return None + + return str(value) diff --git a/python/cactus/transpiler/IR/simplify_ir.py b/python/cactus/transpiler/IR/simplify_ir.py index aa9589005..2852fb3f9 100644 --- a/python/cactus/transpiler/IR/simplify_ir.py +++ b/python/cactus/transpiler/IR/simplify_ir.py @@ -15,16 +15,32 @@ def simplify( input_modalities: tuple[str, ...] = (), fusion_fields: tuple[str, ...] = (), max_fusions: int | None = None, + max_passes: int = 3, ) -> CModels.LayerMap: graph = models.Graph.from_map(layer_map) - simplified_graph = rev_top_sort( - graph, - inference_mode=layer_map.task, - input_modalities=input_modalities, - fusion_fields=fusion_fields, - max_fusions=max_fusions, - ) - return simplified_graph.to_layer_map() + total_fusions = 0 + + for _ in range(max_passes): + remaining_fusions = None if max_fusions is None else max_fusions - total_fusions + + if remaining_fusions is not None and remaining_fusions <= 0: + break + + simplified_graph = rev_top_sort( + graph, + inference_mode=layer_map.task, + input_modalities=input_modalities, + fusion_fields=fusion_fields, + max_fusions=remaining_fusions, + ) + + if not simplified_graph.fusions: + break + + total_fusions += len(simplified_graph.fusions) + graph = simplified_graph + + return graph.to_layer_map() def write_simplified_json( @@ -34,12 +50,14 @@ def write_simplified_json( input_modalities: tuple[str, ...] = (), fusion_fields: tuple[str, ...] = (), max_fusions: int | None = None, + max_passes: int = 3, ) -> CModels.LayerMap: simplified = simplify( layer_map, input_modalities=input_modalities, fusion_fields=fusion_fields, max_fusions=max_fusions, + max_passes=max_passes, ) path = Path(output_path) path.parent.mkdir(parents=True, exist_ok=True) @@ -111,6 +129,9 @@ def try_match_from_node( result = fusion_result_from_bindings(fusion, node, bindings) + if is_noop_fusion(node, result): + continue + if result.consumed_node_names.intersection(consumed_names): continue @@ -119,6 +140,10 @@ def try_match_from_node( return None +def is_noop_fusion(node: models.Node, result: models.FusionResult) -> bool: + return node.target == result.target and len(result.matched_nodes) == 1 + + def bind_matching_fusion( node: models.Node, graph: models.Graph, From 1d1837a0c5ae92fc6931ed52dfeaddb3e5352b61 Mon Sep 17 00:00:00 2001 From: Parkirat Date: Wed, 29 Jul 2026 11:31:11 -0400 Subject: [PATCH 64/84] LFM convs --- python/cactus/transpiler/Fusions/edges.py | 5 ++ python/cactus/transpiler/Fusions/fusions.py | 77 +++++++++++++++++++-- python/cactus/transpiler/Fusions/nodes.py | 7 +- 3 files changed, 79 insertions(+), 10 deletions(-) diff --git a/python/cactus/transpiler/Fusions/edges.py b/python/cactus/transpiler/Fusions/edges.py index a32dbcd63..5f256b22c 100644 --- a/python/cactus/transpiler/Fusions/edges.py +++ b/python/cactus/transpiler/Fusions/edges.py @@ -122,6 +122,9 @@ def _index_edges_by_dest(edges: dict[str, M.FusionEdge]) -> dict[str, tuple[str, "rope_sin_mul_to_add": _edge("rope_sin_mul", "rope_add", 1), "rope_rotate_cat_to_sin_mul": _edge("rope_rotate_cat", "rope_sin_mul", 0), "conv_to_bias_add": _edge("conv", "conv_bias_add", 0), + "conv_to_slice": _edge("conv", "slice", 0), + "conv_cache_append_to_multiply": _edge("conv_cache_append", "multiply", 0), + "multiply_to_sum": _edge("multiply", "sum", 0), "lstm_gate_mm_to_add": _edge("lstm_gate_mm", "lstm_gate_add", 0), "lstm_recurrent_mm_to_add": _edge("lstm_recurrent_mm", "lstm_gate_add", 1), "lstm_add_to_sigmoid": _edge("lstm_gate_add", "lstm_sigmoid", 0), @@ -304,6 +307,8 @@ def _index_edges_by_dest(edges: dict[str, M.FusionEdge]) -> dict[str, tuple[str, "rope_sin_mul_to_add", ), "conv_bias": ("conv_to_bias_add",), + "short_conv_prefill": ("conv_to_slice",), + "short_conv_decode": ("multiply_to_sum",), "lstm_cell": ( "lstm_gate_mm_to_add", "lstm_recurrent_mm_to_add", diff --git a/python/cactus/transpiler/Fusions/fusions.py b/python/cactus/transpiler/Fusions/fusions.py index d094a1ae8..a70d7625b 100644 --- a/python/cactus/transpiler/Fusions/fusions.py +++ b/python/cactus/transpiler/Fusions/fusions.py @@ -609,6 +609,41 @@ def _index_by_root_op( attr_captures=(M.AttrCapture("stride", "conv", "stride", default=1, required=False), M.AttrCapture("padding", "conv", "padding", default=0, required=False), M.AttrCapture("dilation", "conv", "dilation", default=1, required=False), M.AttrCapture("groups", "conv", "groups", default=1, required=False)), ) +LFM_SHORT_CONV_PREFILL_GRAPH = _graph( + "lfm_short_conv_prefill", + "slice", + ("conv", "slice"), + edge_names=E.EDGE_GROUPS["short_conv_prefill"], + inputs=(_input("x", "conv", 0), _input("weight", "conv", 1), _input("bias", "conv", 2, optional=True)), + attr_captures=( + M.AttrCapture("kernel_size", default=3, required=False), + M.AttrCapture("dilation", "conv", "dilation", default=1, required=False), + M.AttrCapture("causal", default=True, required=False), + M.AttrCapture("layout", default="batch_hidden_sequence", required=False), + ), + constraints={ + "short_conv_prefill_structure": {"conv_node": "conv", "slice_node": "slice", "x_role": "x", "weight_role": "weight"}, + }, +) + +LFM_SHORT_CONV_DECODE_GRAPH = _graph( + "lfm_short_conv_decode", + "sum", + ("multiply", "sum"), + edge_names=E.EDGE_GROUPS["short_conv_decode"], + inputs=( + _input("cache_window", "multiply", 0, allowed_value_kinds=(M.ValueKind.CACHE_OUTPUT, M.ValueKind.CACHE_STATE, M.ValueKind.ACTIVATION)), + _input("weight", "multiply", 1, allowed_value_kinds=(M.ValueKind.PARAMETER, M.ValueKind.BUFFER, M.ValueKind.ACTIVATION)), + ), + attr_captures=( + M.AttrCapture("window_size", default=3, required=False), + M.AttrCapture("layout", default="batch_hidden_sequence", required=False), + ), + constraints={ + "short_conv_decode_structure": {"cache_window_role": "cache_window", "multiply_node": "multiply", "sum_node": "sum", "weight_role": "weight"}, + }, +) + CONV_BIAS_GRAPH = _graph( "conv_bias", "conv_bias_add", @@ -623,12 +658,31 @@ def _index_by_root_op( "kv_cache_append", ("kv_cache_append",), inputs=( - _input("new_kv", "kv_cache_append", 0), - _input("cache_state", "kv_cache_append", 1, allowed_value_kinds=(M.ValueKind.CACHE_INPUT, M.ValueKind.CACHE_STATE, M.ValueKind.CACHE_OUTPUT)), + _input("new_kv", "kv_cache_append", 1), + _input("cache_state", "kv_cache_append", 0, allowed_value_kinds=(M.ValueKind.CACHE_INPUT, M.ValueKind.CACHE_STATE, M.ValueKind.CACHE_OUTPUT)), ), - cache_inputs=(_cache_input("cache_state", "kv_cache_append", 1, tensor_role=M.CacheTensorRole.STATE),), - cache_outputs=(_cache_output("updated_cache_state", "kv_cache_append", tensor_role=M.CacheTensorRole.STATE),), + cache_inputs=(_cache_input("cache_state", "kv_cache_append", 0),), + cache_outputs=(_cache_output("updated_cache_state", "kv_cache_append"),), cache_mutations=(_cache_mutation("append_kv", read_roles=("cache_state",), write_roles=("updated_cache_state",)),), + constraints={ + "cache_output_consumers": {"node": "kv_cache_append", "allow_output_child": True}, + }, +) + +KV_CACHE_INITIAL_APPEND_GRAPH = _graph( + "kv_cache_initial_append", + "kv_cache_append", + ("kv_cache_append",), + inputs=( + _input("new_kv", "kv_cache_append", 1), + _input("empty_cache", "kv_cache_append", 0, metadata={"drop_after_fusion": True}), + ), + cache_outputs=(_cache_output("initialized_cache_state", "kv_cache_append"),), + cache_mutations=(_cache_mutation("initialize_append_kv", write_roles=("initialized_cache_state",)),), + constraints={ + "empty_cache_initializer": {"empty_role": "empty_cache", "new_value_role": "new_kv", "output_node": "kv_cache_append"}, + "cache_output_consumers": {"node": "kv_cache_append", "allow_output_child": True}, + }, ) ATTENTION_CACHED_GRAPH = _graph( @@ -894,11 +948,14 @@ def _index_by_root_op( "conv_cache_initialize", ("conv_cache_initialize",), inputs=( - _input("rows", "conv_cache_initialize", 0), - _input("cache_state", "conv_cache_initialize", 1, allowed_value_kinds=(M.ValueKind.CACHE_INPUT, M.ValueKind.CACHE_STATE, M.ValueKind.CACHE_OUTPUT)), + _input("rows", "conv_cache_initialize", 1), + _input("empty_cache", "conv_cache_initialize", 0, metadata={"drop_after_fusion": True}), ), - cache_inputs=(_cache_input("cache_state", "conv_cache_initialize", 1, cache_kind=M.CacheKind.CONV, tensor_role=M.CacheTensorRole.STATE),), cache_outputs=(_cache_output("initialized_cache_state", "conv_cache_initialize", cache_kind=M.CacheKind.CONV, tensor_role=M.CacheTensorRole.STATE),), + cache_mutations=(_cache_mutation("initialize_conv_state", cache_kind=M.CacheKind.CONV, write_roles=("initialized_cache_state",)),), + constraints={ + "conv_cache_initialize_structure": {"rows_role": "rows", "empty_role": "empty_cache", "output_node": "conv_cache_initialize"}, + }, ) RECURRENT_CACHE_STATE_GRAPH = _graph( @@ -984,8 +1041,11 @@ def _index_by_root_op( "lfm_bmm_masked_attention": LFM_BMM_MASKED_ATTENTION_GRAPH, "rope": ROPE_GRAPH, "conv": CONV_GRAPH, + "lfm_short_conv_prefill": LFM_SHORT_CONV_PREFILL_GRAPH, + "lfm_short_conv_decode": LFM_SHORT_CONV_DECODE_GRAPH, "conv_bias": CONV_BIAS_GRAPH, "kv_cache_append": KV_CACHE_APPEND_GRAPH, + "kv_cache_initial_append": KV_CACHE_INITIAL_APPEND_GRAPH, "attention_cached": ATTENTION_CACHED_GRAPH, "moe_layer_gated": MOE_GATED_GRAPH, "lfm_grouped_moe": LFM_GROUPED_MOE_GRAPH, @@ -1035,6 +1095,8 @@ def _index_by_root_op( "rope": _definition("rope", "rope", ROPE_GRAPH, fusion_fields=("generic", "rope", "gemma4_rope", "qwen2_5_rope")), "conv": _definition("conv", "conv1d", CONV_GRAPH, fusion_fields=("generic", "conv", "audio", "vision"), metadata=_required_attrs(padding=0, dilation=1, groups=1)), "conv_bias": _definition("conv_bias", "conv1d", CONV_BIAS_GRAPH, fusion_fields=("generic", "conv", "audio", "vision")), + "lfm_short_conv_prefill": _definition("lfm_short_conv_prefill", "conv1d_causal", LFM_SHORT_CONV_PREFILL_GRAPH, fusion_fields=("generic", "conv", "cache", "lfm_moe"), supported_inference_modes=("prefill_with_cache",)), + "lfm_short_conv_decode": _definition("lfm_short_conv_decode", "lfm_short_conv_decode", LFM_SHORT_CONV_DECODE_GRAPH, fusion_fields=("generic", "conv", "cache", "lfm_moe"), supported_inference_modes=("decode_with_cache",)), "conv1d_k3": _definition("conv1d_k3", "conv1d_k3", CONV_GRAPH, fusion_fields=("generic", "conv", "audio"), metadata=_required_attrs(kernel_size=3, stride=1, padding=0, dilation=1, groups=1)), "conv1d_k7s3": _definition("conv1d_k7s3", "conv1d_k7s3", CONV_BIAS_GRAPH, fusion_fields=("generic", "conv", "audio"), metadata=_required_attrs(kernel_size=7, stride=3)), "conv1d_causal": _definition("conv1d_causal", "conv1d_causal", CONV_GRAPH, fusion_fields=("generic", "conv", "cache"), metadata=_required_attrs(causal=True)), @@ -1044,6 +1106,7 @@ def _index_by_root_op( "conv2d_depthwise_k3s2p1": _definition("conv2d_depthwise_k3s2p1", "conv2d_depthwise_k3s2p1", CONV_GRAPH, fusion_fields=("generic", "conv", "vision"), metadata=_required_attrs(kernel_size=3, stride=2, padding=1, depthwise=True)), "conv2d_pointwise_1x1": _definition("conv2d_pointwise_1x1", "conv2d_pointwise_1x1", CONV_GRAPH, fusion_fields=("generic", "conv", "vision"), metadata=_required_attrs(kernel_size=1)), "kv_cache_append": _definition("kv_cache_append", "kv_cache_append", KV_CACHE_APPEND_GRAPH, fusion_fields=("generic", "cache"), supported_inference_modes=("prefill_with_cache", "decode_with_cache")), + "kv_cache_initial_append": _definition("kv_cache_initial_append", "kv_cache_append", KV_CACHE_INITIAL_APPEND_GRAPH, fusion_fields=("generic", "cache"), supported_inference_modes=("prefill_with_cache",)), "attention_cached": _definition("attention_cached", "attention_cached", ATTENTION_CACHED_GRAPH, fusion_fields=("generic", "attention", "cache"), supported_inference_modes=("decode_with_cache",)), "moe_layer_gated": _definition("moe_layer_gated", "moe_layer_gated", MOE_GATED_GRAPH, fusion_fields=("generic", "moe")), "lfm_grouped_moe": _definition("lfm_grouped_moe", "moe_layer_gated", LFM_GROUPED_MOE_GRAPH, fusion_fields=("generic", "moe", "lfm_moe")), diff --git a/python/cactus/transpiler/Fusions/nodes.py b/python/cactus/transpiler/Fusions/nodes.py index 8a81c67d2..b7cecbc43 100644 --- a/python/cactus/transpiler/Fusions/nodes.py +++ b/python/cactus/transpiler/Fusions/nodes.py @@ -27,7 +27,7 @@ "sin": ("aten.sin.default",), "log": ("aten.log.default",), "mean": ("aten.mean.dim", "aten.mean.default"), - "sum": ("aten.sum.dim_IntList", "aten.sum.default"), + "sum": ("aten.sum.dim_IntList", "aten.sum.default", "cactus.sum"), "var": ("aten.var.dim", "aten.var.correction"), "max": ("aten.max.dim", "aten.amax.default"), "min": ("aten.min.dim", "aten.amin.default"), @@ -44,6 +44,7 @@ "gather": ("aten.gather.default", "aten.index_select.default", "cactus.gather"), "index_put": ("aten.index_put.default", "cactus.index_put"), "cat": ("aten.cat.default", "cactus.cat"), + "conv_cache_append_op": ("aten.cat.default", "cactus.conv_cache_append"), "scatter": ("aten.scatter.src", "aten.scatter.value", "aten.scatter_reduce.two"), "histc": ("aten.histc.default",), "cumsum": ("aten.cumsum.default",), @@ -253,8 +254,8 @@ def _expert_node(name: str, op_key: str) -> M.FusionNode: "rope_add": _node("rope_add", "add"), "conv": _node("conv", "conv"), "conv_bias_add": _node("conv_bias_add", "add"), - "conv_cache_append": _node("conv_cache_append", "cat"), - "conv_cache_initialize": _node("conv_cache_initialize", "copy"), + "conv_cache_append": _node("conv_cache_append", "conv_cache_append_op"), + "conv_cache_initialize": _node("conv_cache_initialize", "copy_write"), "lfm_conv_cache_clone": _node("lfm_conv_cache_clone", "copy"), "lfm_conv_cache_roll_arange": _node("lfm_conv_cache_roll_arange", "arange"), "lfm_conv_cache_roll_add": _node("lfm_conv_cache_roll_add", "add"), From e07c15801d1b79051d996aae6edb89984a7da745 Mon Sep 17 00:00:00 2001 From: Parkirat Date: Wed, 29 Jul 2026 12:19:40 -0400 Subject: [PATCH 65/84] Added runtime plan --- python/cactus/transpiler/Converter/models.py | 1 + .../cactus/transpiler/Generator/constants.py | 1 + .../cactus/transpiler/Generator/generate.py | 384 +++++++++++++++- python/cactus/transpiler/Generator/models.py | 75 +++- .../transpiler/ModelProfiles/combinations.py | 8 +- .../cactus/transpiler/RuntimePlan/models.py | 420 ++++++++++++++++++ 6 files changed, 856 insertions(+), 33 deletions(-) create mode 100644 python/cactus/transpiler/RuntimePlan/models.py diff --git a/python/cactus/transpiler/Converter/models.py b/python/cactus/transpiler/Converter/models.py index 989a4c2bd..96cb36f5e 100644 --- a/python/cactus/transpiler/Converter/models.py +++ b/python/cactus/transpiler/Converter/models.py @@ -51,6 +51,7 @@ class LayerRecord(BaseModel): users: list[str] tensor_output_meta: Any | None module_stack: Any | None + ir_metadata: dict[str, Any] = Field(default_factory=dict) #Builds a LayerRecord from a torch.fx.Node. @classmethod diff --git a/python/cactus/transpiler/Generator/constants.py b/python/cactus/transpiler/Generator/constants.py index 5ad7b2030..520b2e790 100644 --- a/python/cactus/transpiler/Generator/constants.py +++ b/python/cactus/transpiler/Generator/constants.py @@ -288,6 +288,7 @@ "cactus.lstm_cell", "cactus.gated_deltanet_decode", "cactus.gated_deltanet_prefill", + "cactus.lfm_short_conv_decode", "cactus.rel_pos_bias", "cactus.sample", "cactus.scatter_topk", diff --git a/python/cactus/transpiler/Generator/generate.py b/python/cactus/transpiler/Generator/generate.py index f620ebb33..037e81586 100644 --- a/python/cactus/transpiler/Generator/generate.py +++ b/python/cactus/transpiler/Generator/generate.py @@ -8,7 +8,9 @@ from . import constants from . import models from ..Converter import models as CModels +from ..Fusions import models as FModels from ..IR import models as IRModels +from ..RuntimePlan import models as RPModels class UnsupportedLoweringError(NotImplementedError): @@ -258,23 +260,67 @@ def lower_node(context: models.GenerationContext, node: IRModels.Node) -> None: def lower_placeholder(context: models.GenerationContext, node: IRModels.Node) -> Any: + if should_lower_cache_placeholder_as_state(node): + return lower_cache_placeholder(context, node) + shape, dynamic_dims = graph_input_shape(node) dtype = cactus_precision(context.graph, tensor_dtype(node)) tensor = context.graph.input(shape, dtype=dtype, dynamic_dims=dynamic_dims if any(dynamic_dims) else None) + logical_name = logical_input_name(context.component.ir_graph, node) if node.value_kind in constants.WEIGHT_VALUE_KINDS: - bind_weight_placeholder(context, node, tensor) + bind_weight_placeholder(context, node, tensor, logical_name) return tensor - context.component.add_runtime_input(tensor) + context.component.add_runtime_input(tensor, logical_name) return tensor -def bind_weight_placeholder(context: models.GenerationContext, node: IRModels.Node, tensor: Any) -> None: +def should_lower_cache_placeholder_as_state(node: IRModels.Node) -> bool: + if node.cache is None: + return False + + if node.cache.kind == FModels.CacheKind.CONV: + return any(child.target in {"cactus.conv_cache_append", "cactus.conv_cache_initialize"} for child in node.children) + + if node.cache.kind == FModels.CacheKind.KV: + return bool(node.children) and all(child.target in {"cactus.kv_cache_append", "cactus.attention_cached"} for child in node.children) + + return False + + +def lower_cache_placeholder(context: models.GenerationContext, node: IRModels.Node) -> Any: + annotation = require_cache_annotation(node) + + if annotation.kind == FModels.CacheKind.KV: + cache_state = context.graph.kv_cache_state( + kv_cache_capacity(context, annotation), + required_cache_int(annotation.num_kv_heads, node, "num_kv_heads"), + required_cache_int(annotation.head_dim, node, "head_dim"), + window_size=int(annotation.window_size or 0), + sink_size=0, + num_slots=1, + ) + record_cache_state_binding(context, node, cache_state, annotation) + return cache_state + + if annotation.kind == FModels.CacheKind.CONV: + cache_state = context.graph.conv_cache_state( + required_cache_int(annotation.window_size, node, "window_size"), + required_cache_int(annotation.hidden_dim, node, "hidden_dim"), + ) + record_cache_state_binding(context, node, cache_state, annotation) + return cache_state + + raise UnsupportedLoweringError(f"{node.name}: unsupported cache placeholder kind {annotation.kind}") + + +def bind_weight_placeholder(context: models.GenerationContext, node: IRModels.Node, tensor: Any, logical_name: str | None = None) -> None: resolver = context.component.weight_resolver if resolver is None: context.component.warn(f"{node.name}: weight lowered as graph input because no weights_dir was provided") + context.component.add_runtime_input(tensor, logical_name) return source_target = resolver.source_target_for(node.name) or node.target @@ -287,7 +333,7 @@ def bind_weight_placeholder(context: models.GenerationContext, node: IRModels.No raise UnsupportedLoweringError(message) context.component.warn(message) - context.component.add_runtime_input(tensor) + context.component.add_runtime_input(tensor, logical_name) return node_id = models.tensor_node_id(tensor) @@ -317,13 +363,16 @@ def bind_weight_placeholder(context: models.GenerationContext, node: IRModels.No def lower_output(context: models.GenerationContext, node: IRModels.Node) -> Any: refs = IRModels.extract_node_refs((node.args, node.kwargs)) values = tuple(context.require(ref) for ref in refs) + output_index = 0 - for value in values: + for ref, value in zip(refs, values): if isinstance(value, (tuple, list)): for item in value: - context.component.add_output(item) + context.component.add_output(item, logical_output_name(context.component.ir_graph, output_index, ref)) + output_index += 1 else: - context.component.add_output(value) + context.component.add_output(value, logical_output_name(context.component.ir_graph, output_index, ref)) + output_index += 1 if len(values) == 1: return values[0] @@ -331,6 +380,98 @@ def lower_output(context: models.GenerationContext, node: IRModels.Node) -> Any: return values +def logical_input_name(graph: IRModels.Graph, node: IRModels.Node) -> str: + metadata_name = logical_name_from_metadata(node.ir_metadata, ("logical_input", "logical_name", "runtime_input")) + + if metadata_name: + return metadata_name + + spec_name = logical_input_name_from_specs(graph, node.name) + + if spec_name: + return spec_name + + if node.cache is not None: + return cache_logical_name(node) + + target_name = clean_logical_name(node.target) + + if target_name is not None: + return target_name + + return node.name + + +def logical_output_name(graph: IRModels.Graph, output_index: int, source_ref: str | None = None) -> str: + source = graph.nodes_map.get(source_ref) if source_ref is not None else None + + if source is not None: + metadata_name = logical_name_from_metadata(source.ir_metadata, ("logical_output", "logical_name", "runtime_output")) + + if metadata_name: + return metadata_name + + if source.cache is not None: + return cache_logical_name(source) + + if "logits" in source.name or "lm_head" in source.target: + return "logits" + + spec_name = logical_output_name_from_specs(graph, output_index) + + if spec_name: + return spec_name + + return f"output_{output_index}" + + +def logical_input_name_from_specs(graph: IRModels.Graph, node_name: str) -> str | None: + for spec in graph.input_specs: + if spec.arg_name != node_name: + continue + + return clean_logical_name(spec.target) or clean_logical_name(spec.arg_name) + + return None + + +def logical_output_name_from_specs(graph: IRModels.Graph, output_index: int) -> str | None: + if output_index < 0 or output_index >= len(graph.output_specs): + return None + + spec = graph.output_specs[output_index] + return clean_logical_name(spec.target) or clean_logical_name(spec.arg_name) + + +def logical_name_from_metadata(metadata: dict[str, Any], keys: tuple[str, ...]) -> str | None: + for key in keys: + value = clean_logical_name(metadata.get(key)) + + if value is not None: + return value + + return None + + +def clean_logical_name(value: Any) -> str | None: + if value is None: + return None + + name = str(value) + + if not name or name == "None": + return None + + return name + + +def cache_logical_name(node: IRModels.Node) -> str: + if node.cache is not None and node.cache.tensor_index is not None: + return f"past_key_values_{node.cache.tensor_index}" + + return node.name + + def lower_binary(context: models.GenerationContext, node: IRModels.Node) -> Any: inputs = context.inputs_for(node) method = constants.BINARY_TARGETS[node.target] @@ -509,7 +650,13 @@ def lower_cat(context: models.GenerationContext, node: IRModels.Node) -> Any: if not inputs: raise unsupported_arity(node, 0, "at least one input") - return context.graph.cat(inputs, axis=int(node.attrs.get("axis", node.attrs.get("dim", 0)))) + passthrough = empty_cat_passthrough(node, inputs) + + if passthrough is not None: + return passthrough + + axis = normalize_dim(int(node.attrs.get("axis", node.attrs.get("dim", 0))), cat_rank(node)) + return context.graph.cat(inputs, axis=axis) def lower_matmul(context: models.GenerationContext, node: IRModels.Node) -> Any: @@ -525,6 +672,39 @@ def lower_matmul(context: models.GenerationContext, node: IRModels.Node) -> Any: ) +def empty_cat_passthrough(node: IRModels.Node, inputs: tuple[Any, ...]) -> Any | None: + target_shape = meta_shape(node) + non_empty_inputs = [] + + for input_value, parent in zip(inputs, node.parents): + parent_shape = meta_shape(parent) + parent_elements = element_count(parent_shape) + + if parent_elements == 0: + continue + + non_empty_inputs.append((input_value, parent_shape)) + + if len(non_empty_inputs) != 1: + return None + + input_value, input_shape = non_empty_inputs[0] + + if tuple(input_shape) == tuple(target_shape): + return input_value + + return None + + +def cat_rank(node: IRModels.Node) -> int: + target_shape = meta_shape(node) + + if target_shape: + return len(target_shape) + + return len(meta_shape(node.parents[0])) if node.parents else 0 + + def lower_bmm(context: models.GenerationContext, node: IRModels.Node, inputs: tuple[Any, ...]) -> Any: left_shape = meta_shape(node.parents[0]) right_shape = meta_shape(node.parents[1]) @@ -653,7 +833,7 @@ def lower_constant_input(context: models.GenerationContext, node: IRModels.Node) shape, dynamic_dims = graph_input_shape(node) dtype = cactus_precision(context.graph, tensor_dtype(node)) tensor = context.graph.input(shape, dtype=dtype, dynamic_dims=dynamic_dims if any(dynamic_dims) else None) - context.component.add_runtime_input(tensor) + context.component.add_runtime_input(tensor, node.name) return tensor @@ -664,7 +844,12 @@ def lower_neg(context: models.GenerationContext, node: IRModels.Node) -> Any: def lower_softmax(context: models.GenerationContext, node: IRModels.Node) -> Any: inputs = require_input_count(context, node, 1) - return context.graph.softmax(inputs[0], axis=axis_attr(node, default=-1)) + axis = axis_attr(node, default=-1) + + if axis is None: + raise UnsupportedLoweringError(f"{node.name}: softmax lowering missing axis/dim attr") + + return context.graph.softmax(inputs[0], axis=normalize_dim(axis, len(meta_shape(node.parents[0])))) def lower_topk(context: models.GenerationContext, node: IRModels.Node) -> Any: @@ -759,7 +944,17 @@ def lower_named_conv(context: models.GenerationContext, node: IRModels.Node, met return lower_generic_conv1d(context, node, inputs, bias) if method == "conv1d_causal": - return context.graph.conv1d_causal(inputs[0], inputs[1], kernel_size=int(node.attrs.get("kernel_size", 0)), dilation=first_int(node.attrs.get("dilation"), 1)) + x = inputs[0] + + if node.attrs.get("layout") == "batch_hidden_sequence": + x = context.graph.permute(x, (0, 2, 1)) + + output = context.graph.conv1d_causal(x, inputs[1], kernel_size=int(node.attrs.get("kernel_size", 0)), dilation=first_int(node.attrs.get("dilation"), 1)) + + if node.attrs.get("layout") == "batch_hidden_sequence": + return context.graph.permute(output, (0, 2, 1)) + + return output if method == "conv1d_k3": require_plain_conv1d_attrs(node, method) @@ -865,18 +1060,33 @@ def lower_cache(context: models.GenerationContext, node: IRModels.Node) -> Any: target = node.target if target == "cactus.kv_cache_append": + if len(inputs) == 1: + cache_state = create_cache_state_for_cache_output(context, node) + context.graph.kv_cache_append(inputs[0], cache_state, window_size=int(node.attrs.get("window_size", 0)), sink_size=int(node.attrs.get("sink_size", 0))) + return cache_state + require_len(node, inputs, 2) require_cache_state_parent(node, 1, "cactus.kv_cache_state") - return context.graph.kv_cache_append(inputs[0], inputs[1], window_size=int(node.attrs.get("window_size", 0)), sink_size=int(node.attrs.get("sink_size", 0))) + context.graph.kv_cache_append(inputs[0], inputs[1], window_size=int(node.attrs.get("window_size", 0)), sink_size=int(node.attrs.get("sink_size", 0))) + return inputs[1] if target == "cactus.conv_cache_append": require_len(node, inputs, 2) require_cache_state_parent(node, 1, "cactus.conv_cache_state") - return context.graph.conv_cache_append(inputs[0], inputs[1]) + window = context.graph.conv_cache_append(inputs[0], inputs[1]) + return conv_cache_window_to_ir_layout(context, node, window) if target == "cactus.conv_cache_initialize": + if len(inputs) == 1: + cache_state = create_cache_state_for_cache_output(context, node) + rows = conv_cache_rows_for_native(context, node, inputs[0]) + context.graph.conv_cache_initialize(rows, cache_state) + return cache_state + require_len(node, inputs, 2) - return context.graph.conv_cache_initialize(inputs[0], inputs[1]) + rows = conv_cache_rows_for_native(context, node, inputs[0]) + context.graph.conv_cache_initialize(rows, inputs[1]) + return inputs[1] if target == "cactus.recurrent_cache_write": require_len(node, inputs, 2) @@ -902,12 +1112,113 @@ def lower_cache(context: models.GenerationContext, node: IRModels.Node) -> Any: def require_cache_state_parent(node: IRModels.Node, parent_index: int, expected_target: str) -> None: - if parent_index >= len(node.parents) or node.parents[parent_index].target != expected_target: - actual = node.parents[parent_index].target if parent_index < len(node.parents) else "" - raise UnsupportedLoweringError( - f"{node.name}: {node.target} requires parent {parent_index} to be {expected_target}; got {actual}. " - "Raw HF cache tensors need a cache-state bridge before native cache append lowering." + if parent_index >= len(node.parents): + raise UnsupportedLoweringError(f"{node.name}: {node.target} missing cache parent {parent_index}") + + parent = node.parents[parent_index] + + if parent.target == expected_target: + return + + if expected_target == "cactus.conv_cache_state" and parent.cache is not None and parent.cache.kind == FModels.CacheKind.CONV: + return + + if expected_target == "cactus.kv_cache_state" and parent.cache is not None and parent.cache.kind == FModels.CacheKind.KV: + return + + raise UnsupportedLoweringError( + f"{node.name}: {node.target} requires parent {parent_index} to be {expected_target}; got {parent.target}. " + "Raw HF cache tensors need a cache-state bridge before native cache append lowering." + ) + + +def require_cache_annotation(node: IRModels.Node) -> IRModels.CacheAnnotation: + if node.cache is None: + raise UnsupportedLoweringError(f"{node.name}: cache lowering requires cache annotation metadata") + + return node.cache + + +def required_cache_int(value: int | None, node: IRModels.Node, name: str) -> int: + if value is None: + raise UnsupportedLoweringError(f"{node.name}: cache lowering missing {name}") + + return int(value) + + +def kv_cache_capacity(context: models.GenerationContext, annotation: IRModels.CacheAnnotation) -> int: + capacity = int(annotation.sequence_length or 1) + + if context.component.ir_graph.task == "decode_with_cache": + capacity += 1 + + return max(capacity, 1) + + +def create_cache_state_for_cache_output(context: models.GenerationContext, node: IRModels.Node) -> Any: + annotation = require_cache_annotation(node) + + if annotation.kind == FModels.CacheKind.KV: + cache_state = context.graph.kv_cache_state( + kv_cache_capacity(context, annotation), + required_cache_int(annotation.num_kv_heads, node, "num_kv_heads"), + required_cache_int(annotation.head_dim, node, "head_dim"), + window_size=int(annotation.window_size or 0), + sink_size=0, + num_slots=1, + ) + record_cache_state_binding(context, node, cache_state, annotation) + return cache_state + + if annotation.kind == FModels.CacheKind.CONV: + cache_state = context.graph.conv_cache_state( + required_cache_int(annotation.window_size, node, "window_size"), + required_cache_int(annotation.hidden_dim, node, "hidden_dim"), ) + record_cache_state_binding(context, node, cache_state, annotation) + return cache_state + + raise UnsupportedLoweringError(f"{node.name}: unsupported cache output kind {annotation.kind}") + + +def record_cache_state_binding(context: models.GenerationContext, node: IRModels.Node, cache_state: Any, annotation: IRModels.CacheAnnotation) -> None: + node_id = models.tensor_node_id(cache_state) + + if node_id is None: + raise UnsupportedLoweringError(f"{node.name}: cache state tensor does not expose a Cactus node id") + + context.component.add_cache_state_binding(RPModels.cache_state_binding_from_annotation(annotation, node_id)) + + +def conv_cache_rows_for_native(context: models.GenerationContext, node: IRModels.Node, rows: Any) -> Any: + rows_shape = meta_shape(node.parents[0]) if node.parents else output_shape(node) + + if len(rows_shape) == 3: + batch, hidden_dim, window_size = rows_shape + + if batch != 1: + raise UnsupportedLoweringError(f"{node.name}: conv cache initialize currently supports batch size 1") + + return context.graph.reshape(context.graph.permute(rows, (0, 2, 1)), (int(window_size), int(hidden_dim))) + + if len(rows_shape) == 2: + return rows + + raise UnsupportedLoweringError(f"{node.name}: conv cache initialize rows must be rank 2 or 3, got shape {rows_shape}") + + +def conv_cache_window_to_ir_layout(context: models.GenerationContext, node: IRModels.Node, window: Any) -> Any: + target_shape = output_shape(node) + + if len(target_shape) != 3: + return window + + batch, hidden_dim, window_size = target_shape + + if batch != 1: + raise UnsupportedLoweringError(f"{node.name}: conv cache append layout bridge currently supports batch size 1") + + return context.graph.reshape(context.graph.transpose(window), (int(batch), int(hidden_dim), int(window_size))) def lower_moe(context: models.GenerationContext, node: IRModels.Node) -> Any: @@ -986,6 +1297,12 @@ def lower_special_cactus(context: models.GenerationContext, node: IRModels.Node) require_len(node, inputs, 6) return context.graph.gated_deltanet_prefill(*inputs[:6], chunk_size=int(node.attrs.get("chunk_size", 1)), scale=float(node.attrs.get("scale", 1.0))) + if target == "cactus.lfm_short_conv_decode": + require_len(node, inputs, 2) + weight = lfm_short_conv_decode_weight(context, node, inputs[1]) + product = context.graph.multiply(inputs[0], weight) + return context.graph.sum(product, axis=normalize_dim(-1, len(meta_shape(node.parents[0])))) + if target == "cactus.rel_pos_bias": require_len(node, inputs, 2) return context.graph.rel_pos_bias(inputs[0], inputs[1], scale=float(node.attrs.get("scale") or 1.0)) @@ -1058,6 +1375,35 @@ def lower_special_cactus(context: models.GenerationContext, node: IRModels.Node) raise UnsupportedLoweringError(f"{node.name}: unsupported special Cactus target {target}") +def lfm_short_conv_decode_weight(context: models.GenerationContext, node: IRModels.Node, weight: Any) -> Any: + cache_window_shape = meta_shape(node.parents[0]) if node.parents else () + weight_shape = meta_shape(node.parents[1]) if len(node.parents) > 1 else () + + if len(cache_window_shape) != 3: + return weight + + batch_size, hidden_dim, window_size = cache_window_shape + + if not all(isinstance(dim, int) for dim in cache_window_shape): + return weight + + if tuple(weight_shape) == tuple(cache_window_shape): + return weight + + if len(weight_shape) == 2 and tuple(weight_shape) == (hidden_dim, window_size): + return context.graph.reshape(weight, (batch_size, hidden_dim, window_size)) + + if len(weight_shape) == 3 and tuple(weight_shape) in { + (1, hidden_dim, window_size), + (hidden_dim, 1, window_size), + }: + return context.graph.reshape(weight, (batch_size, hidden_dim, window_size)) + + raise UnsupportedLoweringError( + f"{node.name}: lfm short conv decode weight shape {weight_shape} is incompatible with cache window {cache_window_shape}" + ) + + def lower_unsupported_semantic(context: models.GenerationContext, node: IRModels.Node) -> Any: raise UnsupportedLoweringError(f"{node.name}: {node.target} has no safe Cactus lowering yet") diff --git a/python/cactus/transpiler/Generator/models.py b/python/cactus/transpiler/Generator/models.py index 2728e428f..9267d56eb 100644 --- a/python/cactus/transpiler/Generator/models.py +++ b/python/cactus/transpiler/Generator/models.py @@ -7,6 +7,7 @@ from typing import Any, Callable, Mapping from ..IR import models as IRModels +from ..RuntimePlan import models as RPModels GraphTensor = Any @@ -54,8 +55,12 @@ class ComponentGraph: graph: CactusGraph | None = None weight_resolver: "WeightResolver | None" = None weight_bindings: list["WeightBinding"] = field(default_factory=list) + cache_state_bindings: list[RPModels.CacheStateBinding] = field(default_factory=list) runtime_input_ids: list[int] = field(default_factory=list) + logical_inputs: list[str] = field(default_factory=list) output_node_ids: list[int] = field(default_factory=list) + logical_outputs: list[str] = field(default_factory=list) + metadata: dict[str, str] = field(default_factory=dict) unsupported_nodes: list[str] = field(default_factory=list) warnings: list[str] = field(default_factory=list) @@ -78,11 +83,14 @@ def warn(self, message: str) -> None: def add_weight_binding(self, binding: "WeightBinding") -> None: add_component_weight_binding(self, binding) - def add_runtime_input(self, tensor: GraphTensor) -> None: - add_component_runtime_input(self, tensor) + def add_cache_state_binding(self, binding: RPModels.CacheStateBinding) -> None: + add_component_cache_state_binding(self, binding) - def add_output(self, tensor: GraphTensor) -> None: - add_component_output(self, tensor) + def add_runtime_input(self, tensor: GraphTensor, logical_name: str | None = None) -> None: + add_component_runtime_input(self, tensor, logical_name) + + def add_output(self, tensor: GraphTensor, logical_name: str | None = None) -> None: + add_component_output(self, tensor, logical_name) @dataclass(slots=True) @@ -175,8 +183,12 @@ class ComponentGraphManifest: component: str graph_path: str weight_bindings: tuple[WeightBinding, ...] = () + cache_state_node_ids: tuple[RPModels.CacheStateBinding, ...] = () runtime_input_node_ids: tuple[int, ...] = () + logical_inputs: tuple[str, ...] = () output_node_ids: tuple[int, ...] = () + logical_outputs: tuple[str, ...] = () + metadata: dict[str, str] | None = None unsupported_nodes: tuple[str, ...] = () warnings: tuple[str, ...] = () @@ -269,18 +281,42 @@ def add_component_weight_binding(component: ComponentGraph, binding: WeightBindi component.weight_bindings.append(binding) -def add_component_runtime_input(component: ComponentGraph, tensor: GraphTensor) -> None: +def add_component_cache_state_binding(component: ComponentGraph, binding: RPModels.CacheStateBinding) -> None: + RPModels.merge_cache_state_binding(component.cache_state_bindings, binding) + + +def add_component_runtime_input(component: ComponentGraph, tensor: GraphTensor, logical_name: str | None = None) -> None: node_id = tensor_node_id(tensor) - if node_id is not None and node_id not in component.runtime_input_ids: - component.runtime_input_ids.append(node_id) + if node_id is None: + return + + name = logical_name or f"input_{len(component.runtime_input_ids)}" + if node_id in component.runtime_input_ids: + index = component.runtime_input_ids.index(node_id) + component.logical_inputs[index] = name + return -def add_component_output(component: ComponentGraph, tensor: GraphTensor) -> None: + component.runtime_input_ids.append(node_id) + component.logical_inputs.append(name) + + +def add_component_output(component: ComponentGraph, tensor: GraphTensor, logical_name: str | None = None) -> None: node_id = tensor_node_id(tensor) - if node_id is not None and node_id not in component.output_node_ids: - component.output_node_ids.append(node_id) + if node_id is None: + return + + name = logical_name or f"output_{len(component.output_node_ids)}" + + if node_id in component.output_node_ids: + index = component.output_node_ids.index(node_id) + component.logical_outputs[index] = name + return + + component.output_node_ids.append(node_id) + component.logical_outputs.append(name) def tensor_node_id(tensor: GraphTensor) -> int | None: @@ -483,8 +519,12 @@ def component_manifest_from_component(cls: type[ComponentGraphManifest], compone component=component.name, graph_path=component.output_path.name, weight_bindings=tuple(component.weight_bindings), + cache_state_node_ids=tuple(component.cache_state_bindings), runtime_input_node_ids=tuple(component.runtime_input_ids), + logical_inputs=tuple(component.logical_inputs), output_node_ids=tuple(component.output_node_ids), + logical_outputs=tuple(component.logical_outputs), + metadata=dict(component.metadata), unsupported_nodes=tuple(component.unsupported_nodes), warnings=tuple(component.warnings), ) @@ -493,6 +533,7 @@ def component_manifest_from_component(cls: type[ComponentGraphManifest], compone def component_manifest_to_dict(manifest: ComponentGraphManifest) -> dict[str, Any]: return { "component": manifest.component, + "graph": manifest.graph_path, "graph_path": manifest.graph_path, "weight_bindings": [ { @@ -507,8 +548,22 @@ def component_manifest_to_dict(manifest: ComponentGraphManifest) -> dict[str, An } for binding in manifest.weight_bindings ], + "bound_constant_bindings": [ + { + "node_id": binding.node_id, + "path": binding.path, + } + for binding in manifest.weight_bindings + ], + "cache_state_node_ids": [ + RPModels.cache_state_binding_to_dict(binding) + for binding in manifest.cache_state_node_ids + ], "runtime_input_node_ids": list(manifest.runtime_input_node_ids), + "logical_inputs": list(manifest.logical_inputs), "output_node_ids": list(manifest.output_node_ids), + "logical_outputs": list(manifest.logical_outputs), + "metadata": dict(manifest.metadata or {}), "unsupported_nodes": list(manifest.unsupported_nodes), "warnings": list(manifest.warnings), } diff --git a/python/cactus/transpiler/ModelProfiles/combinations.py b/python/cactus/transpiler/ModelProfiles/combinations.py index fa53ea279..ac09b5b05 100644 --- a/python/cactus/transpiler/ModelProfiles/combinations.py +++ b/python/cactus/transpiler/ModelProfiles/combinations.py @@ -1,12 +1,12 @@ from .models import Combinations # The input and output strings are component names from components.py. -TEXT_EMBED_TO_DECODER = Combinations(input=("text_embed"), output="text_decoder") -DECODER_TO_LM_HEAD = Combinations(input=("text_decoder"), output="lm_head") +TEXT_EMBED_TO_DECODER = Combinations(input=("text_embed",), output="text_decoder") +DECODER_TO_LM_HEAD = Combinations(input=("text_decoder",), output="lm_head") TEXT_VISION_TO_TOKEN_MERGE = Combinations(input=("text_embed", "vision_encoder"), output="token_merge") TEXT_AUDIO_TO_TOKEN_MERGE = Combinations(input=("text_embed", "audio_encoder"), output="token_merge") TEXT_VISION_AUDIO_TO_TOKEN_MERGE = Combinations(input=("text_embed", "vision_encoder", "audio_encoder"), output="token_merge") -TOKEN_MERGE_TO_DECODER = Combinations(input=("token_merge"), output="text_decoder") +TOKEN_MERGE_TO_DECODER = Combinations(input=("token_merge",), output="text_decoder") AUDIO_TEXT_TO_DECODER = Combinations(input=("audio_encoder", "text_embed"), output="text_decoder") -AUDIO_ENCODER_TO_ASR_DECODER = Combinations(input=("audio_encoder"), output="asr_decoder") +AUDIO_ENCODER_TO_ASR_DECODER = Combinations(input=("audio_encoder",), output="asr_decoder") AUDIO_ASR_DECODER_TO_ASR_HEAD = Combinations(input=("audio_encoder", "asr_decoder"), output="asr_head") diff --git a/python/cactus/transpiler/RuntimePlan/models.py b/python/cactus/transpiler/RuntimePlan/models.py new file mode 100644 index 000000000..bfb09c7d2 --- /dev/null +++ b/python/cactus/transpiler/RuntimePlan/models.py @@ -0,0 +1,420 @@ +from __future__ import annotations + +import json +import os +from dataclasses import dataclass +from pathlib import Path +from typing import Any + + +@dataclass(slots=True, frozen=True) +class RuntimeTensorBinding: + node_id: int + logical_name: str + + +@dataclass(slots=True, frozen=True) +class ConstantBinding: + node_id: int + path: str + + +@dataclass(slots=True) +class CacheStateBinding: + layer_key: str + key_node_id: int = -1 + value_node_id: int = -1 + cache_kind: str = "kv" + tensor_indices: tuple[int, ...] = () + + +@dataclass(slots=True, frozen=True) +class RuntimeRouteEdge: + inputs: tuple[str, ...] + output: str + + +@dataclass(slots=True, frozen=True) +class RuntimeRoute: + name: str + edges: tuple[RuntimeRouteEdge, ...] = () + + +@dataclass(slots=True, frozen=True) +class RuntimeComponent: + component: str + graph: str + runtime_input_node_ids: tuple[int, ...] = () + logical_inputs: tuple[str, ...] = () + output_node_ids: tuple[int, ...] = () + logical_outputs: tuple[str, ...] = () + bound_constant_bindings: tuple[ConstantBinding, ...] = () + cache_state_node_ids: tuple[CacheStateBinding, ...] = () + metadata: dict[str, str] | None = None + unsupported_nodes: tuple[str, ...] = () + warnings: tuple[str, ...] = () + + def to_engine_dict(self) -> dict[str, Any]: + return runtime_component_to_engine_dict(self) + + def to_plan_dict(self) -> dict[str, Any]: + return runtime_component_to_plan_dict(self) + + +@dataclass(slots=True, frozen=True) +class RuntimePlan: + family: str = "" + components: tuple[RuntimeComponent, ...] = () + routes: tuple[RuntimeRoute, ...] = () + metadata: dict[str, str] | None = None + + def to_engine_manifest(self) -> dict[str, Any]: + return runtime_plan_to_engine_manifest(self) + + def to_plan_dict(self) -> dict[str, Any]: + return runtime_plan_to_dict(self) + + def write(self, bundle_dir: str | Path) -> tuple[Path, Path]: + return write_runtime_plan(self, bundle_dir) + + +################################################# Runtime Plan Utils!!!!!!! ################################################# + + +def cache_state_binding_from_annotation(annotation: Any, node_id: int) -> CacheStateBinding: + cache_kind = str(getattr(annotation, "kind", "kv")) + role = str(getattr(annotation, "role", "state")) + tensor_index = getattr(annotation, "tensor_index", None) + tensor_indices = (int(tensor_index),) if tensor_index is not None else () + + key_node_id = -1 + value_node_id = -1 + + if cache_kind == "conv" or role == "state": + key_node_id = node_id + value_node_id = node_id + elif role == "key": + key_node_id = node_id + elif role == "value": + value_node_id = node_id + else: + key_node_id = node_id + value_node_id = node_id + + return CacheStateBinding( + layer_key=cache_layer_key(annotation), + key_node_id=key_node_id, + value_node_id=value_node_id, + cache_kind=cache_kind, + tensor_indices=tensor_indices, + ) + + +def cache_layer_key(annotation: Any) -> str: + cache_kind = str(getattr(annotation, "kind", "cache")) + layer_index = getattr(annotation, "layer_index", None) + tensor_index = getattr(annotation, "tensor_index", None) + + if layer_index is not None: + return f"{cache_kind}:{int(layer_index)}" + + if tensor_index is not None: + return f"{cache_kind}:tensor:{int(tensor_index)}" + + return f"{cache_kind}:unknown" + + +def add_runtime_tensor_binding(bindings: list[RuntimeTensorBinding], binding: RuntimeTensorBinding) -> None: + for index, existing in enumerate(bindings): + if existing.node_id != binding.node_id: + continue + + if existing.logical_name == binding.logical_name: + return + + bindings[index] = binding + return + + bindings.append(binding) + + +def merge_cache_state_binding(bindings: list[CacheStateBinding], binding: CacheStateBinding) -> None: + for existing in bindings: + if existing.layer_key != binding.layer_key: + continue + + if binding.key_node_id >= 0: + existing.key_node_id = binding.key_node_id + + if binding.value_node_id >= 0: + existing.value_node_id = binding.value_node_id + + existing.tensor_indices = unique_ints((*existing.tensor_indices, *binding.tensor_indices)) + return + + bindings.append(binding) + + +def cache_state_binding_from_dict(data: dict[str, Any]) -> CacheStateBinding: + return CacheStateBinding( + layer_key=str(data.get("layer_key", "")), + key_node_id=int(data.get("key", data.get("key_node_id", -1))), + value_node_id=int(data.get("value", data.get("value_node_id", -1))), + cache_kind=str(data.get("cache_kind", "kv")), + tensor_indices=tuple(int(value) for value in data.get("tensor_indices", ())), + ) + + +def cache_state_binding_to_dict(binding: CacheStateBinding) -> dict[str, Any]: + return { + "layer_key": binding.layer_key, + "key": binding.key_node_id, + "value": binding.value_node_id, + "cache_kind": binding.cache_kind, + "tensor_indices": list(binding.tensor_indices), + } + + +def constant_binding_from_dict(data: dict[str, Any]) -> ConstantBinding: + return ConstantBinding( + node_id=int(data["node_id"]), + path=str(data["path"]), + ) + + +def constant_binding_to_dict(binding: ConstantBinding) -> dict[str, Any]: + return { + "node_id": binding.node_id, + "path": binding.path, + } + + +def runtime_component_from_generator_manifest( + manifest_path: str | Path, + *, + bundle_dir: str | Path | None = None, + component_name: str | None = None, + metadata: dict[str, str] | None = None, +) -> RuntimeComponent: + path = Path(manifest_path) + data = json.loads(path.read_text(encoding="utf-8")) + graph_path = str(data.get("graph") or data.get("graph_path") or "") + + return RuntimeComponent( + component=component_name or str(data["component"]), + graph=runtime_graph_path(path, graph_path, Path(bundle_dir) if bundle_dir is not None else None), + runtime_input_node_ids=tuple(int(value) for value in data.get("runtime_input_node_ids", ())), + logical_inputs=tuple(str(value) for value in data.get("logical_inputs", ())), + output_node_ids=tuple(int(value) for value in data.get("output_node_ids", ())), + logical_outputs=tuple(str(value) for value in data.get("logical_outputs", ())), + bound_constant_bindings=constant_bindings_from_generator_manifest(data), + cache_state_node_ids=tuple( + cache_state_binding_from_dict(value) + for value in data.get("cache_state_node_ids", ()) + if isinstance(value, dict) + ), + metadata=string_dict({**string_dict(data.get("metadata")), **string_dict(metadata)}), + unsupported_nodes=tuple(str(value) for value in data.get("unsupported_nodes", ())), + warnings=tuple(str(value) for value in data.get("warnings", ())), + ) + + +def runtime_plan_from_generator_manifests( + component_manifest_paths: dict[str, str | Path], + *, + bundle_dir: str | Path | None = None, + model_profile: Any | None = None, + component_name_map: dict[str, str] | None = None, + component_metadata: dict[str, dict[str, str]] | None = None, + metadata: dict[str, str] | None = None, +) -> RuntimePlan: + components = tuple( + runtime_component_from_generator_manifest( + manifest_path, + bundle_dir=bundle_dir, + component_name=(component_name_map or {}).get(name, name), + metadata=(component_metadata or {}).get(name), + ) + for name, manifest_path in component_manifest_paths.items() + ) + + return RuntimePlan( + family=str(getattr(model_profile, "model_profiles", "") or ""), + components=components, + routes=runtime_routes_from_model_profile(model_profile), + metadata=string_dict(metadata), + ) + + +def runtime_plan_from_generation_result( + generation_result: Any, + *, + bundle_dir: str | Path | None = None, + model_profile: Any | None = None, + component_name_map: dict[str, str] | None = None, + component_metadata: dict[str, dict[str, str]] | None = None, + metadata: dict[str, str] | None = None, +) -> RuntimePlan: + return runtime_plan_from_generator_manifests( + dict(getattr(generation_result, "component_manifest_paths", {})), + bundle_dir=bundle_dir, + model_profile=model_profile, + component_name_map=component_name_map, + component_metadata=component_metadata, + metadata=metadata, + ) + + +def runtime_routes_from_model_profile(model_profile: Any | None) -> tuple[RuntimeRoute, ...]: + if model_profile is None: + return () + + routes = getattr(model_profile, "inference_type", {}) or {} + return tuple(runtime_route_from_profile_route(route) for route in routes.values()) + + +def runtime_route_from_profile_route(route: Any) -> RuntimeRoute: + return RuntimeRoute( + name=str(getattr(route, "name", "")), + edges=tuple( + RuntimeRouteEdge( + inputs=tuple(str(value) for value in getattr(edge, "input", ())), + output=str(getattr(edge, "output", "")), + ) + for edge in getattr(route, "route", ()) + ), + ) + + +def runtime_plan_to_engine_manifest(plan: RuntimePlan) -> dict[str, Any]: + manifest: dict[str, Any] = { + "components": [component.to_engine_dict() for component in plan.components], + } + + if plan.family: + manifest["family"] = plan.family + + for key, value in string_dict(plan.metadata).items(): + manifest[key] = value + + return manifest + + +def runtime_plan_to_dict(plan: RuntimePlan) -> dict[str, Any]: + return { + "family": plan.family, + "components": [component.to_plan_dict() for component in plan.components], + "routes": [runtime_route_to_dict(route) for route in plan.routes], + "metadata": string_dict(plan.metadata), + } + + +def runtime_component_to_engine_dict(component: RuntimeComponent) -> dict[str, Any]: + return { + "component": component.component, + "graph": component.graph, + "runtime_input_node_ids": list(component.runtime_input_node_ids), + "logical_inputs": list(component.logical_inputs), + "output_node_ids": list(component.output_node_ids), + "logical_outputs": list(component.logical_outputs), + "metadata": string_dict(component.metadata), + "bound_constant_bindings": [ + constant_binding_to_dict(binding) + for binding in component.bound_constant_bindings + ], + "cache_state_node_ids": [ + cache_state_binding_to_dict(binding) + for binding in component.cache_state_node_ids + ], + } + + +def runtime_component_to_plan_dict(component: RuntimeComponent) -> dict[str, Any]: + data = runtime_component_to_engine_dict(component) + data["unsupported_nodes"] = list(component.unsupported_nodes) + data["warnings"] = list(component.warnings) + return data + + +def runtime_route_to_dict(route: RuntimeRoute) -> dict[str, Any]: + return { + "name": route.name, + "edges": [ + { + "inputs": list(edge.inputs), + "output": edge.output, + } + for edge in route.edges + ], + } + + +def write_runtime_plan(plan: RuntimePlan, bundle_dir: str | Path) -> tuple[Path, Path]: + bundle_path = Path(bundle_dir) + components_dir = bundle_path / "components" + components_dir.mkdir(parents=True, exist_ok=True) + + engine_manifest_path = components_dir / "manifest.json" + runtime_plan_path = bundle_path / "runtime_plan.json" + + write_json(engine_manifest_path, plan.to_engine_manifest()) + write_json(runtime_plan_path, plan.to_plan_dict()) + return engine_manifest_path, runtime_plan_path + + +def write_json(path: Path, data: dict[str, Any]) -> None: + path.parent.mkdir(parents=True, exist_ok=True) + path.write_text(json.dumps(data, indent=4), encoding="utf-8") + + +def runtime_graph_path(manifest_path: Path, graph_path: str, bundle_dir: Path | None) -> str: + if not graph_path or bundle_dir is None: + return graph_path + + path = Path(graph_path) + + if not path.is_absolute(): + path = manifest_path.parent / path + + return os.path.relpath(path, bundle_dir) + + +def constant_bindings_from_generator_manifest(data: dict[str, Any]) -> tuple[ConstantBinding, ...]: + if "bound_constant_bindings" in data: + return tuple( + constant_binding_from_dict(value) + for value in data.get("bound_constant_bindings", ()) + if isinstance(value, dict) + ) + + return tuple( + ConstantBinding( + node_id=int(value["node_id"]), + path=str(value["path"]), + ) + for value in data.get("weight_bindings", ()) + if isinstance(value, dict) and "node_id" in value and "path" in value + ) + + +def string_dict(data: Any) -> dict[str, str]: + if not isinstance(data, dict): + return {} + + return { + str(key): str(value) + for key, value in data.items() + if value is not None + } + + +def unique_ints(values: tuple[int, ...]) -> tuple[int, ...]: + unique: list[int] = [] + + for value in values: + if value in unique: + continue + + unique.append(value) + + return tuple(unique) From 3b0b9cfa8c4604c960359c3dc8fc8ae8b2bb5fbf Mon Sep 17 00:00:00 2001 From: Parkirat Date: Wed, 29 Jul 2026 15:28:25 -0400 Subject: [PATCH 66/84] New run operations with runtime plans --- python/cactus/bindings/cactus.py | 197 ++++++- python/cactus/cli/__init__.py | 4 + python/cactus/cli/common.py | 7 +- python/cactus/cli/model.py | 48 +- python/cactus/cli/serve.py | 2 +- .../convert/cactus_adapters/tensor_io.py | 2 +- python/cactus/convert/cli.py | 21 + .../cactus/convert/model_adapters/naming.py | 10 + python/cactus/py.typed | 0 python/cactus/transpiler/Fusions/fusions.py | 7 +- .../transpiler/Generator/component_split.py | 485 +++++++++++++++++ .../cactus/transpiler/Generator/constants.py | 65 ++- .../cactus/transpiler/Generator/generate.py | 515 +++++++++++++++++- python/cactus/transpiler/Generator/models.py | 12 + python/cactus/transpiler/IR/constants.py | 3 + python/cactus/transpiler/main.py | 58 +- 16 files changed, 1368 insertions(+), 68 deletions(-) delete mode 100644 python/cactus/py.typed create mode 100644 python/cactus/transpiler/Generator/component_split.py diff --git a/python/cactus/bindings/cactus.py b/python/cactus/bindings/cactus.py index 327afdb5b..9a4a5418d 100644 --- a/python/cactus/bindings/cactus.py +++ b/python/cactus/bindings/cactus.py @@ -138,6 +138,25 @@ class cactus_tensor_info_t(ctypes.Structure): [cactus_graph_t, cactus_node_t, cactus_node_t, ctypes.POINTER(cactus_node_t)], ctypes.c_int, ) +for _binary_name in ( + "cactus_graph_equal", + "cactus_graph_less", + "cactus_graph_less_equal", + "cactus_graph_greater", + "cactus_graph_greater_equal", + "cactus_graph_bitwise_and", + "cactus_graph_bitwise_or", +): + _bind_optional( + _binary_name, + [cactus_graph_t, cactus_node_t, cactus_node_t, ctypes.POINTER(cactus_node_t)], + ctypes.c_int, + ) +_bind_optional( + "cactus_graph_where", + [cactus_graph_t, cactus_node_t, cactus_node_t, cactus_node_t, ctypes.POINTER(cactus_node_t)], + ctypes.c_int, +) _lib.cactus_graph_precision_cast.argtypes = [ cactus_graph_t, cactus_node_t, ctypes.c_int32, ctypes.POINTER(cactus_node_t) @@ -175,6 +194,27 @@ class cactus_tensor_info_t(ctypes.Structure): [cactus_graph_t, cactus_node_t, ctypes.c_float, ctypes.POINTER(cactus_node_t)], ctypes.c_int, ) +for _scalar_name in ( + "cactus_graph_scalar_equal", + "cactus_graph_scalar_less", + "cactus_graph_scalar_less_equal", + "cactus_graph_scalar_greater", + "cactus_graph_scalar_greater_equal", +): + _bind_optional( + _scalar_name, + [cactus_graph_t, cactus_node_t, ctypes.c_float, ctypes.POINTER(cactus_node_t)], + ctypes.c_int, + ) +for _unary_bool_name in ( + "cactus_graph_logical_not", + "cactus_graph_bitwise_not", +): + _bind_optional( + _unary_bool_name, + [cactus_graph_t, cactus_node_t, ctypes.POINTER(cactus_node_t)], + ctypes.c_int, + ) _lib.cactus_graph_scalar_exp.argtypes = [ cactus_graph_t, cactus_node_t, ctypes.POINTER(cactus_node_t) ] @@ -247,10 +287,25 @@ class cactus_tensor_info_t(ctypes.Structure): cactus_graph_t, cactus_node_t, ctypes.c_int32, ctypes.c_size_t, ctypes.c_size_t, ctypes.POINTER(cactus_node_t) ] _lib.cactus_graph_slice.restype = ctypes.c_int +_bind_optional( + "cactus_graph_strided_slice", + [cactus_graph_t, cactus_node_t, ctypes.c_int32, ctypes.c_size_t, ctypes.c_size_t, ctypes.c_size_t, ctypes.POINTER(cactus_node_t)], + ctypes.c_int, +) _lib.cactus_graph_index.argtypes = [ cactus_graph_t, cactus_node_t, ctypes.c_size_t, ctypes.c_int32, ctypes.POINTER(cactus_node_t) ] _lib.cactus_graph_index.restype = ctypes.c_int +_bind_optional( + "cactus_graph_unfold", + [cactus_graph_t, cactus_node_t, ctypes.c_int32, ctypes.c_size_t, ctypes.c_size_t, ctypes.POINTER(cactus_node_t)], + ctypes.c_int, +) +_bind_optional( + "cactus_graph_pad", + [cactus_graph_t, cactus_node_t, ctypes.POINTER(ctypes.c_size_t), ctypes.c_size_t, ctypes.c_float, ctypes.POINTER(cactus_node_t)], + ctypes.c_int, +) _lib.cactus_graph_concat.argtypes = [ cactus_graph_t, cactus_node_t, cactus_node_t, ctypes.c_int32, ctypes.POINTER(cactus_node_t) @@ -270,6 +325,11 @@ class cactus_tensor_info_t(ctypes.Structure): cactus_graph_t, cactus_node_t, cactus_node_t, ctypes.POINTER(cactus_node_t) ] _lib.cactus_graph_gather.restype = ctypes.c_int +_bind_optional( + "cactus_graph_gather_dim", + [cactus_graph_t, cactus_node_t, cactus_node_t, ctypes.c_int32, ctypes.POINTER(cactus_node_t)], + ctypes.c_int, +) _lib.cactus_graph_embedding_from_tensor.argtypes = [ cactus_graph_t, cactus_node_t, cactus_node_t, ctypes.POINTER(cactus_node_t) ] @@ -531,6 +591,11 @@ class cactus_tensor_info_t(ctypes.Structure): cactus_graph_t, cactus_node_t, cactus_node_t, ctypes.c_bool, cactus_node_t, ctypes.c_size_t, ctypes.POINTER(cactus_node_t) ] _lib.cactus_graph_conv1d.restype = ctypes.c_int +_bind_optional( + "cactus_graph_conv1d_depthwise", + [cactus_graph_t, cactus_node_t, cactus_node_t, ctypes.c_bool, cactus_node_t, ctypes.c_size_t, ctypes.POINTER(cactus_node_t)], + ctypes.c_int, +) _lib.cactus_graph_conv1d_same_depthwise_k9.argtypes = [ cactus_graph_t, cactus_node_t, cactus_node_t, ctypes.c_bool, cactus_node_t, ctypes.POINTER(cactus_node_t) ] @@ -1627,6 +1692,43 @@ def divide(self, a, b): def not_equal(self, a, b): return self._binary("cactus_graph_not_equal", a, b) + def equal(self, a, b): + return self._binary("cactus_graph_equal", a, b) + + def less(self, a, b): + return self._binary("cactus_graph_less", a, b) + + def less_equal(self, a, b): + return self._binary("cactus_graph_less_equal", a, b) + + def greater(self, a, b): + return self._binary("cactus_graph_greater", a, b) + + def greater_equal(self, a, b): + return self._binary("cactus_graph_greater_equal", a, b) + + def bitwise_and(self, a, b): + return self._binary("cactus_graph_bitwise_and", a, b) + + def bitwise_or(self, a, b): + return self._binary("cactus_graph_bitwise_or", a, b) + + def where(self, condition, true_value, false_value): + condition = self._ensure_tensor(condition) + true_value = self._ensure_tensor(true_value) + false_value = self._ensure_tensor(false_value) + out = cactus_node_t() + rc = _lib.cactus_graph_where( + self.h, + cactus_node_t(condition.id), + cactus_node_t(true_value.id), + cactus_node_t(false_value.id), + ctypes.byref(out), + ) + if rc != 0: + raise RuntimeError(_err("graph_where failed")) + return self._tensor_from_node(out.value) + def abs(self, x): x = self._ensure_tensor(x) out = cactus_node_t() @@ -1689,6 +1791,27 @@ def scalar_floor_divide(self, x, value): def scalar_not_equal(self, x, value): return self._scalar("cactus_graph_scalar_not_equal", x, value) + def scalar_equal(self, x, value): + return self._scalar("cactus_graph_scalar_equal", x, value) + + def scalar_less(self, x, value): + return self._scalar("cactus_graph_scalar_less", x, value) + + def scalar_less_equal(self, x, value): + return self._scalar("cactus_graph_scalar_less_equal", x, value) + + def scalar_greater(self, x, value): + return self._scalar("cactus_graph_scalar_greater", x, value) + + def scalar_greater_equal(self, x, value): + return self._scalar("cactus_graph_scalar_greater_equal", x, value) + + def logical_not(self, x): + return self._scalar("cactus_graph_logical_not", x) + + def bitwise_not(self, x): + return self._scalar("cactus_graph_bitwise_not", x) + def scalar_exp(self, x): return self._scalar("cactus_graph_scalar_exp", x) @@ -1796,6 +1919,22 @@ def slice(self, x, axis, start, length=0): raise RuntimeError(_err("graph_slice failed")) return self._tensor_from_node(out.value) + def strided_slice(self, x, axis, start, length, step): + x = self._ensure_tensor(x) + out = cactus_node_t() + rc = _lib.cactus_graph_strided_slice( + self.h, + cactus_node_t(x.id), + ctypes.c_int32(int(axis)), + ctypes.c_size_t(int(start)), + ctypes.c_size_t(int(length)), + ctypes.c_size_t(int(step)), + ctypes.byref(out), + ) + if rc != 0: + raise RuntimeError(_err("graph_strided_slice failed")) + return self._tensor_from_node(out.value) + def index(self, x, index_value, axis=0): x = self._ensure_tensor(x) out = cactus_node_t() @@ -1810,6 +1949,38 @@ def index(self, x, index_value, axis=0): raise RuntimeError(_err("graph_index failed")) return self._tensor_from_node(out.value) + def unfold(self, x, dimension, size, step): + x = self._ensure_tensor(x) + out = cactus_node_t() + rc = _lib.cactus_graph_unfold( + self.h, + cactus_node_t(x.id), + ctypes.c_int32(int(dimension)), + ctypes.c_size_t(int(size)), + ctypes.c_size_t(int(step)), + ctypes.byref(out), + ) + if rc != 0: + raise RuntimeError(_err("graph_unfold failed")) + return self._tensor_from_node(out.value) + + def pad(self, x, pads, value=0.0): + x = self._ensure_tensor(x) + pads = tuple(int(pad) for pad in pads) + arr = (ctypes.c_size_t * len(pads))(*pads) + out = cactus_node_t() + rc = _lib.cactus_graph_pad( + self.h, + cactus_node_t(x.id), + arr, + ctypes.c_size_t(len(pads)), + ctypes.c_float(float(value)), + ctypes.byref(out), + ) + if rc != 0: + raise RuntimeError(_err("graph_pad failed")) + return self._tensor_from_node(out.value) + def transpose(self, x, backend=CPU): x = self._ensure_tensor(x) out = cactus_node_t() @@ -1856,16 +2027,25 @@ def matmul(self, a, b, pretransposed_rhs=False, backend=CPU, output_dtype=None): raise RuntimeError(_err("graph_matmul failed")) return self._tensor_from_node(out.value) - def gather(self, tensor, indices): + def gather(self, tensor, indices, axis=0): tensor = self._ensure_tensor(tensor) indices = self._ensure_tensor(indices) out = cactus_node_t() - rc = _lib.cactus_graph_gather( - self.h, - cactus_node_t(tensor.id), - cactus_node_t(indices.id), - ctypes.byref(out), - ) + if hasattr(_lib, "cactus_graph_gather_dim"): + rc = _lib.cactus_graph_gather_dim( + self.h, + cactus_node_t(tensor.id), + cactus_node_t(indices.id), + ctypes.c_int32(int(axis)), + ctypes.byref(out), + ) + else: + rc = _lib.cactus_graph_gather( + self.h, + cactus_node_t(tensor.id), + cactus_node_t(indices.id), + ctypes.byref(out), + ) if rc != 0: raise RuntimeError(_err("graph_gather failed")) return self._tensor_from_node(out.value) @@ -2445,6 +2625,9 @@ def conv1d_k7s3(self, x, weight, bias): def conv1d(self, x, weight, bias=None, stride=1): return self._conv_with_optional_bias("cactus_graph_conv1d", x, weight, bias, ctypes.c_size_t(int(stride))) + def conv1d_depthwise(self, x, weight, bias=None, stride=1): + return self._conv_with_optional_bias("cactus_graph_conv1d_depthwise", x, weight, bias, ctypes.c_size_t(int(stride))) + def conv1d_same_depthwise_k9(self, x, weight, bias=None): return self._conv_with_optional_bias("cactus_graph_conv1d_same_depthwise_k9", x, weight, bias) diff --git a/python/cactus/cli/__init__.py b/python/cactus/cli/__init__.py index eea9a875d..39cb18b42 100644 --- a/python/cactus/cli/__init__.py +++ b/python/cactus/cli/__init__.py @@ -1,5 +1,6 @@ import sys import argparse +from pathlib import Path from .. import __version__ from .common import ( @@ -81,6 +82,9 @@ def _unit_float(value): def _hf_id_or_path(value): v = (value or "").strip() + if Path(v).expanduser().exists(): + return v + if "/" not in v: raise argparse.ArgumentTypeError( f"invalid model {value!r}. Use a HuggingFace id like 'openai/whisper-base' " diff --git a/python/cactus/cli/common.py b/python/cactus/cli/common.py index fac99ed54..3a6fc3e27 100644 --- a/python/cactus/cli/common.py +++ b/python/cactus/cli/common.py @@ -84,9 +84,12 @@ def apply_runtime_env(args) -> None: def is_valid_bundle(path) -> bool: - """A runnable v2 bundle has both config.txt and the components manifest.""" + """A runnable v2 bundle has config.txt plus an engine manifest or runtime plan.""" path = Path(path) - return (path / "config.txt").exists() and (path / "components" / "manifest.json").exists() + return (path / "config.txt").exists() and ( + (path / "components" / "manifest.json").exists() + or (path / "runtime_plan.json").exists() + ) def launch_binary(name, *args) -> int: diff --git a/python/cactus/cli/model.py b/python/cactus/cli/model.py index db9e7bdfc..71d330764 100644 --- a/python/cactus/cli/model.py +++ b/python/cactus/cli/model.py @@ -2,6 +2,7 @@ from __future__ import annotations import os +import json import shutil from pathlib import Path @@ -49,16 +50,61 @@ def ensure_weights(model_id, *, bits=4, platform=None, token=None, reconvert=Fal def _has_runnable_bundle(path): path = Path(path) - return (path / "components" / "manifest.json").exists() + return (path / "components" / "manifest.json").exists() or (path / "runtime_plan.json").exists() def resolve_bundle_dir(model_id): path = Path(model_id).expanduser() + if path.is_file() and path.name == "runtime_plan.json": + path = path.parent + if path.is_dir() and _has_runnable_bundle(path): + materialize_engine_manifest_from_runtime_plan(path) return path return None +def materialize_engine_manifest_from_runtime_plan(bundle_dir): + bundle_path = Path(bundle_dir) + runtime_plan_path = bundle_path / "runtime_plan.json" + engine_manifest_path = bundle_path / "components" / "manifest.json" + + if not runtime_plan_path.exists(): + return engine_manifest_path if engine_manifest_path.exists() else None + + if engine_manifest_path.exists() and engine_manifest_path.stat().st_mtime >= runtime_plan_path.stat().st_mtime: + return engine_manifest_path + + plan = json.loads(runtime_plan_path.read_text(encoding="utf-8")) + manifest = engine_manifest_from_runtime_plan(plan) + engine_manifest_path.parent.mkdir(parents=True, exist_ok=True) + engine_manifest_path.write_text(json.dumps(manifest, indent=4), encoding="utf-8") + return engine_manifest_path + + +def engine_manifest_from_runtime_plan(plan): + components = plan.get("components") + + if not isinstance(components, list): + raise ValueError("runtime_plan.json must contain a components list") + + manifest = {"components": components} + family = plan.get("family") + + if family: + manifest["family"] = str(family) + + metadata = plan.get("metadata") + + if isinstance(metadata, dict): + for key, value in metadata.items(): + if value is None: + continue + manifest[str(key)] = str(value) + + return manifest + + def _local_build_unavailable(model_id) -> RuntimeError: return RuntimeError( f"No runnable prebuilt bundle is available for {model_id!r}. " diff --git a/python/cactus/cli/serve.py b/python/cactus/cli/serve.py index 005203c79..edce062ca 100644 --- a/python/cactus/cli/serve.py +++ b/python/cactus/cli/serve.py @@ -31,7 +31,7 @@ def cmd_serve(args): model_path, model_name = built, built.name if model_path is not None and not is_valid_bundle(model_path): print_color(RED, f"Error: not a valid v2 Cactus bundle: {model_path}") - print("Expected config.txt and components/manifest.json.") + print("Expected config.txt and either components/manifest.json or runtime_plan.json.") return 1 try: diff --git a/python/cactus/convert/cactus_adapters/tensor_io.py b/python/cactus/convert/cactus_adapters/tensor_io.py index ba978c6aa..e3113bb2a 100644 --- a/python/cactus/convert/cactus_adapters/tensor_io.py +++ b/python/cactus/convert/cactus_adapters/tensor_io.py @@ -311,7 +311,7 @@ def save_tensor_with_header(tensor, output_path, precision='INT8', transpose=Fal model_type == 'gemma4' and len(shape) == 3 and 'vision' not in output_path.name - and not output_path.name.startswith('audio_') + and not output_path.name.startswith(('audio_', 'model_audio_')) ): data = data.transpose(0, 2, 1) shape = [shape[0], shape[2], shape[1]] diff --git a/python/cactus/convert/cli.py b/python/cactus/convert/cli.py index 20f8dc94e..5218398b6 100644 --- a/python/cactus/convert/cli.py +++ b/python/cactus/convert/cli.py @@ -195,7 +195,27 @@ def _augment_state_dict_for_family(state_dict: dict[str, Any], family: str) -> d return adapter_for_family(family).normalize_state_dict(state_dict).state_dict +def _augment_state_dict_with_model_buffers(state_dict: Mapping[str, Any], model: Any) -> dict[str, Any]: + augmented = dict(state_dict) + if model is None or not hasattr(model, "named_buffers"): + return augmented + + for name, tensor in model.named_buffers(): + if name in augmented or not hasattr(tensor, "numel"): + continue + if int(tensor.numel()) > 1_000_000: + continue + augmented[name] = tensor.detach().cpu() if hasattr(tensor, "detach") else tensor + + return augmented + + def _save_fallback_tensor(tensor, out_path: Path, precision: str, family: str) -> None: + if len(_tensor_shape(tensor)) == 0: + if torch is not None and isinstance(tensor, torch.Tensor): + tensor = tensor.reshape(1) + else: + tensor = np.asarray(tensor).reshape(1) if len(_tensor_shape(tensor)) > 4: if torch is not None and isinstance(tensor, torch.Tensor): tensor = tensor.reshape(tensor.shape[0], -1) @@ -406,6 +426,7 @@ def convert(args: argparse.Namespace) -> None: if model is None: raise RuntimeError(f"could not load model object or checkpoint tensors from {args.model}") checkpoint_state = model.state_dict() + checkpoint_state = _augment_state_dict_with_model_buffers(checkpoint_state, model) normalized_state = adapter.normalize_state_dict(checkpoint_state) state_dict = normalized_state.state_dict model_config = adapter.runtime_config(cfg) diff --git a/python/cactus/convert/model_adapters/naming.py b/python/cactus/convert/model_adapters/naming.py index b44a4a072..78d964be5 100644 --- a/python/cactus/convert/model_adapters/naming.py +++ b/python/cactus/convert/model_adapters/naming.py @@ -26,6 +26,13 @@ class NameMatch: "embed_vision_proj", "embed_audio_proj", } +GEMMA4_RUNTIME_BUFFER_SUFFIXES = ( + "embed_scale", + "inv_freq", + "original_inv_freq", + "inv_timescales", + "softcap", +) def remap_gemma4_audio_key(key: str) -> str: @@ -391,6 +398,9 @@ def cactus_name_for_tensor(name: str, family: str, num_layers: int | None = None adapter_name = adapter_key_for_family(name, family) if family in {"parakeet", "parakeet_tdt"} and adapter_name.endswith(".conv.norm.num_batches_tracked"): return NameMatch(name, None, "transcription", True, hf_name=hf_name, adapter_name=adapter_name) + if family == "gemma4" and adapter_name.endswith(GEMMA4_RUNTIME_BUFFER_SUFFIXES): + out = adapter_name.replace(".", "_") + ".weights" + return NameMatch(name, out, component_for_name(hf_name, out), True, hf_name=hf_name, adapter_name=adapter_name) global_name = _global_match(adapter_name, family) if global_name: return NameMatch(name, global_name, _component_for_family(hf_name, global_name, family), True, hf_name=hf_name, adapter_name=adapter_name) diff --git a/python/cactus/py.typed b/python/cactus/py.typed deleted file mode 100644 index e69de29bb..000000000 diff --git a/python/cactus/transpiler/Fusions/fusions.py b/python/cactus/transpiler/Fusions/fusions.py index a70d7625b..be8cc2723 100644 --- a/python/cactus/transpiler/Fusions/fusions.py +++ b/python/cactus/transpiler/Fusions/fusions.py @@ -376,12 +376,6 @@ def _index_by_root_op( M.AttrCapture("hi", "clamp", "max", required=False), ), ), - "glu": _single_node_graph( - "glu", - "glu", - ("x",), - attr_captures=(M.AttrCapture("axis", "glu", "dim", default=-1, required=False),), - ), "groupnorm": _single_node_graph( "groupnorm", "groupnorm", @@ -1091,6 +1085,7 @@ def _index_by_root_op( "scaled_dot_product_attention": _definition("scaled_dot_product_attention", "attention", ATTENTION_DIRECT_GRAPH, fusion_fields=("generic", "attention")), "attention_core": _definition("attention_core", "attention", ATTENTION_CORE_GRAPH, fusion_fields=("generic", "attention", "gemma4_attention", "qwen2_5_attention", "lfm_attention")), "attention_masked": _definition("attention_masked", "attention", ATTENTION_MASKED_GRAPH, fusion_fields=("generic", "attention", "gemma4_attention")), + "gemma4_bmm_masked_attention": _definition("gemma4_bmm_masked_attention", "attention", LFM_BMM_MASKED_ATTENTION_GRAPH, fusion_fields=("generic", "attention", "gemma4_attention")), "lfm_bmm_masked_attention": _definition("lfm_bmm_masked_attention", "attention", LFM_BMM_MASKED_ATTENTION_GRAPH, fusion_fields=("generic", "attention", "lfm_attention", "lfm_moe")), "rope": _definition("rope", "rope", ROPE_GRAPH, fusion_fields=("generic", "rope", "gemma4_rope", "qwen2_5_rope")), "conv": _definition("conv", "conv1d", CONV_GRAPH, fusion_fields=("generic", "conv", "audio", "vision"), metadata=_required_attrs(padding=0, dilation=1, groups=1)), diff --git a/python/cactus/transpiler/Generator/component_split.py b/python/cactus/transpiler/Generator/component_split.py new file mode 100644 index 000000000..f7b4c14f5 --- /dev/null +++ b/python/cactus/transpiler/Generator/component_split.py @@ -0,0 +1,485 @@ +from __future__ import annotations + +from collections.abc import Mapping +from dataclasses import dataclass, field +from typing import Any + +from ..Fusions import models as FModels +from ..IR import models as IRModels + + +PREFILL_WITH_CACHE_TASK = "prefill_with_cache" +DECODE_WITH_CACHE_TASK = "decode_with_cache" + + +@dataclass(slots=True, frozen=True) +class PlaceholderSpec: + name: str + logical_name: str + source_node: str | None = None + tensor_node: str | None = None + target: str | None = None + value_kind: str = FModels.ValueKind.USER_INPUT + + +@dataclass(slots=True, frozen=True) +class OutputSpec: + node: str + logical_name: str + + +@dataclass(slots=True, frozen=True) +class ComponentSplitSpec: + name: str + graph: IRModels.Graph + outputs: tuple[OutputSpec, ...] + placeholders: tuple[PlaceholderSpec, ...] = () + ref_aliases: Mapping[str, str] = field(default_factory=dict) + input_aliases: Mapping[str, str] = field(default_factory=dict) + + +@dataclass(slots=True, frozen=True) +class Gemma4Boundaries: + vision_features: str + audio_features: str + merged_inputs_embeds: str + prefill_text_inputs_embeds: str + prefill_token_ids: str + prefill_per_layer_inputs: str + decode_text_inputs_embeds: str + decode_token_ids: str + decode_per_layer_inputs: str + decode_position_ids: str + prefill_logits: str + decode_logits: str + + +def split_component_graphs( + graphs: Mapping[str, IRModels.Graph], + model_profile: Any | None = None, +) -> dict[str, IRModels.Graph] | None: + if not is_gemma4_profile(model_profile): + return None + + prefill = find_graph_by_task(graphs, PREFILL_WITH_CACHE_TASK) + decode = find_graph_by_task(graphs, DECODE_WITH_CACHE_TASK) + + if prefill is None or decode is None: + return None + + return split_gemma4_components(prefill, decode) + + +def split_gemma4_components(prefill: IRModels.Graph, decode: IRModels.Graph) -> dict[str, IRModels.Graph]: + boundaries = find_gemma4_boundaries(prefill, decode) + position_step = PlaceholderSpec("position_ids", "position_ids", source_node=boundaries.decode_position_ids) + position_chunk = PlaceholderSpec("position_ids", "position_ids", tensor_node="input_ids") + inputs_embeds_step = PlaceholderSpec("inputs_embeds", "inputs_embeds", source_node=boundaries.decode_text_inputs_embeds) + inputs_embeds_chunk = PlaceholderSpec("inputs_embeds", "inputs_embeds", source_node=boundaries.merged_inputs_embeds) + per_layer_step = PlaceholderSpec("per_layer_inputs", "per_layer_inputs", source_node=boundaries.decode_per_layer_inputs) + per_layer_chunk = PlaceholderSpec("per_layer_inputs", "per_layer_inputs", source_node=boundaries.prefill_per_layer_inputs) + + specs = ( + ComponentSplitSpec( + name="vision_encoder", + graph=prefill, + outputs=(OutputSpec(boundaries.vision_features, "image_features"),), + input_aliases={"image_position_ids": "pixel_position_ids"}, + ), + ComponentSplitSpec( + name="audio_encoder", + graph=prefill, + outputs=(OutputSpec(boundaries.audio_features, "audio_features"),), + ), + ComponentSplitSpec( + name="lm_encoder_step", + graph=decode, + outputs=( + OutputSpec(boundaries.decode_text_inputs_embeds, "inputs_embeds"), + OutputSpec(boundaries.decode_per_layer_inputs, "per_layer_inputs"), + OutputSpec(position_step.name, "position_ids"), + ), + placeholders=(position_step,), + ref_aliases={boundaries.decode_token_ids: "input_ids"}, + ), + ComponentSplitSpec( + name="lm_encoder_text_chunk", + graph=prefill, + outputs=( + OutputSpec(boundaries.prefill_text_inputs_embeds, "inputs_embeds"), + OutputSpec(boundaries.prefill_per_layer_inputs, "per_layer_inputs"), + OutputSpec(position_chunk.name, "position_ids"), + ), + placeholders=(position_chunk,), + ref_aliases={ + boundaries.merged_inputs_embeds: boundaries.prefill_text_inputs_embeds, + boundaries.prefill_token_ids: "input_ids", + }, + ), + ComponentSplitSpec( + name="lm_encoder_media_step", + graph=decode, + outputs=( + OutputSpec(inputs_embeds_step.name, "inputs_embeds"), + OutputSpec(boundaries.decode_per_layer_inputs, "per_layer_inputs"), + OutputSpec(position_step.name, "position_ids"), + ), + placeholders=(inputs_embeds_step, position_step), + ref_aliases={boundaries.decode_token_ids: "input_ids"}, + ), + ComponentSplitSpec( + name="decoder_prefill_chunk", + graph=prefill, + outputs=(OutputSpec(boundaries.prefill_logits, "logits"),), + placeholders=(inputs_embeds_chunk, per_layer_chunk, position_chunk), + ), + ComponentSplitSpec( + name="decoder_step", + graph=decode, + outputs=(OutputSpec(boundaries.decode_logits, "logits"),), + placeholders=(inputs_embeds_step, per_layer_step, position_step), + ), + ) + + return {spec.name: extract_component_graph(spec) for spec in specs} + + +def is_gemma4_profile(model_profile: Any | None) -> bool: + profile_name = str(getattr(model_profile, "model_profiles", "") or "").lower() + return "gemma4" in profile_name or "gemma_4" in profile_name + + +def find_graph_by_task(graphs: Mapping[str, IRModels.Graph], task: str) -> IRModels.Graph | None: + for graph in graphs.values(): + if graph.task == task: + return graph + + for name, graph in graphs.items(): + if task in name: + return graph + + return None + + +def find_gemma4_boundaries(prefill: IRModels.Graph, decode: IRModels.Graph) -> Gemma4Boundaries: + vision_merge, audio_merge = find_gemma4_masked_scatter_merges(prefill) + + return Gemma4Boundaries( + vision_features=feature_parent(vision_merge, "vision"), + audio_features=feature_parent(audio_merge, "audio"), + merged_inputs_embeds=audio_merge.name, + prefill_text_inputs_embeds=find_text_inputs_embeds(prefill), + prefill_token_ids=find_embedding_token_ids(prefill), + prefill_per_layer_inputs=find_per_layer_inputs(prefill), + decode_text_inputs_embeds=find_text_inputs_embeds(decode), + decode_token_ids=find_embedding_token_ids(decode), + decode_per_layer_inputs=find_per_layer_inputs(decode), + decode_position_ids=find_decode_position_ids(decode), + prefill_logits=find_logits_output(prefill), + decode_logits=find_logits_output(decode), + ) + + +def find_gemma4_masked_scatter_merges(graph: IRModels.Graph) -> tuple[IRModels.Node, IRModels.Node]: + merges = [ + node + for node in graph.nodes + if node.target == "aten.masked_scatter.default" + and tensor_rank(node) == 3 + and tensor_last_dim(node) == 1536 + and len(node.parents) >= 3 + ] + + if len(merges) < 2: + raise ValueError("Gemma4 component split requires the vision/audio masked_scatter merge nodes") + + return tuple(sorted(merges, key=lambda node: node.index)[:2]) + + +def feature_parent(merge: IRModels.Node, role: str) -> str: + if len(merge.parents) < 3: + raise ValueError(f"{merge.name}: Gemma4 {role} merge does not have a feature parent") + + return merge.parents[2].name + + +def find_text_inputs_embeds(graph: IRModels.Graph) -> str: + for node in graph.nodes: + if node.target not in {"aten.mul.Tensor", "cactus.multiply"} or tensor_rank(node) != 3 or tensor_last_dim(node) != 1536: + continue + + if any(parent.target in {"aten.embedding.default", "cactus.embedding_from_tensor"} for parent in node.parents): + return node.name + + raise ValueError("Gemma4 component split could not find text inputs_embeds") + + +def find_embedding_token_ids(graph: IRModels.Graph) -> str: + text_inputs_embeds = graph.nodes_map[find_text_inputs_embeds(graph)] + + for parent in text_inputs_embeds.parents: + if parent.target not in {"aten.embedding.default", "cactus.embedding_from_tensor"}: + continue + + for embedding_parent in parent.parents: + if embedding_parent.value_kind != FModels.ValueKind.PARAMETER: + return embedding_parent.name + + raise ValueError("Gemma4 component split could not find embedding token-id input") + + +def find_per_layer_inputs(graph: IRModels.Graph) -> str: + candidates = [ + node + for node in graph.nodes + if node.target in {"aten.mul.Tensor", "aten.mul.Scalar", "cactus.multiply", "cactus.scalar_multiply"} + and tensor_shape(node)[-2:] == [35, 256] + and len(node.children) >= 10 + ] + + if not candidates: + raise ValueError("Gemma4 component split could not find per_layer_inputs") + + return max(candidates, key=lambda node: len(node.children)).name + + +def find_decode_position_ids(graph: IRModels.Graph) -> str: + for node in graph.nodes: + if node.target == "aten.unsqueeze.default" and tensor_shape(node) == [1, 1]: + return node.name + + for node in graph.nodes: + if node.name == "cache_position": + return node.name + + raise ValueError("Gemma4 component split could not find decode position_ids") + + +def find_logits_output(graph: IRModels.Graph) -> str: + output = graph.outputs[0] if graph.outputs else None + refs = IRModels.extract_node_refs((output.args, output.kwargs)) if output is not None else () + + if refs: + return refs[0] + + for node in reversed(graph.nodes): + if tensor_rank(node) == 3 and tensor_last_dim(node) == 262144: + return node.name + + raise ValueError("Gemma4 component split could not find logits output") + + +def extract_component_graph(spec: ComponentSplitSpec) -> IRModels.Graph: + placeholder_specs = {placeholder.name: placeholder for placeholder in spec.placeholders} + ref_replacements = { + **dict(spec.ref_aliases), + **{ + placeholder.source_node: placeholder.name + for placeholder in spec.placeholders + if placeholder.source_node is not None + }, + } + required_names = collect_required_node_names(spec, placeholder_specs, ref_replacements) + output_names = tuple(ref_replacements.get(output.node, output.node) for output in spec.outputs) + output_logical_names = { + ref_replacements.get(output.node, output.node): output.logical_name + for output in spec.outputs + } + used_names: set[str] = set() + nodes: list[IRModels.Node] = [] + + for placeholder in spec.placeholders: + if placeholder.name not in required_names and placeholder.name not in output_names: + continue + + nodes.append(create_placeholder_node(spec, placeholder, output_logical_names, used_names)) + + for node in spec.graph.nodes: + if node.name not in required_names: + continue + + if node.name in ref_replacements: + continue + + nodes.append(clone_component_node(spec, node, ref_replacements, output_logical_names, used_names)) + + nodes.append(create_output_node(spec.name, output_names, nodes)) + return IRModels.rebuild_graph(tuple(nodes), spec.graph) + + +def collect_required_node_names( + spec: ComponentSplitSpec, + placeholder_specs: Mapping[str, PlaceholderSpec], + ref_replacements: Mapping[str, str], +) -> frozenset[str]: + required: set[str] = set() + + for output in spec.outputs: + collect_node_ancestors(output.node, spec.graph, required, placeholder_specs, ref_replacements) + + for placeholder in spec.placeholders: + if placeholder.source_node is None: + required.add(placeholder.name) + + return frozenset(required) + + +def collect_node_ancestors( + node_name: str, + graph: IRModels.Graph, + required: set[str], + placeholder_specs: Mapping[str, PlaceholderSpec], + ref_replacements: Mapping[str, str], +) -> None: + stack = [node_name] + + while stack: + current_name = stack.pop() + effective_name = ref_replacements.get(current_name, current_name) + + if effective_name in required: + continue + + required.add(effective_name) + + if effective_name in placeholder_specs: + continue + + node = graph.nodes_map.get(effective_name) + + if node is None: + raise ValueError(f"Component split references missing node {effective_name}") + + stack.extend(parent.name for parent in node.parents) + + +def create_placeholder_node( + spec: ComponentSplitSpec, + placeholder: PlaceholderSpec, + output_logical_names: Mapping[str, str], + used_names: set[str], +) -> IRModels.Node: + template = source_template_node(spec.graph, placeholder) + name = unique_component_node_name(placeholder.name, used_names) + metadata = dict(template.ir_metadata) if template is not None else {} + metadata["logical_input"] = placeholder.logical_name + + if placeholder.name in output_logical_names: + metadata["logical_output"] = output_logical_names[placeholder.name] + + return IRModels.Node( + index=template.index if template is not None else -1, + name=name, + node_type="placeholder", + target=placeholder.target or placeholder.logical_name, + args=[], + kwargs={}, + users=(), + tensor_output_meta=template.tensor_output_meta if template is not None else None, + module_stack=template.module_stack if template is not None else None, + value_kind=placeholder.value_kind, + attrs={}, + ir_metadata=metadata, + cache=template.cache if template is not None else None, + ) + + +def source_template_node(graph: IRModels.Graph, placeholder: PlaceholderSpec) -> IRModels.Node | None: + for name in (placeholder.tensor_node, placeholder.source_node, placeholder.name): + if name is not None and name in graph.nodes_map: + return graph.nodes_map[name] + + return None + + +def clone_component_node( + spec: ComponentSplitSpec, + node: IRModels.Node, + ref_replacements: Mapping[str, str], + output_logical_names: Mapping[str, str], + used_names: set[str], +) -> IRModels.Node: + metadata = dict(node.ir_metadata) + + if node.is_placeholder: + logical_input = spec.input_aliases.get(node.name) or spec.input_aliases.get(node.target) + + if logical_input is not None: + metadata["logical_input"] = logical_input + + if node.name in output_logical_names: + metadata["logical_output"] = output_logical_names[node.name] + + name = unique_component_node_name(node.name, used_names) + + return IRModels.Node( + index=node.index, + name=name, + node_type=node.node_type, + target=node.target, + args=IRModels.rewrite_node_refs(node.args, dict(ref_replacements)), + kwargs=IRModels.rewrite_node_refs(node.kwargs, dict(ref_replacements)), + users=(), + tensor_output_meta=node.tensor_output_meta, + module_stack=node.module_stack, + value_kind=node.value_kind, + attrs=dict(node.attrs), + ir_metadata=metadata, + cache=node.cache, + ) + + +def create_output_node( + component_name: str, + output_names: tuple[str, ...], + nodes: list[IRModels.Node], +) -> IRModels.Node: + max_index = max((node.index for node in nodes), default=-1) + + return IRModels.Node( + index=max_index + 1, + name=f"{component_name}_output", + node_type="output", + target="output", + args=[[{"node": name} for name in output_names]], + kwargs={}, + users=(), + tensor_output_meta=None, + module_stack=None, + value_kind=FModels.ValueKind.OUTPUT, + attrs={}, + ir_metadata={}, + ) + + +def unique_component_node_name(name: str, used_names: set[str]) -> str: + if name not in used_names: + used_names.add(name) + return name + + suffix = 1 + candidate = f"{name}_{suffix}" + + while candidate in used_names: + suffix += 1 + candidate = f"{name}_{suffix}" + + used_names.add(candidate) + return candidate + + +def tensor_shape(node: IRModels.Node) -> list[Any]: + if not isinstance(node.tensor_output_meta, dict): + return [] + + shape = node.tensor_output_meta.get("shape") + return list(shape) if isinstance(shape, list) else [] + + +def tensor_rank(node: IRModels.Node) -> int: + return len(tensor_shape(node)) + + +def tensor_last_dim(node: IRModels.Node) -> Any | None: + shape = tensor_shape(node) + return shape[-1] if shape else None diff --git a/python/cactus/transpiler/Generator/constants.py b/python/cactus/transpiler/Generator/constants.py index 520b2e790..fc64a644e 100644 --- a/python/cactus/transpiler/Generator/constants.py +++ b/python/cactus/transpiler/Generator/constants.py @@ -8,13 +8,13 @@ "torch.bfloat16": "FP16", "torch.float32": "FP32", "torch.float": "FP32", - "torch.int8": "INT8", - "torch.uint8": "INT8", - "torch.int16": "INT8", - "torch.int32": "INT8", - "torch.int64": "INT8", - "torch.long": "INT8", - "torch.bool": "INT8", + "torch.int8": "FP32", + "torch.uint8": "FP32", + "torch.int16": "FP32", + "torch.int32": "FP32", + "torch.int64": "FP32", + "torch.long": "FP32", + "torch.bool": "FP32", } WEIGHT_VALUE_KINDS = { @@ -62,11 +62,25 @@ "cactus.multiply": "multiply", "cactus.divide": "divide", "cactus.not_equal": "not_equal", + "cactus.equal": "equal", + "cactus.less": "less", + "cactus.less_equal": "less_equal", + "cactus.greater": "greater", + "cactus.greater_equal": "greater_equal", + "cactus.bitwise_and": "bitwise_and", + "cactus.bitwise_or": "bitwise_or", "aten.add.Tensor": "add", "aten.sub.Tensor": "subtract", "aten.mul.Tensor": "multiply", "aten.div.Tensor": "divide", "aten.ne.Tensor": "not_equal", + "aten.eq.Tensor": "equal", + "aten.lt.Tensor": "less", + "aten.le.Tensor": "less_equal", + "aten.gt.Tensor": "greater", + "aten.ge.Tensor": "greater_equal", + "aten.bitwise_and.Tensor": "bitwise_and", + "aten.bitwise_or.Tensor": "bitwise_or", } SCALAR_TARGETS: dict[str, tuple[str, str]] = { @@ -76,12 +90,22 @@ "cactus.scalar_divide": ("scalar_divide", "value"), "cactus.scalar_floor_divide": ("scalar_floor_divide", "value"), "cactus.scalar_not_equal": ("scalar_not_equal", "value"), + "cactus.scalar_equal": ("scalar_equal", "value"), + "cactus.scalar_less": ("scalar_less", "value"), + "cactus.scalar_less_equal": ("scalar_less_equal", "value"), + "cactus.scalar_greater": ("scalar_greater", "value"), + "cactus.scalar_greater_equal": ("scalar_greater_equal", "value"), "aten.add.Scalar": ("scalar_add", "other"), "aten.sub.Scalar": ("scalar_subtract", "other"), "aten.mul.Scalar": ("scalar_multiply", "other"), "aten.div.Scalar": ("scalar_divide", "other"), "aten.floor_divide.default": ("scalar_floor_divide", "other"), "aten.ne.Scalar": ("scalar_not_equal", "other"), + "aten.eq.Scalar": ("scalar_equal", "other"), + "aten.lt.Scalar": ("scalar_less", "other"), + "aten.le.Scalar": ("scalar_less_equal", "other"), + "aten.gt.Scalar": ("scalar_greater", "other"), + "aten.ge.Scalar": ("scalar_greater_equal", "other"), } UNARY_TARGETS: dict[str, str] = { @@ -97,6 +121,8 @@ "cactus.gelu_erf": "gelu_erf", "cactus.sigmoid": "sigmoid", "cactus.tanh": "tanh", + "cactus.logical_not": "logical_not", + "cactus.bitwise_not": "bitwise_not", "aten.abs.default": "abs", "aten.exp.default": "scalar_exp", "aten.sqrt.default": "scalar_sqrt", @@ -108,6 +134,12 @@ "aten.gelu.default": "gelu", "aten.sigmoid.default": "sigmoid", "aten.tanh.default": "tanh", + "aten.logical_not.default": "logical_not", + "aten.bitwise_not.default": "bitwise_not", +} + +LOG1P_TARGETS = { + "aten.log1p.default", } POW_TARGETS = { @@ -127,6 +159,7 @@ "aten.var.correction": "variance", "aten.amax.default": "max", "aten.amin.default": "min", + "aten.any.dim": "max", } SHAPE_TARGETS: dict[str, str] = { @@ -164,6 +197,15 @@ INDEX_TARGETS = { "cactus.index", + "aten.index.Tensor", +} + +WHERE_TARGETS = { + "aten.where.self", +} + +UNFOLD_TARGETS = { + "aten.unfold.default", } CAT_TARGETS = { @@ -304,17 +346,8 @@ } UNSUPPORTED_SEMANTIC_TARGETS = { - "aten.any.dim", - "aten.bitwise_and.Tensor", - "aten.bitwise_or.Tensor", - "aten.eq.Scalar", - "aten.eq.Tensor", "aten.fmod.Scalar", - "aten.index.Tensor", - "aten.le.Tensor", - "aten.logical_not.default", "aten.slice_scatter.default", - "aten.where.self", } OPEN_SLICE_END = 9_223_372_036_854_775_807 diff --git a/python/cactus/transpiler/Generator/generate.py b/python/cactus/transpiler/Generator/generate.py index 037e81586..dc227ea72 100644 --- a/python/cactus/transpiler/Generator/generate.py +++ b/python/cactus/transpiler/Generator/generate.py @@ -2,9 +2,11 @@ import json from collections.abc import Mapping +from dataclasses import replace from pathlib import Path from typing import Any +from . import component_split from . import constants from . import models from ..Converter import models as CModels @@ -78,6 +80,40 @@ def generate_from_json( ) +def generate_bundle( + ir_output: ComponentInput, + bundle_dir: str | Path, + *, + model_profile: Any | None = None, + component_name: str | None = None, + weights_dir: str | Path | None = None, + weights_manifest_path: str | Path | None = None, + metadata: dict[str, str] | None = None, + strict: bool = True, + allow_unsupported_ops: bool = False, +) -> models.GenerationResult: + bundle_path = Path(bundle_dir) + components_dir = bundle_path / "components" + result = generate( + ir_output, + output_dir=components_dir, + model_profile=model_profile, + component_name=component_name, + weights_dir=weights_dir, + weights_manifest_path=weights_manifest_path, + strict=strict, + allow_unsupported_ops=allow_unsupported_ops, + ) + plan = RPModels.runtime_plan_from_generation_result( + result, + bundle_dir=bundle_path, + model_profile=model_profile, + metadata=metadata, + ) + engine_manifest_path, runtime_plan_path = plan.write(bundle_path) + return replace(result, engine_manifest_path=engine_manifest_path, runtime_plan_path=runtime_plan_path) + + ################################################# Generator Utils!!!!!!! ################################################# @@ -93,9 +129,18 @@ def component_graphs_from_input( model_profile: Any | None, ) -> tuple[models.ComponentGraph, ...]: if isinstance(ir_output, Mapping): + graphs = {name: graph_from_input(component_input) for name, component_input in ir_output.items()} + split_graphs = component_split.split_component_graphs(graphs, model_profile) + + if split_graphs is not None: + return tuple( + models.ComponentGraph.from_ir(name, graph, config) + for name, graph in split_graphs.items() + ) + return tuple( - models.ComponentGraph.from_ir(name, graph_from_input(component_input), config) - for name, component_input in ir_output.items() + models.ComponentGraph.from_ir(name, graph, config) + for name, graph in graphs.items() ) graph = graph_from_input(ir_output) @@ -160,6 +205,7 @@ def build_lowering_rules( add_rules(rules, constants.BINARY_TARGETS, lower_binary) add_rules(rules, constants.SCALAR_TARGETS, lower_scalar) add_rules(rules, constants.UNARY_TARGETS, lower_unary) + add_rules(rules, constants.LOG1P_TARGETS, lower_log1p) add_rules(rules, constants.POW_TARGETS, lower_pow) add_rules(rules, constants.REDUCE_TARGETS, lower_reduce) add_rules(rules, constants.SHAPE_TARGETS, lower_shape) @@ -169,6 +215,8 @@ def build_lowering_rules( add_rules(rules, constants.TRANSPOSE_TARGETS, lower_transpose) add_rules(rules, constants.SLICE_TARGETS, lower_slice) add_rules(rules, constants.INDEX_TARGETS, lower_index) + add_rules(rules, constants.WHERE_TARGETS, lower_where) + add_rules(rules, constants.UNFOLD_TARGETS, lower_unfold) add_rules(rules, constants.CAT_TARGETS, lower_cat) add_rules(rules, constants.MATMUL_TARGETS, lower_matmul) add_rules(rules, constants.ADDM_CONST_TARGETS, lower_addmm) @@ -264,7 +312,7 @@ def lower_placeholder(context: models.GenerationContext, node: IRModels.Node) -> return lower_cache_placeholder(context, node) shape, dynamic_dims = graph_input_shape(node) - dtype = cactus_precision(context.graph, tensor_dtype(node)) + dtype = placeholder_precision(context, node) tensor = context.graph.input(shape, dtype=dtype, dynamic_dims=dynamic_dims if any(dynamic_dims) else None) logical_name = logical_input_name(context.component.ir_graph, node) @@ -276,6 +324,15 @@ def lower_placeholder(context: models.GenerationContext, node: IRModels.Node) -> return tensor +def placeholder_precision(context: models.GenerationContext, node: IRModels.Node) -> int: + if node.value_kind in constants.WEIGHT_VALUE_KINDS and context.component.weight_resolver is not None: + record = context.component.weight_resolver.resolve(node.name) + if record is not None and record.precision and hasattr(context.graph, record.precision): + return int(getattr(context.graph, record.precision)) + + return cactus_precision(context.graph, tensor_dtype(node)) + + def should_lower_cache_placeholder_as_state(node: IRModels.Node) -> bool: if node.cache is None: return False @@ -495,6 +552,15 @@ def lower_scalar(context: models.GenerationContext, node: IRModels.Node) -> Any: value = scalar_attr(node, attr_name) if value is None: + value = scalar_attr(node, "arg_1") + + if value is None: + if method == "scalar_equal": + return context.graph.scalar_multiply(inputs[0], 0.0) + + if method == "scalar_not_equal": + return context.graph.scalar_add(context.graph.scalar_multiply(inputs[0], 0.0), 1.0) + raise UnsupportedLoweringError(f"{node.name}: scalar lowering {node.target} missing attr {attr_name}") return getattr(context.graph, method)(inputs[0], value) @@ -504,10 +570,21 @@ def align_binary_inputs(context: models.GenerationContext, node: IRModels.Node, target_shape = output_shape(node) left = align_input_to_shape(context, inputs[0], node.parents[0], target_shape) right = align_input_to_shape(context, inputs[1], node.parents[1], target_shape) - return left, right + return align_binary_precision(context, left, right) + + +def align_binary_precision(context: models.GenerationContext, left: Any, right: Any) -> tuple[Any, Any]: + left_precision = getattr(left, "dtype", None) + right_precision = getattr(right, "dtype", None) + + if left_precision is None or right_precision is None or left_precision == right_precision: + return left, right + + return left, cast_to_precision(context, right, left_precision) def align_input_to_shape(context: models.GenerationContext, value: Any, source_node: IRModels.Node, target_shape: tuple[Any, ...]) -> Any: + value = align_value_to_declared_shape(context, value, source_node) source_shape = meta_shape(source_node) if shape_matches_tensor(source_node, target_shape): @@ -519,12 +596,56 @@ def align_input_to_shape(context: models.GenerationContext, value: Any, source_n return value +def align_value_to_declared_shape(context: models.GenerationContext, value: Any, source_node: IRModels.Node) -> Any: + actual_shape = tuple(getattr(value, "shape", ())) + declared_shape = meta_shape(source_node) + + if not actual_shape or not declared_shape or len(declared_shape) <= len(actual_shape): + return value + + resolved_shape = resolve_declared_shape_from_actual(declared_shape, actual_shape) + + if resolved_shape is None or tuple(resolved_shape) == actual_shape: + return value + + return context.graph.reshape(value, resolved_shape) + + +def resolve_declared_shape_from_actual(declared_shape: tuple[Any, ...], actual_shape: tuple[int, ...]) -> tuple[int, ...] | None: + actual_index = len(actual_shape) - 1 + resolved = [1 for _ in declared_shape] + + for declared_index in range(len(declared_shape) - 1, -1, -1): + declared_dim = concrete_dim(declared_shape[declared_index]) + + if actual_index >= 0 and (declared_dim is None or declared_dim == actual_shape[actual_index]): + resolved[declared_index] = actual_shape[actual_index] + actual_index -= 1 + continue + + if declared_dim == 1: + resolved[declared_index] = 1 + continue + + return None + + if actual_index >= 0: + return None + + return tuple(resolved) + + def lower_unary(context: models.GenerationContext, node: IRModels.Node) -> Any: inputs = require_input_count(context, node, 1) method = constants.UNARY_TARGETS[node.target] return getattr(context.graph, method)(inputs[0]) +def lower_log1p(context: models.GenerationContext, node: IRModels.Node) -> Any: + inputs = require_input_count(context, node, 1) + return context.graph.scalar_log(context.graph.scalar_add(inputs[0], 1.0)) + + def lower_pow(context: models.GenerationContext, node: IRModels.Node) -> Any: inputs = require_input_count(context, node, 1) exponent = scalar_attr(node, "exponent") @@ -543,7 +664,18 @@ def lower_reduce(context: models.GenerationContext, node: IRModels.Node) -> Any: if axis is None: raise UnsupportedLoweringError(f"{node.name}: reduce lowering missing axis/dim attr") - return getattr(context.graph, method)(inputs[0], axis) + normalized_axis = normalize_dim(axis, len(meta_shape(node.parents[0]))) + reduced = getattr(context.graph, method)(inputs[0], normalized_axis) + target_shape = output_shape(node) + dropped_shape = reduction_dropped_shape(meta_shape(node.parents[0]), normalized_axis) + + if bool(node.attrs.get("keepdim", False)) or ( + tuple(dropped_shape) != tuple(target_shape) + and element_count(dropped_shape) == element_count(target_shape) + ): + return context.graph.reshape(reduced, target_shape) + + return reduced def lower_shape(context: models.GenerationContext, node: IRModels.Node) -> Any: @@ -589,6 +721,10 @@ def lower_flatten(context: models.GenerationContext, node: IRModels.Node) -> Any def lower_transpose(context: models.GenerationContext, node: IRModels.Node) -> Any: inputs = require_input_count(context, node, 1) + + if can_alias_quantized_weight_transpose(context, node, inputs[0]): + return inputs[0] + permutation = node.attrs.get("permutation") if permutation is not None: @@ -606,6 +742,54 @@ def lower_transpose(context: models.GenerationContext, node: IRModels.Node) -> A return context.graph.permute(inputs[0], swap_permutation(parent_rank(node), int(dim0), int(dim1))) +def can_alias_quantized_weight_transpose(context: models.GenerationContext, node: IRModels.Node, value: Any) -> bool: + return ( + is_weight_transpose_node(node) + and is_cq_tensor(context, value) + and bool(node.children) + and all(child.target in constants.MATMUL_TARGETS for child in node.children) + ) + + +def is_weight_transpose_node(node: IRModels.Node) -> bool: + if not node.parents: + return False + + parent = node.parents[0] + if not parent.is_placeholder or parent.value_kind not in constants.WEIGHT_VALUE_KINDS: + return False + + if node.target == "aten.t.default": + return True + + rank = len(meta_shape(parent)) + if rank < 2: + return False + + permutation = node.attrs.get("permutation") + if permutation is not None: + return tuple_ints(permutation) == swap_permutation(rank, rank - 2, rank - 1) + + dim0 = node.attrs.get("dim0") + dim1 = node.attrs.get("dim1") + if dim0 is None or dim1 is None: + return False + + return { + normalize_dim(int(dim0), rank), + normalize_dim(int(dim1), rank), + } == {rank - 2, rank - 1} + + +def is_cq_tensor(context: models.GenerationContext, value: Any) -> bool: + return getattr(value, "dtype", None) in { + context.graph.CQ1, + context.graph.CQ2, + context.graph.CQ3, + context.graph.CQ4, + } + + def lower_slice(context: models.GenerationContext, node: IRModels.Node) -> Any: inputs = require_input_count(context, node, 1) axis_value = node.attrs.get("axis") @@ -628,13 +812,18 @@ def lower_slice(context: models.GenerationContext, node: IRModels.Node) -> Any: step = node.attrs.get("step", 1) if step is not None and int(step) != 1: - raise UnsupportedLoweringError(f"{node.name}: Cactus slice lowering only supports step 1") + target_shape = output_shape(node) + length = target_shape[axis] if axis < len(target_shape) else slice_length(node, start) + return context.graph.strided_slice(inputs[0], axis, start, length, int(step)) length = slice_length(node, start) return context.graph.slice(inputs[0], axis, start, length=length) def lower_index(context: models.GenerationContext, node: IRModels.Node) -> Any: + if node.target == "aten.index.Tensor": + return lower_index_tensor(context, node) + inputs = require_input_count(context, node, 1) index_value = node.attrs.get("index_value", node.attrs.get("index")) @@ -644,6 +833,60 @@ def lower_index(context: models.GenerationContext, node: IRModels.Node) -> Any: return context.graph.index(inputs[0], int(index_value), axis=int(node.attrs.get("axis", node.attrs.get("dim", 0)))) +def lower_index_tensor(context: models.GenerationContext, node: IRModels.Node) -> Any: + inputs = context.inputs_for(node) + + if len(inputs) < 2: + raise unsupported_arity(node, len(inputs), "source tensor plus at least one index tensor") + + source = inputs[0] + source_shape = meta_shape(node.parents[0]) + target_shape = meta_shape(node) + + if len(source_shape) == 3 and len(target_shape) == 2 and concrete_dim(source_shape[0]) == 1: + feature_count = concrete_dim(source_shape[1]) + feature_dim = concrete_dim(source_shape[2]) + + if feature_count is not None and feature_dim is not None: + return context.graph.reshape(source, (feature_count, feature_dim)) + + if len(source_shape) == 2 and len(target_shape) == 4 and concrete_dim(source_shape[0]) == 1: + source_length = concrete_dim(source_shape[1]) + target_length = concrete_dim(target_shape[-1]) + + if source_length is not None and target_length is not None: + reshaped = context.graph.reshape(source, (1, 1, 1, source_length)) + + if target_length < source_length: + return context.graph.slice(reshaped, 3, 0, length=target_length) + + if target_length == source_length: + return reshaped + + concrete_target_shape = concrete_shape(target_shape) + + if concrete_target_shape is not None and element_count(source_shape) == element_count(concrete_target_shape): + return context.graph.reshape(source, concrete_target_shape) + + raise UnsupportedLoweringError(f"{node.name}: unsupported aten.index.Tensor pattern") + + +def lower_where(context: models.GenerationContext, node: IRModels.Node) -> Any: + inputs = require_input_count(context, node, 3) + output_precision = cactus_precision(context.graph, tensor_dtype(node)) + true_value = cast_to_precision(context, inputs[1], output_precision) + false_value = cast_to_precision(context, inputs[2], output_precision) + return context.graph.where(inputs[0], true_value, false_value) + + +def lower_unfold(context: models.GenerationContext, node: IRModels.Node) -> Any: + inputs = require_input_count(context, node, 1) + dimension = int(node.attrs.get("dimension", node.attrs.get("dim", node.attrs.get("arg_1", 0)))) + size = int(node.attrs.get("size", node.attrs.get("arg_2", 1))) + step = int(node.attrs.get("step", node.attrs.get("arg_3", 1))) + return context.graph.unfold(inputs[0], dimension, size, step) + + def lower_cat(context: models.GenerationContext, node: IRModels.Node) -> Any: inputs = context.inputs_for(node) @@ -665,13 +908,33 @@ def lower_matmul(context: models.GenerationContext, node: IRModels.Node) -> Any: if node.target == "aten.bmm.default": return lower_bmm(context, node, inputs) + lhs = matmul_activation_operand(context, inputs[0]) + rhs = inputs[1] + + if len(node.parents) >= 2 and is_weight_transpose_node(node.parents[1]): + original_weight_node = node.parents[1].parents[0] + original_weight = context.lookup(original_weight_node.name) + + if original_weight is not None and is_cq_tensor(context, original_weight): + return context.graph.matmul(lhs, original_weight, pretransposed_rhs=True) + + if not is_cq_tensor(context, rhs): + rhs = matmul_activation_operand(context, rhs) + return context.graph.matmul( - inputs[0], - inputs[1], + lhs, + rhs, pretransposed_rhs=bool(node.attrs.get("pretransposed_rhs", False)), ) +def matmul_activation_operand(context: models.GenerationContext, value: Any) -> Any: + if getattr(value, "dtype", None) == context.graph.FP32: + return cast_to_precision(context, value, context.graph.FP16) + + return value + + def empty_cat_passthrough(node: IRModels.Node, inputs: tuple[Any, ...]) -> Any | None: target_shape = meta_shape(node) non_empty_inputs = [] @@ -712,18 +975,37 @@ def lower_bmm(context: models.GenerationContext, node: IRModels.Node, inputs: tu if len(left_shape) != 3 or len(right_shape) != 3: raise UnsupportedLoweringError(f"{node.name}: bmm lowering requires rank-3 inputs") - if left_shape[0] != 1 or right_shape[0] != 1: - raise UnsupportedLoweringError(f"{node.name}: bmm lowering currently supports only batch size 1") + if left_shape[0] != right_shape[0]: + raise UnsupportedLoweringError(f"{node.name}: bmm lowering requires matching batch dimensions") + + batch_size = concrete_dim(left_shape[0]) - left_2d = context.graph.reshape(inputs[0], (left_shape[1], left_shape[2])) - right_2d = context.graph.reshape(inputs[1], (right_shape[1], right_shape[2])) - product = context.graph.matmul(left_2d, right_2d) - return context.graph.reshape(product, output_shape(node)) + if batch_size is None: + raise UnsupportedLoweringError(f"{node.name}: bmm lowering requires a concrete batch size") + + if batch_size == 1: + left_2d = matmul_activation_operand(context, context.graph.reshape(inputs[0], (left_shape[1], left_shape[2]))) + right_2d = matmul_activation_operand(context, context.graph.reshape(inputs[1], (right_shape[1], right_shape[2]))) + product = context.graph.matmul(left_2d, right_2d) + return context.graph.reshape(product, output_shape(node)) + + batch_outputs = [] + + for batch_index in range(batch_size): + left_2d = matmul_activation_operand(context, context.graph.index(inputs[0], batch_index, axis=0)) + right_2d = matmul_activation_operand(context, context.graph.index(inputs[1], batch_index, axis=0)) + product = context.graph.matmul(left_2d, right_2d) + batch_outputs.append(context.graph.reshape(product, (1, left_shape[1], right_shape[2]))) + + return context.graph.cat(batch_outputs, axis=0) def lower_addmm(context: models.GenerationContext, node: IRModels.Node) -> Any: inputs = require_input_count(context, node, 3) - product = context.graph.matmul(inputs[1], inputs[2]) + product = context.graph.matmul( + matmul_activation_operand(context, inputs[1]), + matmul_activation_operand(context, inputs[2]), + ) return context.graph.add(inputs[0], product) @@ -788,22 +1070,30 @@ def lower_copy(context: models.GenerationContext, node: IRModels.Node) -> Any: def lower_constant_pad_nd(context: models.GenerationContext, node: IRModels.Node) -> Any: inputs = require_input_count(context, node, 1) pads = node.attrs.get("pad", node.attrs.get("arg_1")) + value = scalar_attr(node, "value") + + if value is None: + value = scalar_attr(node, "arg_2") + + if value is None: + value = 0.0 if not isinstance(pads, (list, tuple)) or len(pads) % 2 != 0: raise UnsupportedLoweringError(f"{node.name}: constant_pad_nd lowering requires an even pad list") - if any(int(pad) > 0 for pad in pads): - raise UnsupportedLoweringError(f"{node.name}: positive constant_pad_nd requires a native pad op") - result = inputs[0] parent_shape = meta_shape(node.parents[0]) + current_shape = list(parent_shape) rank = len(parent_shape) + positive_pads = [0 for _ in pads] for pair_index in range(0, len(pads), 2): left_pad = int(pads[pair_index]) right_pad = int(pads[pair_index + 1]) + positive_pads[pair_index] = max(left_pad, 0) + positive_pads[pair_index + 1] = max(right_pad, 0) - if left_pad == 0 and right_pad == 0: + if left_pad >= 0 and right_pad >= 0: continue axis = rank - 1 - (pair_index // 2) @@ -811,7 +1101,7 @@ def lower_constant_pad_nd(context: models.GenerationContext, node: IRModels.Node if axis < 0: raise UnsupportedLoweringError(f"{node.name}: pad rank exceeds input rank") - axis_size = parent_shape[axis] + axis_size = current_shape[axis] if not isinstance(axis_size, int): raise UnsupportedLoweringError(f"{node.name}: cannot crop symbolic pad dimension {axis_size}") @@ -824,6 +1114,10 @@ def lower_constant_pad_nd(context: models.GenerationContext, node: IRModels.Node raise UnsupportedLoweringError(f"{node.name}: pad crop removes more than the full axis") result = context.graph.slice(result, axis, start, length=length) + current_shape[axis] = length + + if any(positive_pads): + return context.graph.pad(result, positive_pads, float(value)) return result @@ -864,7 +1158,8 @@ def lower_topk(context: models.GenerationContext, node: IRModels.Node) -> Any: def lower_gather(context: models.GenerationContext, node: IRModels.Node) -> Any: inputs = require_input_count(context, node, 2) - return context.graph.gather(inputs[0], inputs[1]) + axis = axis_attr(node, default=0) + return context.graph.gather(inputs[0], inputs[1], axis=axis if axis is not None else 0) def lower_embedding(context: models.GenerationContext, node: IRModels.Node) -> Any: @@ -874,13 +1169,29 @@ def lower_embedding(context: models.GenerationContext, node: IRModels.Node) -> A def lower_clamp(context: models.GenerationContext, node: IRModels.Node) -> Any: inputs = require_input_count(context, node, 1) - lo = node.attrs.get("lo", node.attrs.get("min")) - hi = node.attrs.get("hi", node.attrs.get("max")) + lo = node.attrs.get("lo", node.attrs.get("min", node.attrs.get("arg_1"))) + hi = node.attrs.get("hi", node.attrs.get("max", node.attrs.get("arg_2"))) - if lo is None or hi is None: - raise UnsupportedLoweringError(f"{node.name}: clamp lowering requires min/lo and max/hi attrs") + if lo is None and hi is None: + return inputs[0] - return context.graph.clamp(inputs[0], lo, hi) + result = inputs[0] + + if lo is not None: + keep = context.graph.scalar_greater_equal(result, lo) + result = context.graph.add( + context.graph.multiply(result, keep), + context.graph.scalar_multiply(context.graph.logical_not(keep), lo), + ) + + if hi is not None: + keep = context.graph.scalar_less_equal(result, hi) + result = context.graph.add( + context.graph.multiply(result, keep), + context.graph.scalar_multiply(context.graph.logical_not(keep), hi), + ) + + return result def lower_norm(context: models.GenerationContext, node: IRModels.Node) -> Any: @@ -923,6 +1234,11 @@ def lower_conv(context: models.GenerationContext, node: IRModels.Node) -> Any: return lower_generic_conv1d(context, node, inputs, bias) if target == "aten.conv2d.default" or rank == 4: + specialized = lower_specialized_conv2d(context, node, inputs, bias) + + if specialized is not None: + return specialized + return context.graph.conv2d( inputs[0], inputs[1], @@ -968,6 +1284,11 @@ def lower_named_conv(context: models.GenerationContext, node: IRModels.Node, met return getattr(context.graph, method)(inputs[0], inputs[1], bias=bias) if method == "conv2d": + specialized = lower_specialized_conv2d(context, node, inputs, bias) + + if specialized is not None: + return specialized + return context.graph.conv2d( inputs[0], inputs[1], @@ -982,8 +1303,53 @@ def lower_named_conv(context: models.GenerationContext, node: IRModels.Node, met def lower_generic_conv1d(context: models.GenerationContext, node: IRModels.Node, inputs: tuple[Any, ...], bias: Any | None) -> Any: + stride = first_int(node.attrs.get("stride"), 1) + padding = first_int(node.attrs.get("padding"), 0) + dilation = first_int(node.attrs.get("dilation"), 1) + groups = first_int(node.attrs.get("groups"), 1) + input_shape = meta_shape(node.parents[0]) if node.parents else () + weight_shape = meta_shape(node.parents[1]) if len(node.parents) > 1 else () + + if ( + padding == 0 + and dilation == 1 + and len(input_shape) == 3 + and len(weight_shape) == 3 + and groups == input_shape[1] + and weight_shape[0] == input_shape[1] + and weight_shape[1] == 1 + ): + return context.graph.conv1d_depthwise(inputs[0], inputs[1], bias=bias, stride=stride) + require_plain_conv1d_attrs(node, "conv1d") - return context.graph.conv1d(inputs[0], inputs[1], bias=bias, stride=first_int(node.attrs.get("stride"), 1)) + return context.graph.conv1d(inputs[0], inputs[1], bias=bias, stride=stride) + + +def lower_specialized_conv2d(context: models.GenerationContext, node: IRModels.Node, inputs: tuple[Any, ...], bias: Any | None) -> Any | None: + weight_shape = meta_shape(node.parents[1]) if len(node.parents) > 1 else () + stride = tuple_int_values(node.attrs.get("stride"), 1) + padding = tuple_int_values(node.attrs.get("padding"), 0) + dilation = tuple_int_values(node.attrs.get("dilation"), 1) + groups = first_int(node.attrs.get("groups"), 1) + + if len(weight_shape) != 4 or dilation != (1, 1): + return None + + kernel = tuple(int(dim) for dim in weight_shape[2:]) + + if kernel == (3, 3) and stride == (2, 2) and padding == (1, 1): + if groups != 1 and weight_shape[1] == 1: + return context.graph.conv2d_depthwise_k3s2p1(inputs[0], inputs[1], bias=bias) + + return context.graph.conv2d_k3s2p1(inputs[0], inputs[1], bias=bias) + + if kernel == (3, 3) and stride == (1, 1) and padding == (1, 1) and groups == 1: + return context.graph.conv2d_k3s1p1(inputs[0], inputs[1], bias=bias) + + if kernel == (1, 1) and stride == (1, 1) and padding == (0, 0) and groups == 1: + return context.graph.conv2d_pointwise_1x1(inputs[0], inputs[1], bias=bias) + + return None def require_plain_conv1d_attrs(node: IRModels.Node, method: str) -> None: @@ -1004,6 +1370,9 @@ def lower_attention(context: models.GenerationContext, node: IRModels.Node) -> A if node.target == "cactus.attention" or node.target == "aten.scaled_dot_product_attention.default": require_len(node, inputs, 3) query, key, value, mask = attention_inputs_for_layout(context, node, inputs) + query = cast_to_precision(context, query, context.graph.FP16) + key = cast_to_precision(context, key, context.graph.FP16) + value = cast_to_precision(context, value, context.graph.FP16) output = context.graph.attention( query, key, @@ -1275,11 +1644,20 @@ def lower_special_cactus(context: models.GenerationContext, node: IRModels.Node) if target == "cactus.rope": require_len(node, inputs, 1) - return context.graph.rope(inputs[0], theta=float(node.attrs.get("theta", 10_000.0)), position_offset=int(node.attrs.get("position_offset", 0))) + return context.graph.rope( + inputs[0], + theta=float(attr_value(node, "theta", 10_000.0)), + position_offset=int(attr_value(node, "position_offset", 0)), + ) if target == "cactus.rope_gptj": require_len(node, inputs, 1) - return context.graph.rope_gptj(inputs[0], theta=float(node.attrs.get("theta", 10_000.0)), position_offset=int(node.attrs.get("position_offset", 0)), rot_dim=int(node.attrs.get("rot_dim", 0))) + return context.graph.rope_gptj( + inputs[0], + theta=float(attr_value(node, "theta", 10_000.0)), + position_offset=int(attr_value(node, "position_offset", 0)), + rot_dim=int(attr_value(node, "rot_dim", 0)), + ) if target == "cactus.glu": require_len(node, inputs, 1) @@ -1434,6 +1812,15 @@ def tensor_dtype(node: IRModels.Node) -> str | None: if dtype is not None: return str(dtype) + if isinstance(node.tensor_output_meta, list) and node.tensor_output_meta: + first_meta = node.tensor_output_meta[0] + + if isinstance(first_meta, dict): + dtype = first_meta.get("dtype") + + if dtype is not None: + return str(dtype) + return None @@ -1442,6 +1829,13 @@ def cactus_precision(graph: Any, dtype: str | None) -> int: return int(getattr(graph, precision_name)) +def cast_to_precision(context: models.GenerationContext, value: Any, precision: int) -> Any: + if getattr(value, "dtype", precision) == precision: + return value + + return context.graph.precision_cast(value, precision) + + def graph_input_shape(node: IRModels.Node) -> tuple[tuple[int, ...], tuple[bool, ...]]: shape = meta_shape(node) @@ -1482,6 +1876,18 @@ def meta_shape(node: IRModels.Node) -> tuple[Any, ...]: if isinstance(shape, tuple): return shape + if isinstance(node.tensor_output_meta, list) and node.tensor_output_meta: + first_meta = node.tensor_output_meta[0] + + if isinstance(first_meta, dict): + shape = first_meta.get("shape") + + if isinstance(shape, list): + return tuple(shape) + + if isinstance(shape, tuple): + return shape + return () @@ -1494,6 +1900,30 @@ def output_shape(node: IRModels.Node) -> tuple[int, ...]: raise UnsupportedLoweringError(f"{node.name}: missing concrete output shape") +def concrete_dim(dim: Any) -> int | None: + if isinstance(dim, int) and dim >= 0: + return dim + + if isinstance(dim, str) and dim.isdigit(): + return int(dim) + + return None + + +def concrete_shape(shape: tuple[Any, ...]) -> tuple[int, ...] | None: + dims = tuple(concrete_dim(dim) for dim in shape) + + if any(dim is None for dim in dims): + return None + + return tuple(int(dim) for dim in dims) + + +def reduction_dropped_shape(shape: tuple[Any, ...], axis: int) -> tuple[Any, ...]: + dropped = tuple(dim for index, dim in enumerate(shape) if index != axis) + return dropped or (1,) + + def shape_matches_tensor(node: IRModels.Node, expected_shape: tuple[int, ...]) -> bool: actual_shape = meta_shape(node) return tuple(actual_shape) == tuple(expected_shape) @@ -1537,7 +1967,7 @@ def shape_attr(node: IRModels.Node) -> tuple[int, ...]: def axis_attr(node: IRModels.Node, default: int | None = None) -> int | None: - axis = node.attrs.get("axis", node.attrs.get("dim", default)) + axis = node.attrs.get("axis", node.attrs.get("dim", node.attrs.get("arg_1", default))) if isinstance(axis, list): if len(axis) != 1: @@ -1562,6 +1992,11 @@ def scalar_attr(node: IRModels.Node, name: str) -> Any | None: return value +def attr_value(node: IRModels.Node, name: str, default: Any) -> Any: + value = node.attrs.get(name) + return default if value is None else value + + def epsilon_attr(node: IRModels.Node) -> float: return float(node.attrs.get("epsilon", node.attrs.get("eps", 1e-5))) @@ -1591,6 +2026,24 @@ def first_int(value: Any, default: int) -> int: return int(value) +def tuple_int_values(value: Any, default: int) -> tuple[int, int]: + if value is None: + return (default, default) + + if isinstance(value, (list, tuple)): + if not value: + return (default, default) + + if len(value) == 1: + item = int(value[0]) + return (item, item) + + return (int(value[0]), int(value[1])) + + item = int(value) + return (item, item) + + def tuple_ints(value: Any) -> tuple[int, ...]: if not isinstance(value, (list, tuple)): raise TypeError(f"Expected list/tuple of ints, got {type(value).__name__}") @@ -1698,6 +2151,8 @@ def dump_result_manifest(result: models.GenerationResult, output_path: str | Pat { "component_paths": {name: str(path_) for name, path_ in result.component_paths.items()}, "component_manifest_paths": {name: str(path_) for name, path_ in result.component_manifest_paths.items()}, + "engine_manifest_path": str(result.engine_manifest_path) if result.engine_manifest_path is not None else None, + "runtime_plan_path": str(result.runtime_plan_path) if result.runtime_plan_path is not None else None, "unsupported_nodes": list(result.unsupported_nodes), "warnings": list(result.warnings), "ok": result.ok, diff --git a/python/cactus/transpiler/Generator/models.py b/python/cactus/transpiler/Generator/models.py index 9267d56eb..fbfd36098 100644 --- a/python/cactus/transpiler/Generator/models.py +++ b/python/cactus/transpiler/Generator/models.py @@ -204,6 +204,8 @@ def to_dict(self) -> dict[str, Any]: class GenerationResult: component_paths: dict[str, Path] = field(default_factory=dict) component_manifest_paths: dict[str, Path] = field(default_factory=dict) + engine_manifest_path: Path | None = None + runtime_plan_path: Path | None = None unsupported_nodes: tuple[str, ...] = () warnings: tuple[str, ...] = () @@ -508,6 +510,16 @@ def weight_name_variants(name: str) -> tuple[str, ...]: current = current[len("model."):] variants.append(current) + if name.endswith("lm_head.weight"): + variants.extend( + ( + "model.language_model.embed_tokens.weight", + "language_model.embed_tokens.weight", + "model.embed_tokens.weight", + "embed_tokens.weight", + ) + ) + if name.startswith("_orig_mod."): variants.append(name[len("_orig_mod."):]) diff --git a/python/cactus/transpiler/IR/constants.py b/python/cactus/transpiler/IR/constants.py index 29a15c903..11b59ec5e 100644 --- a/python/cactus/transpiler/IR/constants.py +++ b/python/cactus/transpiler/IR/constants.py @@ -36,6 +36,7 @@ "aten.amax.default": {1: "dim", 2: "keepdim"}, "aten.min.dim": {1: "dim", 2: "keepdim"}, "aten.amin.default": {1: "dim", 2: "keepdim"}, + "aten.any.dim": {1: "dim", 2: "keepdim"}, "aten.softmax.int": {1: "dim", 2: "dtype"}, "aten._softmax.default": {1: "dim", 2: "half_to_float"}, "aten.topk.default": {1: "k", 2: "dim", 3: "largest", 4: "sorted"}, @@ -43,6 +44,7 @@ "aten.view.default": {1: "shape"}, "aten.reshape.default": {1: "shape"}, "aten.flatten.using_ints": {1: "start_dim", 2: "end_dim"}, + "aten.unfold.default": {1: "dim", 2: "size", 3: "step"}, "aten.unsqueeze.default": {1: "dim"}, "aten.expand.default": {1: "shape", 2: "implicit"}, "aten.t.default": {}, @@ -53,6 +55,7 @@ "aten.narrow.default": {1: "dim", 2: "start", 3: "length"}, "aten.index.Tensor": {1: "indices"}, "aten.index_put.default": {3: "accumulate"}, + "aten.constant_pad_nd.default": {1: "pad", 2: "value"}, "aten.cat.default": {1: "dim"}, "aten.embedding.default": {2: "padding_idx", 3: "scale_grad_by_freq", 4: "sparse"}, "aten.gather.default": {1: "dim", 2: "index"}, diff --git a/python/cactus/transpiler/main.py b/python/cactus/transpiler/main.py index 1a3187734..cacea140a 100644 --- a/python/cactus/transpiler/main.py +++ b/python/cactus/transpiler/main.py @@ -1,33 +1,83 @@ import sys +import os +import shutil from pathlib import Path sys.path.append(str(Path(__file__).resolve().parents[1])) from transpiler.Converter import constants from transpiler.Converter.convert import convert +from transpiler.Converter.models import LayerMap +from transpiler.Generator.generate import generate_bundle +from transpiler.ModelProfiles.profiles import GEMMA4_E2B_PROFILE + + +def materialize_runtime_bundle_files(source_dir: Path, bundle_dir: Path) -> None: + for source in source_dir.iterdir(): + if not source.is_file(): + continue + + target = bundle_dir / source.name + if target.exists() or target.is_symlink(): + try: + if target.samefile(source): + continue + except OSError: + pass + target.unlink() + + try: + os.symlink(source, target) + except OSError: + shutil.copy2(source, target) def main() -> None: model_id = "google/gemma-4-E2B" input_modalities = ("text", "vision", "audio") output_dir = constants.CONVERTER_JSON_DIR + bundle_dir = Path(__file__).resolve().parent / "Generator" / "graphs" / "gemma4_e2b_bundle" output_dir.mkdir(parents=True, exist_ok=True) + bundle_dir.mkdir(parents=True, exist_ok=True) + + prefill_path = output_dir / "output_prefill_with_cache.json" + prefill_simplified_path = output_dir / "output_prefill_with_cache_simplified.json" + decode_path = output_dir / "output_decode_with_cache.json" + decode_simplified_path = output_dir / "output_decode_with_cache_simplified.json" convert( model_id=model_id, input_modalities=input_modalities, - output_path=str(output_dir / "output_prefill_with_cache.json"), - simplified_output_path=str(output_dir / "output_prefill_with_cache_simplified.json"), + output_path=str(prefill_path), + simplified_output_path=str(prefill_simplified_path), inference_mode=constants.PREFILL_WITH_CACHE_MODE, ) convert( model_id=model_id, input_modalities=input_modalities, - output_path=str(output_dir / "output_decode_with_cache.json"), - simplified_output_path=str(output_dir / "output_decode_with_cache_simplified.json"), + output_path=str(decode_path), + simplified_output_path=str(decode_simplified_path), inference_mode=constants.DECODE_WITH_CACHE_MODE, ) + prefill = LayerMap.model_validate_json(prefill_simplified_path.read_text(encoding="utf-8")) + decode = LayerMap.model_validate_json(decode_simplified_path.read_text(encoding="utf-8")) + weights_arg = sys.argv[1] if len(sys.argv) > 1 else os.environ.get("CACTUS_TRANSPILER_WEIGHTS_DIR") + weights_dir = Path(weights_arg).expanduser() if weights_arg else None + + if weights_dir is not None: + materialize_runtime_bundle_files(weights_dir, bundle_dir) + + result = generate_bundle( + {"prefill": prefill, "decode": decode}, + bundle_dir, + model_profile=GEMMA4_E2B_PROFILE, + weights_dir=weights_dir, + strict=True, + ) + print(f"Wrote component bundle to {bundle_dir}") + print(f"Wrote runtime plan to {result.runtime_plan_path}") + if __name__ == "__main__": main() From aa06bf6587480395c333260e2b22b520e2961e12 Mon Sep 17 00:00:00 2001 From: Parkirat Date: Wed, 29 Jul 2026 15:28:44 -0400 Subject: [PATCH 67/84] New run operations with runtime plans --- .gitignore | 5 + cactus-graph/cactus_graph.h | 82 +++- cactus-graph/src/builder.cpp | 195 +++++++++- cactus-graph/src/execute.cpp | 85 ++++- cactus-graph/src/graph_ffi.cpp | 251 +++++++++++++ cactus-graph/src/io.cpp | 9 +- cactus-graph/src/ops_conv.cpp | 23 ++ cactus-graph/src/ops_math.cpp | 184 ++++++++- cactus-graph/src/ops_tensor.cpp | 349 +++++++++++++++++- cactus-graph/src/param_io.cpp | 13 +- python/tests/test_cli_run.py | 38 ++ .../tests/test_generator_component_split.py | 115 ++++++ python/tests/test_graph.py | 67 ++++ python/tests/test_runtime_plan.py | 96 +++++ python/tests/test_server_live.py | 2 +- 15 files changed, 1493 insertions(+), 21 deletions(-) create mode 100644 python/tests/test_generator_component_split.py create mode 100644 python/tests/test_runtime_plan.py diff --git a/.gitignore b/.gitignore index 743d97d84..59a9f0a9a 100644 --- a/.gitignore +++ b/.gitignore @@ -4,6 +4,9 @@ *.dylib *.bin *.cg +*.json +*.cactus +*.weights !cactus-engine/libs/curl/**/*.a !android/mbedtls/**/*.a @@ -28,3 +31,5 @@ python/cactus/bin/ python/cactus/bindings/lib/ site/ site_docs/ + +TODO.md diff --git a/cactus-graph/cactus_graph.h b/cactus-graph/cactus_graph.h index 55aa7076b..7641d2321 100644 --- a/cactus-graph/cactus_graph.h +++ b/cactus-graph/cactus_graph.h @@ -118,7 +118,27 @@ enum class OpType { SCALAR_NOT_EQUAL, RECURRENT_CACHE_STATE, RECURRENT_CACHE_WRITE, - CONV_CACHE_INITIALIZE + CONV_CACHE_INITIALIZE, + EQUAL, + LESS, + LESS_EQUAL, + GREATER, + GREATER_EQUAL, + SCALAR_EQUAL, + SCALAR_LESS, + SCALAR_LESS_EQUAL, + SCALAR_GREATER, + SCALAR_GREATER_EQUAL, + LOGICAL_NOT, + BITWISE_AND, + BITWISE_OR, + BITWISE_NOT, + WHERE, + UNFOLD, + PAD, + SCALAR_FLOOR_DIVIDE, + STRIDED_SLICE, + EXPAND }; struct PrecisionTraits { @@ -481,18 +501,34 @@ class CactusGraph { size_t multiply(size_t input1, size_t input2); size_t divide(size_t input1, size_t input2); size_t not_equal(size_t input1, size_t input2); + size_t equal(size_t input1, size_t input2); + size_t less(size_t input1, size_t input2); + size_t less_equal(size_t input1, size_t input2); + size_t greater(size_t input1, size_t input2); + size_t greater_equal(size_t input1, size_t input2); + size_t bitwise_and(size_t input1, size_t input2); + size_t bitwise_or(size_t input1, size_t input2); + size_t where(size_t condition, size_t true_value, size_t false_value); size_t scalar_add(size_t input, float value); size_t scalar_subtract(size_t input, float value); size_t scalar_multiply(size_t input, float value); size_t scalar_divide(size_t input, float value); + size_t scalar_floor_divide(size_t input, float value); size_t scalar_not_equal(size_t input, float value); + size_t scalar_equal(size_t input, float value); + size_t scalar_less(size_t input, float value); + size_t scalar_less_equal(size_t input, float value); + size_t scalar_greater(size_t input, float value); + size_t scalar_greater_equal(size_t input, float value); size_t scalar_exp(size_t input); size_t scalar_sqrt(size_t input); size_t scalar_cos(size_t input); size_t scalar_sin(size_t input); size_t scalar_log(size_t input); + size_t logical_not(size_t input); + size_t bitwise_not(size_t input); size_t abs(size_t input); size_t pow(size_t input, float exponent); @@ -520,11 +556,15 @@ class CactusGraph { size_t reshape(size_t input, const std::vector& new_shape); size_t view(size_t input, const std::vector& new_shape); + size_t expand(size_t input, const std::vector& new_shape); size_t flatten(size_t input, int start_dim = 0, int end_dim = -1); size_t transpose(size_t input, ComputeBackend backend = ComputeBackend::CPU); size_t transposeN(size_t input, const std::vector& permutation, ComputeBackend backend = ComputeBackend::CPU); size_t slice(size_t input, int axis, size_t start, size_t length); + size_t strided_slice(size_t input, int axis, size_t start, size_t length, size_t step); size_t index(size_t input, size_t index_value, int dim); + size_t unfold(size_t input, int dimension, size_t size, size_t step); + size_t pad(size_t input, const std::vector& pads, float value = 0.0f); size_t concat(size_t input1, size_t input2, int axis = 0); size_t cat(const std::vector& inputs, int axis); @@ -603,6 +643,8 @@ class CactusGraph { size_t conv1d_k7s3(size_t input, size_t weight, size_t bias); size_t conv1d(size_t input, size_t weight, size_t stride); size_t conv1d(size_t input, size_t weight, size_t bias, size_t stride); + size_t conv1d_depthwise(size_t input, size_t weight, size_t stride = 1); + size_t conv1d_depthwise(size_t input, size_t weight, size_t bias, size_t stride); size_t conv1d_same_depthwise_k9(size_t input, size_t weight); size_t conv1d_same_depthwise_k9(size_t input, size_t weight, size_t bias); size_t conv1d_pointwise(size_t input, size_t weight); @@ -685,7 +727,7 @@ class CactusGraph { size_t top_k, const std::unordered_map& logit_bias = {}); size_t scatter_topk(size_t indices, size_t values, size_t num_classes); - size_t gather(size_t embeddings, size_t indices); + size_t gather(size_t tensor, size_t indices, int axis = 0); size_t embedding(const std::string& filename, size_t indices); size_t embedding(size_t embedding_tensor, size_t indices); size_t mmap_embeddings(const std::string& filename); @@ -895,17 +937,41 @@ CACTUS_FFI_EXPORT int cactus_graph_divide(cactus_graph_t graph, cactus_node_t a, cactus_node_t b, cactus_node_t* out); CACTUS_FFI_EXPORT int cactus_graph_not_equal(cactus_graph_t graph, cactus_node_t a, cactus_node_t b, cactus_node_t* out); +CACTUS_FFI_EXPORT int cactus_graph_equal(cactus_graph_t graph, cactus_node_t +a, cactus_node_t b, cactus_node_t* out); +CACTUS_FFI_EXPORT int cactus_graph_less(cactus_graph_t graph, cactus_node_t +a, cactus_node_t b, cactus_node_t* out); +CACTUS_FFI_EXPORT int cactus_graph_less_equal(cactus_graph_t graph, cactus_node_t +a, cactus_node_t b, cactus_node_t* out); +CACTUS_FFI_EXPORT int cactus_graph_greater(cactus_graph_t graph, cactus_node_t +a, cactus_node_t b, cactus_node_t* out); +CACTUS_FFI_EXPORT int cactus_graph_greater_equal(cactus_graph_t graph, cactus_node_t +a, cactus_node_t b, cactus_node_t* out); +CACTUS_FFI_EXPORT int cactus_graph_bitwise_and(cactus_graph_t graph, cactus_node_t +a, cactus_node_t b, cactus_node_t* out); +CACTUS_FFI_EXPORT int cactus_graph_bitwise_or(cactus_graph_t graph, cactus_node_t +a, cactus_node_t b, cactus_node_t* out); +CACTUS_FFI_EXPORT int cactus_graph_where(cactus_graph_t graph, cactus_node_t +condition, cactus_node_t true_value, cactus_node_t false_value, cactus_node_t* out); CACTUS_FFI_EXPORT int cactus_graph_scalar_add(cactus_graph_t graph, cactus_node_t x, float value, cactus_node_t* out); CACTUS_FFI_EXPORT int cactus_graph_scalar_subtract(cactus_graph_t graph, cactus_node_t x, float value, cactus_node_t* out); CACTUS_FFI_EXPORT int cactus_graph_scalar_multiply(cactus_graph_t graph, cactus_node_t x, float value, cactus_node_t* out); CACTUS_FFI_EXPORT int cactus_graph_scalar_divide(cactus_graph_t graph, cactus_node_t x, float value, cactus_node_t* out); +CACTUS_FFI_EXPORT int cactus_graph_scalar_floor_divide(cactus_graph_t graph, cactus_node_t x, float value, cactus_node_t* out); CACTUS_FFI_EXPORT int cactus_graph_scalar_not_equal(cactus_graph_t graph, cactus_node_t x, float value, cactus_node_t* out); +CACTUS_FFI_EXPORT int cactus_graph_scalar_equal(cactus_graph_t graph, cactus_node_t x, float value, cactus_node_t* out); +CACTUS_FFI_EXPORT int cactus_graph_scalar_less(cactus_graph_t graph, cactus_node_t x, float value, cactus_node_t* out); +CACTUS_FFI_EXPORT int cactus_graph_scalar_less_equal(cactus_graph_t graph, cactus_node_t x, float value, cactus_node_t* out); +CACTUS_FFI_EXPORT int cactus_graph_scalar_greater(cactus_graph_t graph, cactus_node_t x, float value, cactus_node_t* out); +CACTUS_FFI_EXPORT int cactus_graph_scalar_greater_equal(cactus_graph_t graph, cactus_node_t x, float value, cactus_node_t* out); CACTUS_FFI_EXPORT int cactus_graph_scalar_exp(cactus_graph_t graph, cactus_node_t x, cactus_node_t* out); CACTUS_FFI_EXPORT int cactus_graph_scalar_sqrt(cactus_graph_t graph, cactus_node_t x, cactus_node_t* out); CACTUS_FFI_EXPORT int cactus_graph_scalar_cos(cactus_graph_t graph, cactus_node_t x, cactus_node_t* out); CACTUS_FFI_EXPORT int cactus_graph_scalar_sin(cactus_graph_t graph, cactus_node_t x, cactus_node_t* out); CACTUS_FFI_EXPORT int cactus_graph_scalar_log(cactus_graph_t graph, cactus_node_t x, cactus_node_t* out); +CACTUS_FFI_EXPORT int cactus_graph_logical_not(cactus_graph_t graph, cactus_node_t x, cactus_node_t* out); +CACTUS_FFI_EXPORT int cactus_graph_bitwise_not(cactus_graph_t graph, cactus_node_t x, cactus_node_t* out); CACTUS_FFI_EXPORT int cactus_graph_abs(cactus_graph_t graph, cactus_node_t x, cactus_node_t* out); @@ -920,14 +986,22 @@ CACTUS_FFI_EXPORT int cactus_graph_flatten( cactus_node_t* out); CACTUS_FFI_EXPORT int cactus_graph_reshape( cactus_graph_t graph, cactus_node_t x, const size_t* shape, size_t rank, cactus_node_t* out); +CACTUS_FFI_EXPORT int cactus_graph_expand( + cactus_graph_t graph, cactus_node_t x, const size_t* shape, size_t rank, cactus_node_t* out); CACTUS_FFI_EXPORT int cactus_graph_transpose( cactus_graph_t graph, cactus_node_t x, int32_t backend, cactus_node_t* out); CACTUS_FFI_EXPORT int cactus_graph_transpose_n( cactus_graph_t graph, cactus_node_t x, const size_t* permutation, size_t rank, int32_t backend, cactus_node_t* out); CACTUS_FFI_EXPORT int cactus_graph_slice( cactus_graph_t graph, cactus_node_t x, int32_t axis, size_t start, size_t length, cactus_node_t* out); +CACTUS_FFI_EXPORT int cactus_graph_strided_slice( + cactus_graph_t graph, cactus_node_t x, int32_t axis, size_t start, size_t length, size_t step, cactus_node_t* out); CACTUS_FFI_EXPORT int cactus_graph_index( cactus_graph_t graph, cactus_node_t x, size_t index_value, int32_t dim, cactus_node_t* out); +CACTUS_FFI_EXPORT int cactus_graph_unfold( + cactus_graph_t graph, cactus_node_t x, int32_t dimension, size_t size, size_t step, cactus_node_t* out); +CACTUS_FFI_EXPORT int cactus_graph_pad( + cactus_graph_t graph, cactus_node_t x, const size_t* pads, size_t pad_count, float value, cactus_node_t* out); CACTUS_FFI_EXPORT int cactus_graph_sum(cactus_graph_t graph, cactus_node_t x, int32_t axis, cactus_node_t* out); CACTUS_FFI_EXPORT int cactus_graph_mean(cactus_graph_t graph, cactus_node_t x, int32_t axis, cactus_node_t* out); CACTUS_FFI_EXPORT int cactus_graph_variance(cactus_graph_t graph, cactus_node_t x, int32_t axis, cactus_node_t* out); @@ -946,6 +1020,8 @@ CACTUS_FFI_EXPORT int cactus_graph_matmul( cactus_graph_t graph, cactus_node_t a, cactus_node_t b, bool pretransposed_rhs, int32_t backend, cactus_node_t* out); CACTUS_FFI_EXPORT int cactus_graph_gather( cactus_graph_t graph, cactus_node_t tensor, cactus_node_t indices, cactus_node_t* out); +CACTUS_FFI_EXPORT int cactus_graph_gather_dim( + cactus_graph_t graph, cactus_node_t tensor, cactus_node_t indices, int32_t axis, cactus_node_t* out); CACTUS_FFI_EXPORT int cactus_graph_embedding_from_tensor( cactus_graph_t graph, cactus_node_t embedding_tensor, cactus_node_t indices, cactus_node_t* out); CACTUS_FFI_EXPORT int cactus_graph_embedding_from_file( @@ -1043,6 +1119,8 @@ CACTUS_FFI_EXPORT int cactus_graph_conv1d_k7s3( cactus_graph_t graph, cactus_node_t input, cactus_node_t weight, cactus_node_t bias, cactus_node_t* out); CACTUS_FFI_EXPORT int cactus_graph_conv1d( cactus_graph_t graph, cactus_node_t input, cactus_node_t weight, bool has_bias, cactus_node_t bias, size_t stride, cactus_node_t* out); +CACTUS_FFI_EXPORT int cactus_graph_conv1d_depthwise( + cactus_graph_t graph, cactus_node_t input, cactus_node_t weight, bool has_bias, cactus_node_t bias, size_t stride, cactus_node_t* out); CACTUS_FFI_EXPORT int cactus_graph_conv1d_same_depthwise_k9( cactus_graph_t graph, cactus_node_t input, cactus_node_t weight, bool has_bias, cactus_node_t bias, cactus_node_t* out); CACTUS_FFI_EXPORT int cactus_graph_conv1d_pointwise( diff --git a/cactus-graph/src/builder.cpp b/cactus-graph/src/builder.cpp index 88778a168..ddf2d5b19 100644 --- a/cactus-graph/src/builder.cpp +++ b/cactus-graph/src/builder.cpp @@ -39,6 +39,25 @@ size_t CactusGraph::subtract(size_t a, size_t b) { return binary_broadcast_op(Op size_t CactusGraph::multiply(size_t a, size_t b) { return binary_broadcast_op(OpType::MULTIPLY, a, b); } size_t CactusGraph::divide(size_t a, size_t b) { return binary_broadcast_op(OpType::DIVIDE, a, b); } size_t CactusGraph::not_equal(size_t a, size_t b) { return binary_broadcast_op(OpType::NOT_EQUAL, a, b); } +size_t CactusGraph::equal(size_t a, size_t b) { return binary_broadcast_op(OpType::EQUAL, a, b); } +size_t CactusGraph::less(size_t a, size_t b) { return binary_broadcast_op(OpType::LESS, a, b); } +size_t CactusGraph::less_equal(size_t a, size_t b) { return binary_broadcast_op(OpType::LESS_EQUAL, a, b); } +size_t CactusGraph::greater(size_t a, size_t b) { return binary_broadcast_op(OpType::GREATER, a, b); } +size_t CactusGraph::greater_equal(size_t a, size_t b) { return binary_broadcast_op(OpType::GREATER_EQUAL, a, b); } +size_t CactusGraph::bitwise_and(size_t a, size_t b) { return binary_broadcast_op(OpType::BITWISE_AND, a, b); } +size_t CactusGraph::bitwise_or(size_t a, size_t b) { return binary_broadcast_op(OpType::BITWISE_OR, a, b); } + +size_t CactusGraph::where(size_t condition, size_t true_value, size_t false_value) { + BroadcastInfo cond_true = BroadcastInfo::compute( + get_output_buffer(condition).shape, + get_output_buffer(true_value).shape); + BroadcastInfo all_inputs = BroadcastInfo::compute( + cond_true.output_shape, + get_output_buffer(false_value).shape); + OpParams params; + params.output_precision = get_output_buffer(true_value).precision; + return add_node(OpType::WHERE, {condition, true_value, false_value}, all_inputs.output_shape, params); +} size_t CactusGraph::abs(size_t input) { const auto& input_buffer = get_output_buffer(input); @@ -176,6 +195,29 @@ size_t CactusGraph::reshape(size_t input, const std::vector& new_shape) return add_node(OpType::RESHAPE, {input}, new_shape, params); } +size_t CactusGraph::expand(size_t input, const std::vector& new_shape) { + const auto& input_buffer = get_output_buffer(input); + const auto& input_shape = input_buffer.shape; + if (new_shape.size() < input_shape.size()) { + throw std::runtime_error("expand target rank must be >= input rank"); + } + + const size_t offset = new_shape.size() - input_shape.size(); + for (size_t i = 0; i < new_shape.size(); ++i) { + size_t input_dim = i < offset ? 1 : input_shape[i - offset]; + size_t target_dim = new_shape[i]; + + if (input_dim != target_dim && input_dim != 1) { + throw std::runtime_error("expand target shape is not broadcast-compatible"); + } + } + + OpParams params; + params.new_shape = new_shape; + params.output_precision = input_buffer.precision; + return add_node(OpType::EXPAND, {input}, new_shape, params); +} + size_t CactusGraph::index(size_t input, size_t index_value, int dim) { const auto& input_buffer = get_output_buffer(input); const auto& shape = input_buffer.shape; @@ -620,6 +662,32 @@ size_t CactusGraph::conv1d(size_t input, size_t weight, size_t bias, size_t stri return add_node(OpType::CONV1D, {input, weight, bias}, {xin.shape[0], w.shape[0], L_out}, params); } +size_t CactusGraph::conv1d_depthwise(size_t input, size_t weight, size_t stride) { + const auto& xin = get_output_buffer(input); + const auto& w = get_output_buffer(weight); + if (xin.shape.size() != 3) throw std::runtime_error("conv1d_depthwise expects N,C,L"); + if (w.shape.size() != 3 || w.shape[1] != 1) throw std::runtime_error("conv1d_depthwise weight expects [C,1,K]"); + if (w.shape[0] != xin.shape[1]) throw std::runtime_error("conv1d_depthwise channel mismatch"); + size_t L_out = (xin.shape[2] - w.shape[2]) / stride + 1; + OpParams params{.output_precision = xin.precision, .stride = stride}; + params.num_groups = xin.shape[1]; + return add_node(OpType::CONV1D, {input, weight}, {xin.shape[0], w.shape[0], L_out}, params); +} + +size_t CactusGraph::conv1d_depthwise(size_t input, size_t weight, size_t bias, size_t stride) { + const auto& b = get_output_buffer(bias); + const auto& w = get_output_buffer(weight); + if (b.total_size != w.shape[0]) throw std::runtime_error("conv1d_depthwise bias size mismatch"); + const auto& xin = get_output_buffer(input); + if (xin.shape.size() != 3) throw std::runtime_error("conv1d_depthwise expects N,C,L"); + if (w.shape.size() != 3 || w.shape[1] != 1) throw std::runtime_error("conv1d_depthwise weight expects [C,1,K]"); + if (w.shape[0] != xin.shape[1]) throw std::runtime_error("conv1d_depthwise channel mismatch"); + size_t L_out = (xin.shape[2] - w.shape[2]) / stride + 1; + OpParams params{.output_precision = xin.precision, .stride = stride}; + params.num_groups = xin.shape[1]; + return add_node(OpType::CONV1D, {input, weight, bias}, {xin.shape[0], w.shape[0], L_out}, params); +} + size_t CactusGraph::conv1d_same_depthwise_k9(size_t input, size_t weight) { const auto& xin = get_output_buffer(input); const auto& w = get_output_buffer(weight); @@ -1043,7 +1111,25 @@ size_t CactusGraph::scalar_add(size_t input, float value) { return scalar_val_op size_t CactusGraph::scalar_subtract(size_t input, float value) { return scalar_val_op(*this, OpType::SCALAR_SUBTRACT, input, value); } size_t CactusGraph::scalar_multiply(size_t input, float value) { return scalar_val_op(*this, OpType::SCALAR_MULTIPLY, input, value); } size_t CactusGraph::scalar_divide(size_t input, float value) { return scalar_val_op(*this, OpType::SCALAR_DIVIDE, input, value); } +size_t CactusGraph::scalar_floor_divide(size_t input, float value) { return scalar_val_op(*this, OpType::SCALAR_FLOOR_DIVIDE, input, value); } size_t CactusGraph::scalar_not_equal(size_t input, float value) { return scalar_val_op(*this, OpType::SCALAR_NOT_EQUAL, input, value); } +size_t CactusGraph::scalar_equal(size_t input, float value) { return scalar_val_op(*this, OpType::SCALAR_EQUAL, input, value); } +size_t CactusGraph::scalar_less(size_t input, float value) { return scalar_val_op(*this, OpType::SCALAR_LESS, input, value); } +size_t CactusGraph::scalar_less_equal(size_t input, float value) { return scalar_val_op(*this, OpType::SCALAR_LESS_EQUAL, input, value); } +size_t CactusGraph::scalar_greater(size_t input, float value) { return scalar_val_op(*this, OpType::SCALAR_GREATER, input, value); } +size_t CactusGraph::scalar_greater_equal(size_t input, float value) { return scalar_val_op(*this, OpType::SCALAR_GREATER_EQUAL, input, value); } + +size_t CactusGraph::logical_not(size_t input) { + const auto& input_buffer = get_output_buffer(input); + OpParams params{.output_precision = input_buffer.precision}; + return add_node(OpType::LOGICAL_NOT, {input}, input_buffer.shape, params); +} + +size_t CactusGraph::bitwise_not(size_t input) { + const auto& input_buffer = get_output_buffer(input); + OpParams params{.output_precision = input_buffer.precision}; + return add_node(OpType::BITWISE_NOT, {input}, input_buffer.shape, params); +} size_t CactusGraph::scalar_exp(size_t input) { return add_node(OpType::SCALAR_EXP, {input}, {}); @@ -1065,6 +1151,55 @@ size_t CactusGraph::scalar_log(size_t input) { return add_node(OpType::SCALAR_LOG, {input}, {}); } +size_t CactusGraph::unfold(size_t input, int dimension, size_t size, size_t step) { + const auto& input_buffer = get_output_buffer(input); + const auto& shape = input_buffer.shape; + if (shape.empty()) throw std::runtime_error("unfold requires a ranked input tensor"); + if (step == 0) throw std::runtime_error("unfold step must be greater than zero"); + + int actual_dim = dimension; + if (actual_dim < 0) actual_dim += static_cast(shape.size()); + if (actual_dim < 0 || static_cast(actual_dim) >= shape.size()) { + throw std::runtime_error("unfold dimension out of range"); + } + + const size_t axis = static_cast(actual_dim); + if (size == 0 || size > shape[axis]) { + throw std::runtime_error("unfold size must be in range for the selected dimension"); + } + + std::vector output_shape = shape; + output_shape[axis] = ((shape[axis] - size) / step) + 1; + output_shape.push_back(size); + + OpParams params; + params.axis = actual_dim; + params.kernel_size = size; + params.stride = step; + params.output_precision = input_buffer.precision; + return add_node(OpType::UNFOLD, {input}, output_shape, params); +} + +size_t CactusGraph::pad(size_t input, const std::vector& pads, float value) { + const auto& input_buffer = get_output_buffer(input); + const auto& shape = input_buffer.shape; + if (shape.empty()) throw std::runtime_error("pad requires a ranked input tensor"); + if (pads.size() % 2 != 0) throw std::runtime_error("pad requires left/right pad pairs"); + if ((pads.size() / 2) > shape.size()) throw std::runtime_error("pad rank exceeds input rank"); + + std::vector output_shape = shape; + for (size_t pair = 0; pair < pads.size(); pair += 2) { + size_t axis = shape.size() - 1 - (pair / 2); + output_shape[axis] += pads[pair] + pads[pair + 1]; + } + + OpParams params; + params.new_shape = pads; + params.scalar = value; + params.output_precision = input_buffer.precision; + return add_node(OpType::PAD, {input}, output_shape, params); +} + size_t CactusGraph::relu(size_t input) { return add_node(OpType::RELU, {input}, {}); } @@ -1125,7 +1260,7 @@ size_t CactusGraph::rope_gptj(size_t input, float theta, size_t position_offset, return add_node(OpType::ROPE_GPTJ, {input}, {}, params); } -size_t CactusGraph::gather(size_t tensor, size_t indices) { +size_t CactusGraph::gather(size_t tensor, size_t indices, int axis) { const auto& tensor_buffer = get_output_buffer(tensor); const auto& idx_shape = get_output_buffer(indices).shape; @@ -1133,12 +1268,25 @@ size_t CactusGraph::gather(size_t tensor, size_t indices) { throw std::runtime_error("Cannot gather from scalar tensor"); } - std::vector output_shape = idx_shape; - for (size_t i = 1; i < tensor_buffer.shape.size(); i++) { - output_shape.push_back(tensor_buffer.shape[i]); + int actual_axis = axis; + if (actual_axis < 0) actual_axis += static_cast(tensor_buffer.shape.size()); + if (actual_axis < 0 || static_cast(actual_axis) >= tensor_buffer.shape.size()) { + throw std::runtime_error("gather axis out of range"); + } + + std::vector output_shape; + + if (idx_shape.size() == tensor_buffer.shape.size()) { + output_shape = idx_shape; + } else { + output_shape = idx_shape; + for (size_t i = 1; i < tensor_buffer.shape.size(); i++) { + output_shape.push_back(tensor_buffer.shape[i]); + } } OpParams params; + params.axis = actual_axis; params.output_precision = tensor_buffer.precision; return add_node(OpType::GATHER, {tensor, indices}, output_shape, params); @@ -1169,6 +1317,45 @@ size_t CactusGraph::slice(size_t input, int axis, size_t start, size_t length) { return add_node(OpType::SLICE, {input}, output_shape, params); } +size_t CactusGraph::strided_slice(size_t input, int axis, size_t start, size_t length, size_t step) { + const auto& input_buffer = get_output_buffer(input); + if (input_buffer.shape.empty()) { + throw std::runtime_error("Cannot strided_slice a scalar tensor"); + } + if (step == 0) { + throw std::runtime_error("strided_slice step must be greater than zero"); + } + + int actual_axis = axis; + if (actual_axis < 0) actual_axis += static_cast(input_buffer.shape.size()); + if (actual_axis < 0 || static_cast(actual_axis) >= input_buffer.shape.size()) { + throw std::runtime_error("strided_slice axis out of range"); + } + + size_t axis_index = static_cast(actual_axis); + size_t axis_size = input_buffer.shape[axis_index]; + + if (length == 0) { + length = start >= axis_size ? 0 : ((axis_size - start + step - 1) / step); + } + + if (length > 0 && start + (length - 1) * step >= axis_size) { + throw std::runtime_error("strided_slice range extends beyond axis size"); + } + + std::vector output_shape = input_buffer.shape; + output_shape[axis_index] = length; + + OpParams params; + params.axis = actual_axis; + params.slice_start = start; + params.slice_length = length; + params.stride = step; + params.output_precision = input_buffer.precision; + + return add_node(OpType::STRIDED_SLICE, {input}, output_shape, params); +} + size_t CactusGraph::embedding(size_t embedding_tensor, size_t indices) { const auto& emb_buffer = get_output_buffer(embedding_tensor); const auto& idx_shape = get_output_buffer(indices).shape; diff --git a/cactus-graph/src/execute.cpp b/cactus-graph/src/execute.cpp index fb1481cbb..ca79893ae 100644 --- a/cactus-graph/src/execute.cpp +++ b/cactus-graph/src/execute.cpp @@ -20,6 +20,7 @@ using ComputeFn = void(*)(GraphNode&, const nodes_vector&, const node_index_map_ DECLARE_COMPUTE(compute_binary_op_node); DECLARE_COMPUTE(compute_unary_op_node); +DECLARE_COMPUTE(compute_where_node); DECLARE_COMPUTE(compute_activation_node); DECLARE_COMPUTE(compute_reduce_node); DECLARE_COMPUTE(compute_reshape_node); @@ -64,6 +65,7 @@ DECLARE_COMPUTE(compute_embedding_node); DECLARE_COMPUTE(compute_concat_node); DECLARE_COMPUTE(compute_cat_node); DECLARE_COMPUTE(compute_index_node); +DECLARE_COMPUTE(compute_unfold_node); DECLARE_COMPUTE(compute_bilinear_interpolation_node); DECLARE_COMPUTE(compute_sample_node); DECLARE_COMPUTE(compute_topk_node); @@ -84,10 +86,13 @@ DECLARE_COMPUTE(compute_rfft_node); DECLARE_COMPUTE(compute_irfft_node); DECLARE_COMPUTE(compute_mel_filter_bank_node); DECLARE_COMPUTE(compute_spectrogram_node); +DECLARE_COMPUTE(compute_pad_node); +DECLARE_COMPUTE(compute_strided_slice_node); +DECLARE_COMPUTE(compute_expand_node); extern void shrink_thread_local_buffers(); #undef DECLARE_COMPUTE -static constexpr int OP_TYPE_COUNT = static_cast(OpType::CONV_CACHE_INITIALIZE) + 1; +static constexpr int OP_TYPE_COUNT = static_cast(OpType::EXPAND) + 1; static_assert(OP_TYPE_COUNT <= 256, "OpType dispatch table overflow"); static ComputeFn dispatch_flat[OP_TYPE_COUNT] = {}; @@ -98,11 +103,27 @@ static bool init_dispatch() { dispatch_flat[static_cast(OpType::MULTIPLY)] = compute_binary_op_node; dispatch_flat[static_cast(OpType::DIVIDE)] = compute_binary_op_node; dispatch_flat[static_cast(OpType::NOT_EQUAL)] = compute_binary_op_node; + dispatch_flat[static_cast(OpType::EQUAL)] = compute_binary_op_node; + dispatch_flat[static_cast(OpType::LESS)] = compute_binary_op_node; + dispatch_flat[static_cast(OpType::LESS_EQUAL)] = compute_binary_op_node; + dispatch_flat[static_cast(OpType::GREATER)] = compute_binary_op_node; + dispatch_flat[static_cast(OpType::GREATER_EQUAL)] = compute_binary_op_node; + dispatch_flat[static_cast(OpType::BITWISE_AND)] = compute_binary_op_node; + dispatch_flat[static_cast(OpType::BITWISE_OR)] = compute_binary_op_node; dispatch_flat[static_cast(OpType::SCALAR_ADD)] = compute_unary_op_node; dispatch_flat[static_cast(OpType::SCALAR_SUBTRACT)] = compute_unary_op_node; dispatch_flat[static_cast(OpType::SCALAR_MULTIPLY)] = compute_unary_op_node; dispatch_flat[static_cast(OpType::SCALAR_DIVIDE)] = compute_unary_op_node; + dispatch_flat[static_cast(OpType::SCALAR_FLOOR_DIVIDE)] = compute_unary_op_node; dispatch_flat[static_cast(OpType::SCALAR_NOT_EQUAL)] = compute_unary_op_node; + dispatch_flat[static_cast(OpType::SCALAR_EQUAL)] = compute_unary_op_node; + dispatch_flat[static_cast(OpType::SCALAR_LESS)] = compute_unary_op_node; + dispatch_flat[static_cast(OpType::SCALAR_LESS_EQUAL)] = compute_unary_op_node; + dispatch_flat[static_cast(OpType::SCALAR_GREATER)] = compute_unary_op_node; + dispatch_flat[static_cast(OpType::SCALAR_GREATER_EQUAL)] = compute_unary_op_node; + dispatch_flat[static_cast(OpType::LOGICAL_NOT)] = compute_unary_op_node; + dispatch_flat[static_cast(OpType::BITWISE_NOT)] = compute_unary_op_node; + dispatch_flat[static_cast(OpType::WHERE)] = compute_where_node; dispatch_flat[static_cast(OpType::SCALAR_EXP)] = compute_unary_op_node; dispatch_flat[static_cast(OpType::SCALAR_SQRT)] = compute_unary_op_node; dispatch_flat[static_cast(OpType::SCALAR_COS)] = compute_unary_op_node; @@ -127,6 +148,7 @@ static bool init_dispatch() { dispatch_flat[static_cast(OpType::FLATTEN)] = compute_reshape_node; dispatch_flat[static_cast(OpType::VIEW)] = compute_reshape_node; dispatch_flat[static_cast(OpType::RESHAPE)] = compute_reshape_node; + dispatch_flat[static_cast(OpType::EXPAND)] = compute_expand_node; dispatch_flat[static_cast(OpType::PRECISION_CAST)] = compute_precision_cast_node; dispatch_flat[static_cast(OpType::MATMUL)] = compute_matmul_node; dispatch_flat[static_cast(OpType::RMS_NORM)] = compute_rms_norm_node; @@ -153,10 +175,13 @@ static bool init_dispatch() { dispatch_flat[static_cast(OpType::TRANSPOSE)] = compute_transpose_node; dispatch_flat[static_cast(OpType::GATHER)] = compute_gather_node; dispatch_flat[static_cast(OpType::SLICE)] = compute_slice_node; + dispatch_flat[static_cast(OpType::STRIDED_SLICE)] = compute_strided_slice_node; dispatch_flat[static_cast(OpType::EMBEDDING)] = compute_embedding_node; dispatch_flat[static_cast(OpType::CONCAT)] = compute_concat_node; dispatch_flat[static_cast(OpType::CAT)] = compute_cat_node; dispatch_flat[static_cast(OpType::INDEX)] = compute_index_node; + dispatch_flat[static_cast(OpType::UNFOLD)] = compute_unfold_node; + dispatch_flat[static_cast(OpType::PAD)] = compute_pad_node; dispatch_flat[static_cast(OpType::BILINEAR_INTERPOLATION)] = compute_bilinear_interpolation_node; dispatch_flat[static_cast(OpType::SAMPLE)] = compute_sample_node; dispatch_flat[static_cast(OpType::TOPK)] = compute_topk_node; @@ -233,7 +258,27 @@ static const char* op_type_names[] = { "NOT_EQUAL", "SCALAR_NOT_EQUAL", "RECURRENT_CACHE_STATE", "RECURRENT_CACHE_WRITE", - "CONV_CACHE_INITIALIZE" + "CONV_CACHE_INITIALIZE", + "EQUAL", + "LESS", + "LESS_EQUAL", + "GREATER", + "GREATER_EQUAL", + "SCALAR_EQUAL", + "SCALAR_LESS", + "SCALAR_LESS_EQUAL", + "SCALAR_GREATER", + "SCALAR_GREATER_EQUAL", + "LOGICAL_NOT", + "BITWISE_AND", + "BITWISE_OR", + "BITWISE_NOT", + "WHERE", + "UNFOLD", + "PAD", + "SCALAR_FLOOR_DIVIDE", + "STRIDED_SLICE", + "EXPAND" }; static const char* get_op_name(OpType op) { @@ -320,7 +365,40 @@ std::vector infer_output_shape(const GraphNode& node, const nodes_vector } case OpType::ADD: case OpType::ADD_CLIPPED: case OpType::SUBTRACT: case OpType::MULTIPLY: case OpType::DIVIDE: case OpType::NOT_EQUAL: + case OpType::EQUAL: case OpType::LESS: case OpType::LESS_EQUAL: + case OpType::GREATER: case OpType::GREATER_EQUAL: + case OpType::BITWISE_AND: case OpType::BITWISE_OR: return BroadcastInfo::compute(in(0), in(1)).output_shape; + case OpType::WHERE: { + auto first = BroadcastInfo::compute(in(0), in(1)).output_shape; + return BroadcastInfo::compute(first, in(2)).output_shape; + } + case OpType::EXPAND: + return node.params.new_shape; + case OpType::UNFOLD: { + std::vector out = in(0); + size_t axis = node.params.axis < 0 ? out.size() + static_cast(node.params.axis) + : static_cast(node.params.axis); + out[axis] = ((out[axis] - node.params.kernel_size) / node.params.stride) + 1; + out.push_back(node.params.kernel_size); + return out; + } + case OpType::STRIDED_SLICE: { + std::vector out = in(0); + size_t axis = node.params.axis < 0 ? out.size() + static_cast(node.params.axis) + : static_cast(node.params.axis); + out[axis] = node.params.slice_length; + return out; + } + case OpType::PAD: { + std::vector out = in(0); + const auto& pads = node.params.new_shape; + for (size_t pair = 0; pair + 1 < pads.size(); pair += 2) { + size_t axis = out.size() - 1 - (pair / 2); + out[axis] += pads[pair] + pads[pair + 1]; + } + return out; + } case OpType::ATTENTION: case OpType::ATTENTION_CACHED: case OpType::ATTENTION_INT8_HYBRID: { std::vector out = in(0); if (node.params.v_head_dim > 0) out.back() = node.params.v_head_dim; @@ -417,6 +495,9 @@ void CactusGraph::infer_shapes() { switch (node.op_type) { case OpType::ADD: case OpType::ADD_CLIPPED: case OpType::SUBTRACT: case OpType::MULTIPLY: case OpType::DIVIDE: case OpType::NOT_EQUAL: + case OpType::EQUAL: case OpType::LESS: case OpType::LESS_EQUAL: + case OpType::GREATER: case OpType::GREATER_EQUAL: + case OpType::BITWISE_AND: case OpType::BITWISE_OR: node.params.broadcast_info = BroadcastInfo::compute( get_input(node, 0, nodes_, node_index_map_).shape, get_input(node, 1, nodes_, node_index_map_).shape); diff --git a/cactus-graph/src/graph_ffi.cpp b/cactus-graph/src/graph_ffi.cpp index 685b6aa3d..d1be261e7 100644 --- a/cactus-graph/src/graph_ffi.cpp +++ b/cactus-graph/src/graph_ffi.cpp @@ -281,6 +281,97 @@ int cactus_graph_not_equal(cactus_graph_t graph, cactus_node_t a, cactus_node_t } } +CACTUS_FFI_EXPORT int cactus_graph_equal(cactus_graph_t graph, cactus_node_t a, cactus_node_t b, cactus_node_t* out) { + if (!graph || !out) return fail_invalid("Invalid args to cactus_graph_equal"); + try { + *out = static_cast(as_graph(graph)->graph.equal(static_cast(a), static_cast(b))); + return 0; + } catch (const std::exception& e) { + last_error_message = e.what(); + return -1; + } +} + +CACTUS_FFI_EXPORT int cactus_graph_less(cactus_graph_t graph, cactus_node_t a, cactus_node_t b, cactus_node_t* out) { + if (!graph || !out) return fail_invalid("Invalid args to cactus_graph_less"); + try { + *out = static_cast(as_graph(graph)->graph.less(static_cast(a), static_cast(b))); + return 0; + } catch (const std::exception& e) { + last_error_message = e.what(); + return -1; + } +} + +CACTUS_FFI_EXPORT int cactus_graph_less_equal(cactus_graph_t graph, cactus_node_t a, cactus_node_t b, cactus_node_t* out) { + if (!graph || !out) return fail_invalid("Invalid args to cactus_graph_less_equal"); + try { + *out = static_cast(as_graph(graph)->graph.less_equal(static_cast(a), static_cast(b))); + return 0; + } catch (const std::exception& e) { + last_error_message = e.what(); + return -1; + } +} + +CACTUS_FFI_EXPORT int cactus_graph_greater(cactus_graph_t graph, cactus_node_t a, cactus_node_t b, cactus_node_t* out) { + if (!graph || !out) return fail_invalid("Invalid args to cactus_graph_greater"); + try { + *out = static_cast(as_graph(graph)->graph.greater(static_cast(a), static_cast(b))); + return 0; + } catch (const std::exception& e) { + last_error_message = e.what(); + return -1; + } +} + +CACTUS_FFI_EXPORT int cactus_graph_greater_equal(cactus_graph_t graph, cactus_node_t a, cactus_node_t b, cactus_node_t* out) { + if (!graph || !out) return fail_invalid("Invalid args to cactus_graph_greater_equal"); + try { + *out = static_cast(as_graph(graph)->graph.greater_equal(static_cast(a), static_cast(b))); + return 0; + } catch (const std::exception& e) { + last_error_message = e.what(); + return -1; + } +} + +CACTUS_FFI_EXPORT int cactus_graph_bitwise_and(cactus_graph_t graph, cactus_node_t a, cactus_node_t b, cactus_node_t* out) { + if (!graph || !out) return fail_invalid("Invalid args to cactus_graph_bitwise_and"); + try { + *out = static_cast(as_graph(graph)->graph.bitwise_and(static_cast(a), static_cast(b))); + return 0; + } catch (const std::exception& e) { + last_error_message = e.what(); + return -1; + } +} + +CACTUS_FFI_EXPORT int cactus_graph_bitwise_or(cactus_graph_t graph, cactus_node_t a, cactus_node_t b, cactus_node_t* out) { + if (!graph || !out) return fail_invalid("Invalid args to cactus_graph_bitwise_or"); + try { + *out = static_cast(as_graph(graph)->graph.bitwise_or(static_cast(a), static_cast(b))); + return 0; + } catch (const std::exception& e) { + last_error_message = e.what(); + return -1; + } +} + +CACTUS_FFI_EXPORT int cactus_graph_where(cactus_graph_t graph, cactus_node_t condition, cactus_node_t true_value, cactus_node_t false_value, cactus_node_t* out) { + if (!graph || !out) return fail_invalid("Invalid args to cactus_graph_where"); + try { + *out = static_cast(as_graph(graph)->graph.where( + static_cast(condition), + static_cast(true_value), + static_cast(false_value))); + return 0; + } catch (const std::exception& e) { + last_error_message = e.what(); + return -1; + } +} + int cactus_graph_scalar_add(cactus_graph_t graph, cactus_node_t x, float value, cactus_node_t *out) { if (!graph || !out) { last_error_message = "Invalid args to cactus_graph_scalar_add"; @@ -334,6 +425,17 @@ int cactus_graph_scalar_divide(cactus_graph_t graph, cactus_node_t x, float valu } } +CACTUS_FFI_EXPORT int cactus_graph_scalar_floor_divide(cactus_graph_t graph, cactus_node_t x, float value, cactus_node_t *out) { + if (!graph || !out) return fail_invalid("Invalid args to cactus_graph_scalar_floor_divide"); + try { + *out = static_cast(as_graph(graph)->graph.scalar_floor_divide(static_cast(x), value)); + return 0; + } catch (const std::exception& e) { + last_error_message = e.what(); + return -1; + } +} + int cactus_graph_scalar_not_equal(cactus_graph_t graph, cactus_node_t x, float value, cactus_node_t *out) { if (!graph || !out) { last_error_message = "Invalid args to cactus_graph_scalar_not_equal"; @@ -348,6 +450,83 @@ int cactus_graph_scalar_not_equal(cactus_graph_t graph, cactus_node_t x, float v } } +CACTUS_FFI_EXPORT int cactus_graph_scalar_equal(cactus_graph_t graph, cactus_node_t x, float value, cactus_node_t *out) { + if (!graph || !out) return fail_invalid("Invalid args to cactus_graph_scalar_equal"); + try { + *out = static_cast(as_graph(graph)->graph.scalar_equal(static_cast(x), value)); + return 0; + } catch (const std::exception& e) { + last_error_message = e.what(); + return -1; + } +} + +CACTUS_FFI_EXPORT int cactus_graph_scalar_less(cactus_graph_t graph, cactus_node_t x, float value, cactus_node_t *out) { + if (!graph || !out) return fail_invalid("Invalid args to cactus_graph_scalar_less"); + try { + *out = static_cast(as_graph(graph)->graph.scalar_less(static_cast(x), value)); + return 0; + } catch (const std::exception& e) { + last_error_message = e.what(); + return -1; + } +} + +CACTUS_FFI_EXPORT int cactus_graph_scalar_less_equal(cactus_graph_t graph, cactus_node_t x, float value, cactus_node_t *out) { + if (!graph || !out) return fail_invalid("Invalid args to cactus_graph_scalar_less_equal"); + try { + *out = static_cast(as_graph(graph)->graph.scalar_less_equal(static_cast(x), value)); + return 0; + } catch (const std::exception& e) { + last_error_message = e.what(); + return -1; + } +} + +CACTUS_FFI_EXPORT int cactus_graph_scalar_greater(cactus_graph_t graph, cactus_node_t x, float value, cactus_node_t *out) { + if (!graph || !out) return fail_invalid("Invalid args to cactus_graph_scalar_greater"); + try { + *out = static_cast(as_graph(graph)->graph.scalar_greater(static_cast(x), value)); + return 0; + } catch (const std::exception& e) { + last_error_message = e.what(); + return -1; + } +} + +CACTUS_FFI_EXPORT int cactus_graph_scalar_greater_equal(cactus_graph_t graph, cactus_node_t x, float value, cactus_node_t *out) { + if (!graph || !out) return fail_invalid("Invalid args to cactus_graph_scalar_greater_equal"); + try { + *out = static_cast(as_graph(graph)->graph.scalar_greater_equal(static_cast(x), value)); + return 0; + } catch (const std::exception& e) { + last_error_message = e.what(); + return -1; + } +} + +CACTUS_FFI_EXPORT int cactus_graph_logical_not(cactus_graph_t graph, cactus_node_t x, cactus_node_t *out) { + if (!graph || !out) return fail_invalid("Invalid args to cactus_graph_logical_not"); + try { + *out = static_cast(as_graph(graph)->graph.logical_not(static_cast(x))); + return 0; + } catch (const std::exception& e) { + last_error_message = e.what(); + return -1; + } +} + +CACTUS_FFI_EXPORT int cactus_graph_bitwise_not(cactus_graph_t graph, cactus_node_t x, cactus_node_t *out) { + if (!graph || !out) return fail_invalid("Invalid args to cactus_graph_bitwise_not"); + try { + *out = static_cast(as_graph(graph)->graph.bitwise_not(static_cast(x))); + return 0; + } catch (const std::exception& e) { + last_error_message = e.what(); + return -1; + } +} + int cactus_graph_scalar_exp(cactus_graph_t graph, cactus_node_t x, cactus_node_t *out) { if (!graph || !out) { last_error_message = "Invalid args to cactus_graph_scalar_exp"; @@ -478,6 +657,18 @@ int cactus_graph_reshape(cactus_graph_t graph, cactus_node_t x, const size_t* sh } } +CACTUS_FFI_EXPORT int cactus_graph_expand(cactus_graph_t graph, cactus_node_t x, const size_t* shape, size_t rank, cactus_node_t* out) { + if (!graph || !shape || rank == 0 || !out) return fail_invalid("Invalid args to cactus_graph_expand"); + try { + std::vector s(shape, shape + rank); + *out = static_cast(as_graph(graph)->graph.expand(static_cast(x), s)); + return 0; + } catch (const std::exception& e) { + last_error_message = e.what(); + return -1; + } +} + int cactus_graph_transpose(cactus_graph_t graph, cactus_node_t x, int32_t backend, cactus_node_t* out) { if (!graph || !out) return fail_invalid("Invalid args to cactus_graph_transpose"); try { @@ -512,6 +703,17 @@ int cactus_graph_slice(cactus_graph_t graph, cactus_node_t x, int32_t axis, size } } +CACTUS_FFI_EXPORT int cactus_graph_strided_slice(cactus_graph_t graph, cactus_node_t x, int32_t axis, size_t start, size_t length, size_t step, cactus_node_t* out) { + if (!graph || !out) return fail_invalid("Invalid args to cactus_graph_strided_slice"); + try { + *out = static_cast(as_graph(graph)->graph.strided_slice(static_cast(x), axis, start, length, step)); + return 0; + } catch (const std::exception& e) { + last_error_message = e.what(); + return -1; + } +} + int cactus_graph_index(cactus_graph_t graph, cactus_node_t x, size_t index_value, int32_t dim, cactus_node_t* out) { if (!graph || !out) return fail_invalid("Invalid args to cactus_graph_index"); try { @@ -523,6 +725,29 @@ int cactus_graph_index(cactus_graph_t graph, cactus_node_t x, size_t index_value } } +CACTUS_FFI_EXPORT int cactus_graph_unfold(cactus_graph_t graph, cactus_node_t x, int32_t dimension, size_t size, size_t step, cactus_node_t* out) { + if (!graph || !out) return fail_invalid("Invalid args to cactus_graph_unfold"); + try { + *out = static_cast(as_graph(graph)->graph.unfold(static_cast(x), dimension, size, step)); + return 0; + } catch (const std::exception& e) { + last_error_message = e.what(); + return -1; + } +} + +CACTUS_FFI_EXPORT int cactus_graph_pad(cactus_graph_t graph, cactus_node_t x, const size_t* pads, size_t pad_count, float value, cactus_node_t* out) { + if (!graph || !pads || !out) return fail_invalid("Invalid args to cactus_graph_pad"); + try { + std::vector pad_values(pads, pads + pad_count); + *out = static_cast(as_graph(graph)->graph.pad(static_cast(x), pad_values, value)); + return 0; + } catch (const std::exception& e) { + last_error_message = e.what(); + return -1; + } +} + int cactus_graph_sum(cactus_graph_t graph, cactus_node_t x, int32_t axis, cactus_node_t* out) { if (!graph || !out) return fail_invalid("Invalid args to cactus_graph_sum"); try { @@ -644,6 +869,17 @@ int cactus_graph_gather(cactus_graph_t graph, cactus_node_t tensor, cactus_node_ } } +CACTUS_FFI_EXPORT int cactus_graph_gather_dim(cactus_graph_t graph, cactus_node_t tensor, cactus_node_t indices, int32_t axis, cactus_node_t* out) { + if (!graph || !out) return fail_invalid("Invalid args to cactus_graph_gather_dim"); + try { + *out = static_cast(as_graph(graph)->graph.gather(static_cast(tensor), static_cast(indices), axis)); + return 0; + } catch (const std::exception& e) { + last_error_message = e.what(); + return -1; + } +} + int cactus_graph_embedding_from_tensor(cactus_graph_t graph, cactus_node_t embedding_tensor, cactus_node_t indices, cactus_node_t* out) { if (!graph || !out) return fail_invalid("Invalid args to cactus_graph_embedding_from_tensor"); try { @@ -1172,6 +1408,21 @@ int cactus_graph_conv1d(cactus_graph_t graph, cactus_node_t input, cactus_node_t } } +CACTUS_FFI_EXPORT int cactus_graph_conv1d_depthwise(cactus_graph_t graph, cactus_node_t input, cactus_node_t weight, bool has_bias, cactus_node_t bias, size_t stride, cactus_node_t* out) { + if (!graph || !out) return fail_invalid("Invalid args to cactus_graph_conv1d_depthwise"); + try { + if (has_bias) { + *out = static_cast(as_graph(graph)->graph.conv1d_depthwise(static_cast(input), static_cast(weight), static_cast(bias), stride)); + } else { + *out = static_cast(as_graph(graph)->graph.conv1d_depthwise(static_cast(input), static_cast(weight), stride)); + } + return 0; + } catch (const std::exception& e) { + last_error_message = e.what(); + return -1; + } +} + int cactus_graph_conv1d_same_depthwise_k9(cactus_graph_t graph, cactus_node_t input, cactus_node_t weight, bool has_bias, cactus_node_t bias, cactus_node_t* out) { if (!graph || !out) return fail_invalid("Invalid args to cactus_graph_conv1d_same_depthwise_k9"); try { diff --git a/cactus-graph/src/io.cpp b/cactus-graph/src/io.cpp index 4f6e5a25f..7a3000553 100644 --- a/cactus-graph/src/io.cpp +++ b/cactus-graph/src/io.cpp @@ -197,6 +197,13 @@ namespace { case OpType::MULTIPLY: case OpType::DIVIDE: case OpType::NOT_EQUAL: + case OpType::EQUAL: + case OpType::LESS: + case OpType::LESS_EQUAL: + case OpType::GREATER: + case OpType::GREATER_EQUAL: + case OpType::BITWISE_AND: + case OpType::BITWISE_OR: return true; default: return false; @@ -241,7 +248,7 @@ namespace { GraphFile::NodeEntry node; node.index = read_u32(in); uint32_t op_type_val = read_u32(in); - if (op_type_val > static_cast(OpType::CONV_CACHE_INITIALIZE)) { + if (op_type_val > static_cast(OpType::EXPAND)) { throw std::runtime_error("Graph file corrupted: invalid op type"); } node.op_type = static_cast(op_type_val); diff --git a/cactus-graph/src/ops_conv.cpp b/cactus-graph/src/ops_conv.cpp index 1033b01a0..717e3b4c5 100644 --- a/cactus-graph/src/ops_conv.cpp +++ b/cactus-graph/src/ops_conv.cpp @@ -234,6 +234,29 @@ void compute_conv1d_node(GraphNode& node, const std::vector(); + const __fp16* input = X.data_as<__fp16>(); + const size_t L_out = Y.shape[2]; + + for (size_t n = 0; n < N; ++n) { + for (size_t c = 0; c < C_in; ++c) { + for (size_t t = 0; t < L_out; ++t) { + float acc = bias_ptr ? static_cast(bias_ptr[c]) : 0.0f; + + for (size_t k = 0; k < K; ++k) { + size_t input_idx = (n * C_in + c) * L + t * stride + k; + size_t weight_idx = c * K + k; + acc += static_cast(input[input_idx]) * static_cast(W_ptr[weight_idx]); + } + + output[(n * C_out + c) * L_out + t] = static_cast<__fp16>(acc); + } + } + } + return; + } + cactus_conv1d_f16(X.data_as<__fp16>(), W_ptr, bias_ptr, Y.data_as<__fp16>(), N, L, C_in, C_out, K, stride); } diff --git a/cactus-graph/src/ops_math.cpp b/cactus-graph/src/ops_math.cpp index e87fc110b..f1263ff2f 100644 --- a/cactus-graph/src/ops_math.cpp +++ b/cactus-graph/src/ops_math.cpp @@ -99,6 +99,33 @@ void dispatch_binary_op_f16(OpType op, const __fp16* lhs, const __fp16* rhs, __f } }); break; + case OpType::EQUAL: + case OpType::LESS: + case OpType::LESS_EQUAL: + case OpType::GREATER: + case OpType::GREATER_EQUAL: + case OpType::BITWISE_AND: + case OpType::BITWISE_OR: + CactusThreading::parallel_for(count, CactusThreading::Thresholds::ELEMENT_WISE, + [&](size_t start_idx, size_t end_idx) { + for (size_t i = start_idx; i < end_idx; ++i) { + float a = static_cast(lhs[i]); + float b = static_cast(rhs[i]); + bool result = false; + switch (op) { + case OpType::EQUAL: result = (a == b); break; + case OpType::LESS: result = (a < b); break; + case OpType::LESS_EQUAL: result = (a <= b); break; + case OpType::GREATER: result = (a > b); break; + case OpType::GREATER_EQUAL: result = (a >= b); break; + case OpType::BITWISE_AND: result = (a != 0.0f && b != 0.0f); break; + case OpType::BITWISE_OR: result = (a != 0.0f || b != 0.0f); break; + default: break; + } + output[i] = static_cast<__fp16>(result ? 1.0f : 0.0f); + } + }); + break; default: break; } @@ -117,6 +144,20 @@ static float apply_binary_op_f32(OpType op, float lhs, float rhs) { return lhs / rhs; case OpType::NOT_EQUAL: return lhs != rhs ? 1.0f : 0.0f; + case OpType::EQUAL: + return lhs == rhs ? 1.0f : 0.0f; + case OpType::LESS: + return lhs < rhs ? 1.0f : 0.0f; + case OpType::LESS_EQUAL: + return lhs <= rhs ? 1.0f : 0.0f; + case OpType::GREATER: + return lhs > rhs ? 1.0f : 0.0f; + case OpType::GREATER_EQUAL: + return lhs >= rhs ? 1.0f : 0.0f; + case OpType::BITWISE_AND: + return (lhs != 0.0f && rhs != 0.0f) ? 1.0f : 0.0f; + case OpType::BITWISE_OR: + return (lhs != 0.0f || rhs != 0.0f) ? 1.0f : 0.0f; default: return lhs; } @@ -132,6 +173,16 @@ void dispatch_binary_op_f32(OpType op, const float* lhs, const float* rhs, float } void dispatch_unary_op_f16(OpType op, const __fp16* input, __fp16* output, size_t count, float param) { + if (op == OpType::SCALAR_FLOOR_DIVIDE) { + CactusThreading::parallel_for(count, CactusThreading::Thresholds::ELEMENT_WISE, + [&](size_t start_idx, size_t end_idx) { + for (size_t i = start_idx; i < end_idx; ++i) { + output[i] = static_cast<__fp16>(std::floor(static_cast(input[i]) / param)); + } + }); + return; + } + ScalarOpType scalar_op; switch (op) { case OpType::SCALAR_ADD: scalar_op = ScalarOpType::ADD; break; @@ -157,6 +208,7 @@ static float apply_scalar_op_f32(OpType op, float value, float param) { case OpType::SCALAR_SUBTRACT: return value - param; case OpType::SCALAR_MULTIPLY: return value * param; case OpType::SCALAR_DIVIDE: return value / param; + case OpType::SCALAR_FLOOR_DIVIDE: return std::floor(value / param); case OpType::SCALAR_EXP: return std::exp(value); case OpType::SCALAR_SQRT: return std::sqrt(value); case OpType::SCALAR_COS: return std::cos(value); @@ -270,6 +322,35 @@ void compute_binary_op_node(GraphNode& node, const std::vector(); + const __fp16* rhs_data = rhs.data_as<__fp16>(); + __fp16* out_data = node.output_buffer.data_as<__fp16>(); + CactusThreading::parallel_for(total_elements, CactusThreading::Thresholds::ELEMENT_WISE, + [&](size_t start_idx, size_t end_idx) { + std::vector coords(output_shape.size()); + size_t tmp = start_idx; + for (int axis = static_cast(output_shape.size()) - 1; axis >= 0; --axis) { + coords[static_cast(axis)] = tmp % output_shape[static_cast(axis)]; + tmp /= output_shape[static_cast(axis)]; + } + for (size_t linear_idx = start_idx; linear_idx < end_idx; ++linear_idx) { + size_t lhs_idx = broadcast_linear_index(coords, lhs_strides); + size_t rhs_idx = broadcast_linear_index(coords, rhs_strides); + float a = static_cast(lhs_data[lhs_idx]); + float b = static_cast(rhs_data[rhs_idx]); + out_data[linear_idx] = static_cast<__fp16>(apply_binary_op_f32(node.op_type, a, b)); + increment_broadcast_coords(coords, output_shape); + } + }); + break; + } default: break; } } else { @@ -292,7 +373,14 @@ void compute_unary_op_node(GraphNode& node, const std::vector(); @@ -300,8 +388,22 @@ void compute_unary_op_node(GraphNode& node, const std::vector( - static_cast(input_data[i]) != scalar ? 1.0f : 0.0f); + float value = static_cast(input_data[i]); + bool result = false; + switch (node.op_type) { + case OpType::SCALAR_NOT_EQUAL: result = value != scalar; break; + case OpType::SCALAR_EQUAL: result = value == scalar; break; + case OpType::SCALAR_LESS: result = value < scalar; break; + case OpType::SCALAR_LESS_EQUAL: result = value <= scalar; break; + case OpType::SCALAR_GREATER: result = value > scalar; break; + case OpType::SCALAR_GREATER_EQUAL: result = value >= scalar; break; + case OpType::LOGICAL_NOT: + case OpType::BITWISE_NOT: + result = value == 0.0f; + break; + default: break; + } + output_data[i] = static_cast<__fp16>(result ? 1.0f : 0.0f); } }); } else { @@ -310,7 +412,22 @@ void compute_unary_op_node(GraphNode& node, const std::vector scalar; break; + case OpType::SCALAR_GREATER_EQUAL: result = value >= scalar; break; + case OpType::LOGICAL_NOT: + case OpType::BITWISE_NOT: + result = value == 0.0f; + break; + default: break; + } + output_data[i] = result ? 1.0f : 0.0f; } }); } @@ -328,6 +445,65 @@ void compute_unary_op_node(GraphNode& node, const std::vector()[index]; + if (buffer.precision == Precision::FP16) return static_cast(buffer.data_as<__fp16>()[index]); + if (buffer.precision == Precision::INT8) return static_cast(buffer.data_as()[index]); + throw std::runtime_error("where only supports FP32/FP16/INT8 condition data"); +} + +static void write_float_value(BufferDesc& buffer, size_t index, float value) { + if (buffer.precision == Precision::FP32) { + buffer.data_as()[index] = value; + return; + } + if (buffer.precision == Precision::FP16) { + buffer.data_as<__fp16>()[index] = static_cast<__fp16>(value); + return; + } + if (buffer.precision == Precision::INT8) { + buffer.data_as()[index] = static_cast(value); + return; + } + throw std::runtime_error("where only supports FP32/FP16/INT8 output data"); +} + +void compute_where_node(GraphNode& node, const std::vector>& nodes, const std::unordered_map& node_index_map) { + const auto& condition = get_input(node, 0, nodes, node_index_map); + const auto& true_value = get_input(node, 1, nodes, node_index_map); + const auto& false_value = get_input(node, 2, nodes, node_index_map); + + if (true_value.precision != false_value.precision) { + throw std::runtime_error("where true/false branches must have matching precision"); + } + + const auto& output_shape = node.output_buffer.shape; + size_t total_elements = node.output_buffer.total_size; + std::vector condition_strides = compute_strides(condition.shape, output_shape); + std::vector true_strides = compute_strides(true_value.shape, output_shape); + std::vector false_strides = compute_strides(false_value.shape, output_shape); + + CactusThreading::parallel_for(total_elements, CactusThreading::Thresholds::ELEMENT_WISE, + [&](size_t start_idx, size_t end_idx) { + std::vector coords(output_shape.size()); + size_t tmp = start_idx; + for (int axis = static_cast(output_shape.size()) - 1; axis >= 0; --axis) { + coords[static_cast(axis)] = tmp % output_shape[static_cast(axis)]; + tmp /= output_shape[static_cast(axis)]; + } + for (size_t linear_idx = start_idx; linear_idx < end_idx; ++linear_idx) { + size_t cond_idx = broadcast_linear_index(coords, condition_strides); + size_t true_idx = broadcast_linear_index(coords, true_strides); + size_t false_idx = broadcast_linear_index(coords, false_strides); + float selected = read_value_as_float(condition, cond_idx) != 0.0f + ? read_value_as_float(true_value, true_idx) + : read_value_as_float(false_value, false_idx); + write_float_value(node.output_buffer, linear_idx, selected); + increment_broadcast_coords(coords, output_shape); + } + }); +} + void compute_activation_node(GraphNode& node, const std::vector>& nodes, const std::unordered_map& node_index_map) { const auto& input = get_input(node, 0, nodes, node_index_map); diff --git a/cactus-graph/src/ops_tensor.cpp b/cactus-graph/src/ops_tensor.cpp index 857c677b9..abe7bf25f 100644 --- a/cactus-graph/src/ops_tensor.cpp +++ b/cactus-graph/src/ops_tensor.cpp @@ -6,28 +6,124 @@ #include #include +static std::vector contiguous_strides(const std::vector& shape); + +template +static void transpose_generic( + const T* input, + T* output, + const std::vector& input_shape, + const std::vector& output_shape, + const std::vector& permutation) { + std::vector input_strides = contiguous_strides(input_shape); + + size_t total = 1; + for (size_t dim : output_shape) total *= dim; + + for (size_t output_linear = 0; output_linear < total; ++output_linear) { + size_t tmp = output_linear; + size_t input_linear = 0; + + for (int dim = static_cast(output_shape.size()) - 1; dim >= 0; --dim) { + size_t dim_index = static_cast(dim); + size_t coord = tmp % output_shape[dim_index]; + tmp /= output_shape[dim_index]; + input_linear += coord * input_strides[permutation[dim_index]]; + } + + output[output_linear] = input[input_linear]; + } +} + void compute_transpose_node(GraphNode& node, const std::vector>& nodes, const std::unordered_map& node_index_map) { if (node.params.backend == ComputeBackend::NPU) { throw std::runtime_error("NPU transpose operation not yet implemented"); } const auto& input_buffer = get_input(node, 0, nodes, node_index_map); + const auto& permutation = node.params.permutation; - if (input_buffer.precision != Precision::FP16) { - throw std::runtime_error("Transpose only supports FP16 precision"); + if (input_buffer.precision == Precision::FP16) { + const __fp16* input = input_buffer.data_as<__fp16>(); + __fp16* output = node.output_buffer.data_as<__fp16>(); + cactus_transpose_f16(input, output, input_buffer.shape.data(), permutation.data(), permutation.size(), 0, input_buffer.total_size); + return; } - const auto& permutation = node.params.permutation; + if (input_buffer.precision == Precision::FP32) { + transpose_generic( + input_buffer.data_as(), + node.output_buffer.data_as(), + input_buffer.shape, + node.output_buffer.shape, + permutation); + return; + } - const __fp16* input = input_buffer.data_as<__fp16>(); - __fp16* output = node.output_buffer.data_as<__fp16>(); - cactus_transpose_f16(input, output, input_buffer.shape.data(), permutation.data(), permutation.size(), 0, input_buffer.total_size); + if (input_buffer.precision == Precision::INT8) { + transpose_generic( + input_buffer.data_as(), + node.output_buffer.data_as(), + input_buffer.shape, + node.output_buffer.shape, + permutation); + return; + } + + throw std::runtime_error("Transpose only supports FP32/FP16/INT8 precision"); } void compute_gather_node(GraphNode& node, const std::vector>& nodes, const std::unordered_map& node_index_map) { const auto& tensor_buffer = get_input(node, 0, nodes, node_index_map); const auto& indices_buffer = get_input(node, 1, nodes, node_index_map); + if (tensor_buffer.shape.size() == indices_buffer.shape.size()) { + int axis_value = node.params.axis; + if (axis_value < 0) axis_value += static_cast(tensor_buffer.shape.size()); + if (axis_value < 0 || static_cast(axis_value) >= tensor_buffer.shape.size()) { + throw std::runtime_error("Gather axis out of range"); + } + + const size_t axis = static_cast(axis_value); + std::vector tensor_strides = contiguous_strides(tensor_buffer.shape); + std::vector output_strides = contiguous_strides(node.output_buffer.shape); + const size_t elem_bytes = PrecisionTraits::packed_size_of(tensor_buffer.precision, 1); + const char* tensor_data = static_cast(tensor_buffer.get_data()); + char* output_data = static_cast(node.output_buffer.get_data()); + + auto read_index = [&](size_t i) -> size_t { + if (indices_buffer.precision == Precision::INT8) { + return static_cast(indices_buffer.data_as()[i]); + } + if (indices_buffer.precision == Precision::FP16) { + return static_cast(indices_buffer.data_as<__fp16>()[i]); + } + return static_cast(indices_buffer.data_as()[i]); + }; + + for (size_t output_linear = 0; output_linear < node.output_buffer.total_size; ++output_linear) { + size_t tmp = output_linear; + size_t input_linear = 0; + size_t gathered_index = read_index(output_linear); + + if (gathered_index >= tensor_buffer.shape[axis]) { + throw std::runtime_error("Gather index out of bounds"); + } + + for (size_t dim = 0; dim < node.output_buffer.shape.size(); ++dim) { + size_t coord = tmp / output_strides[dim]; + tmp %= output_strides[dim]; + input_linear += (dim == axis ? gathered_index : coord) * tensor_strides[dim]; + } + + std::memcpy( + output_data + PrecisionTraits::byte_offset_of(node.output_buffer.precision, output_linear), + tensor_data + PrecisionTraits::byte_offset_of(tensor_buffer.precision, input_linear), + elem_bytes); + } + return; + } + size_t first_dim = tensor_buffer.shape[0]; size_t element_size = 1; for (size_t i = 1; i < tensor_buffer.shape.size(); i++) { @@ -124,6 +220,55 @@ void compute_slice_node(GraphNode& node, const std::vector>& nodes, const std::unordered_map& node_index_map) { + const auto& input_buffer = get_input(node, 0, nodes, node_index_map); + const auto& input_shape = input_buffer.shape; + const auto& output_shape = node.output_buffer.shape; + + int axis_value = node.params.axis; + if (axis_value < 0) axis_value += static_cast(input_shape.size()); + if (axis_value < 0 || static_cast(axis_value) >= input_shape.size()) { + throw std::runtime_error("strided_slice axis out of range"); + } + if (node.params.stride == 0) { + throw std::runtime_error("strided_slice step must be greater than zero"); + } + + const size_t axis = static_cast(axis_value); + const size_t start = node.params.slice_start; + const size_t step = node.params.stride; + const size_t length = node.params.slice_length; + + if (output_shape.size() != input_shape.size() || output_shape[axis] != length) { + throw std::runtime_error("strided_slice shape mismatch"); + } + + const char* input_data = static_cast(input_buffer.get_data()); + char* output_data = static_cast(node.output_buffer.get_data()); + if (!input_data || !output_data) { + throw std::runtime_error("strided_slice input/output buffer is not available"); + } + + auto dims = AxisDims::from_shape(input_shape, axis); + const size_t elem_bytes = PrecisionTraits::packed_size_of(input_buffer.precision, dims.inner); + + for (size_t outer = 0; outer < dims.outer; ++outer) { + for (size_t axis_index = 0; axis_index < length; ++axis_index) { + size_t input_axis = start + axis_index * step; + if (input_axis >= dims.axis_size) { + throw std::runtime_error("strided_slice index out of bounds"); + } + + const size_t input_offset = (outer * dims.axis_size + input_axis) * dims.inner; + const size_t output_offset = (outer * length + axis_index) * dims.inner; + std::memcpy( + output_data + PrecisionTraits::byte_offset_of(input_buffer.precision, output_offset), + input_data + PrecisionTraits::byte_offset_of(input_buffer.precision, input_offset), + elem_bytes); + } + } +} + void compute_embedding_node(GraphNode& node, const std::vector>& nodes, const std::unordered_map& node_index_map) { const auto& embeddings_buffer = get_input(node, 0, nodes, node_index_map); const auto& indices_buffer = get_input(node, 1, nodes, node_index_map); @@ -323,6 +468,198 @@ void compute_index_node(GraphNode& node, const std::vector>& nodes, const std::unordered_map& node_index_map) { + const auto& input_buffer = get_input(node, 0, nodes, node_index_map); + const auto& input_shape = input_buffer.shape; + const auto& output_shape = node.output_buffer.shape; + + if (input_shape.empty() || output_shape.size() != input_shape.size() + 1) { + throw std::runtime_error("unfold shape mismatch"); + } + + int axis_value = node.params.axis; + if (axis_value < 0) axis_value += static_cast(input_shape.size()); + if (axis_value < 0 || static_cast(axis_value) >= input_shape.size()) { + throw std::runtime_error("unfold dimension out of range"); + } + + const size_t axis = static_cast(axis_value); + const size_t window_size = node.params.kernel_size; + const size_t step = node.params.stride; + if (window_size == 0 || step == 0) { + throw std::runtime_error("unfold window size and step must be greater than zero"); + } + + std::vector input_strides(input_shape.size(), 1); + for (int i = static_cast(input_shape.size()) - 2; i >= 0; --i) { + input_strides[static_cast(i)] = + input_strides[static_cast(i + 1)] * input_shape[static_cast(i + 1)]; + } + + const size_t elem_bytes = PrecisionTraits::packed_size_of(input_buffer.precision, 1); + const char* input_data = static_cast(input_buffer.get_data()); + char* output_data = static_cast(node.output_buffer.get_data()); + + if (!input_data || !output_data) { + throw std::runtime_error("unfold input/output buffer is not available"); + } + + for (size_t out_linear = 0; out_linear < node.output_buffer.total_size; ++out_linear) { + size_t tmp = out_linear; + std::vector coords(output_shape.size(), 0); + for (int dim = static_cast(output_shape.size()) - 1; dim >= 0; --dim) { + coords[static_cast(dim)] = tmp % output_shape[static_cast(dim)]; + tmp /= output_shape[static_cast(dim)]; + } + + size_t input_linear = 0; + for (size_t dim = 0; dim < input_shape.size(); ++dim) { + size_t input_coord = coords[dim]; + if (dim == axis) { + input_coord = coords[dim] * step + coords.back(); + } + input_linear += input_coord * input_strides[dim]; + } + + std::memcpy( + output_data + PrecisionTraits::byte_offset_of(node.output_buffer.precision, out_linear), + input_data + PrecisionTraits::byte_offset_of(input_buffer.precision, input_linear), + elem_bytes); + } +} + +static std::vector contiguous_strides(const std::vector& shape) { + std::vector strides(shape.size(), 1); + for (int i = static_cast(shape.size()) - 2; i >= 0; --i) { + strides[static_cast(i)] = strides[static_cast(i + 1)] * shape[static_cast(i + 1)]; + } + return strides; +} + +void compute_expand_node(GraphNode& node, const std::vector>& nodes, const std::unordered_map& node_index_map) { + const auto& input_buffer = get_input(node, 0, nodes, node_index_map); + const auto& input_shape = input_buffer.shape; + const auto& output_shape = node.output_buffer.shape; + + if (output_shape.size() < input_shape.size()) { + throw std::runtime_error("expand output rank is smaller than input rank"); + } + + const char* input_data = static_cast(input_buffer.get_data()); + char* output_data = static_cast(node.output_buffer.get_data()); + if (!input_data || !output_data) { + throw std::runtime_error("expand input/output buffer is not available"); + } + + std::vector input_strides = contiguous_strides(input_shape); + const size_t rank_offset = output_shape.size() - input_shape.size(); + const size_t elem_bytes = PrecisionTraits::packed_size_of(input_buffer.precision, 1); + + for (size_t output_linear = 0; output_linear < node.output_buffer.total_size; ++output_linear) { + size_t tmp = output_linear; + size_t input_linear = 0; + + for (int dim = static_cast(output_shape.size()) - 1; dim >= 0; --dim) { + size_t dim_index = static_cast(dim); + size_t coord = tmp % output_shape[dim_index]; + tmp /= output_shape[dim_index]; + + if (dim_index < rank_offset) { + continue; + } + + size_t input_dim_index = dim_index - rank_offset; + size_t input_coord = input_shape[input_dim_index] == 1 ? 0 : coord; + input_linear += input_coord * input_strides[input_dim_index]; + } + + std::memcpy( + output_data + PrecisionTraits::byte_offset_of(node.output_buffer.precision, output_linear), + input_data + PrecisionTraits::byte_offset_of(input_buffer.precision, input_linear), + elem_bytes); + } +} + +static void fill_buffer(BufferDesc& buffer, float value) { + if (value == 0.0f) { + std::memset(buffer.get_data(), 0, buffer.byte_size); + return; + } + + if (buffer.precision == Precision::FP32) { + float* data = buffer.data_as(); + for (size_t i = 0; i < buffer.total_size; ++i) data[i] = value; + return; + } + + if (buffer.precision == Precision::FP16) { + __fp16* data = buffer.data_as<__fp16>(); + for (size_t i = 0; i < buffer.total_size; ++i) data[i] = static_cast<__fp16>(value); + return; + } + + if (buffer.precision == Precision::INT8) { + int8_t* data = buffer.data_as(); + for (size_t i = 0; i < buffer.total_size; ++i) data[i] = static_cast(value); + return; + } + + throw std::runtime_error("pad only supports FP32/FP16/INT8 precision"); +} + +void compute_pad_node(GraphNode& node, const std::vector>& nodes, const std::unordered_map& node_index_map) { + const auto& input_buffer = get_input(node, 0, nodes, node_index_map); + const auto& input_shape = input_buffer.shape; + const auto& output_shape = node.output_buffer.shape; + const auto& pads = node.params.new_shape; + + if (input_shape.empty() || output_shape.size() != input_shape.size()) { + throw std::runtime_error("pad shape mismatch"); + } + if (pads.size() % 2 != 0 || (pads.size() / 2) > input_shape.size()) { + throw std::runtime_error("pad requires left/right pad pairs within input rank"); + } + if (input_buffer.precision != Precision::FP32 && + input_buffer.precision != Precision::FP16 && + input_buffer.precision != Precision::INT8) { + throw std::runtime_error("pad only supports FP32/FP16/INT8 precision"); + } + + fill_buffer(node.output_buffer, node.params.scalar); + + std::vector left_offsets(input_shape.size(), 0); + for (size_t pair = 0; pair < pads.size(); pair += 2) { + size_t axis = input_shape.size() - 1 - (pair / 2); + left_offsets[axis] = pads[pair]; + } + + std::vector input_strides = contiguous_strides(input_shape); + std::vector output_strides = contiguous_strides(output_shape); + const char* input_data = static_cast(input_buffer.get_data()); + char* output_data = static_cast(node.output_buffer.get_data()); + const size_t elem_bytes = PrecisionTraits::packed_size_of(input_buffer.precision, 1); + + if (!input_data || !output_data) { + throw std::runtime_error("pad input/output buffer is not available"); + } + + for (size_t input_linear = 0; input_linear < input_buffer.total_size; ++input_linear) { + size_t tmp = input_linear; + size_t output_linear = 0; + + for (size_t dim = 0; dim < input_shape.size(); ++dim) { + size_t coord = tmp / input_strides[dim]; + tmp %= input_strides[dim]; + output_linear += (coord + left_offsets[dim]) * output_strides[dim]; + } + + std::memcpy( + output_data + PrecisionTraits::byte_offset_of(node.output_buffer.precision, output_linear), + input_data + PrecisionTraits::byte_offset_of(input_buffer.precision, input_linear), + elem_bytes); + } +} + void compute_bilinear_interpolation_node(GraphNode& node, const std::vector>& nodes, const std::unordered_map& node_index_map) { const auto& pos_embeds_buffer = get_input(node, 0, nodes, node_index_map); diff --git a/cactus-graph/src/param_io.cpp b/cactus-graph/src/param_io.cpp index 2e3fdadd4..61834c1bf 100644 --- a/cactus-graph/src/param_io.cpp +++ b/cactus-graph/src/param_io.cpp @@ -166,7 +166,13 @@ const Schema& op_schema(OpType op_type) { {OpType::SCALAR_SUBTRACT, {{ParamField::Scalar, FieldPersistence::Persistent}}}, {OpType::SCALAR_MULTIPLY, {{ParamField::Scalar, FieldPersistence::Persistent}}}, {OpType::SCALAR_DIVIDE, {{ParamField::Scalar, FieldPersistence::Persistent}}}, + {OpType::SCALAR_FLOOR_DIVIDE, {{ParamField::Scalar, FieldPersistence::Persistent}}}, {OpType::SCALAR_NOT_EQUAL, {{ParamField::Scalar, FieldPersistence::Persistent}}}, + {OpType::SCALAR_EQUAL, {{ParamField::Scalar, FieldPersistence::Persistent}}}, + {OpType::SCALAR_LESS, {{ParamField::Scalar, FieldPersistence::Persistent}}}, + {OpType::SCALAR_LESS_EQUAL, {{ParamField::Scalar, FieldPersistence::Persistent}}}, + {OpType::SCALAR_GREATER, {{ParamField::Scalar, FieldPersistence::Persistent}}}, + {OpType::SCALAR_GREATER_EQUAL, {{ParamField::Scalar, FieldPersistence::Persistent}}}, {OpType::CLAMP, {{ParamField::Scalar, FieldPersistence::Persistent}, {ParamField::Scale, FieldPersistence::Persistent}}}, {OpType::SOFTMAX, {{ParamField::Axis, FieldPersistence::Persistent}}}, {OpType::SUM, {{ParamField::Axis, FieldPersistence::Persistent}}}, @@ -176,14 +182,19 @@ const Schema& op_schema(OpType op_type) { {OpType::MAX, {{ParamField::Axis, FieldPersistence::Persistent}}}, {OpType::CUMSUM, {{ParamField::Axis, FieldPersistence::Persistent}}}, {OpType::INDEX, {{ParamField::Axis, FieldPersistence::Persistent}, {ParamField::IndexValue, FieldPersistence::Persistent}}}, + {OpType::GATHER, {{ParamField::Axis, FieldPersistence::Persistent}}}, {OpType::CONCAT, {{ParamField::Axis, FieldPersistence::Persistent}}}, {OpType::CAT, {{ParamField::Axis, FieldPersistence::Persistent}}}, {OpType::VIEW, {{ParamField::NewShape, FieldPersistence::Persistent}}}, {OpType::RESHAPE, {{ParamField::NewShape, FieldPersistence::Persistent}}}, + {OpType::EXPAND, {{ParamField::NewShape, FieldPersistence::Persistent}}}, {OpType::FLATTEN, {{ParamField::NewShape, FieldPersistence::Persistent}}}, {OpType::MATMUL, {{ParamField::PretransposedRhs, FieldPersistence::Persistent}, {ParamField::Backend, FieldPersistence::Persistent}}}, {OpType::TRANSPOSE, {{ParamField::Permutation, FieldPersistence::Persistent}, {ParamField::Backend, FieldPersistence::Persistent}}}, {OpType::SLICE, {{ParamField::Axis, FieldPersistence::Persistent}, {ParamField::SliceStart, FieldPersistence::Persistent}, {ParamField::SliceLength, FieldPersistence::Persistent}}}, + {OpType::STRIDED_SLICE, {{ParamField::Axis, FieldPersistence::Persistent}, {ParamField::SliceStart, FieldPersistence::Persistent}, {ParamField::SliceLength, FieldPersistence::Persistent}, {ParamField::Stride, FieldPersistence::Persistent}}}, + {OpType::UNFOLD, {{ParamField::Axis, FieldPersistence::Persistent}, {ParamField::KernelSize, FieldPersistence::Persistent}, {ParamField::Stride, FieldPersistence::Persistent}}}, + {OpType::PAD, {{ParamField::NewShape, FieldPersistence::Persistent}, {ParamField::Scalar, FieldPersistence::Persistent}}}, {OpType::RMS_NORM, {{ParamField::Epsilon, FieldPersistence::Persistent}}}, {OpType::LAYERNORM, {{ParamField::Epsilon, FieldPersistence::Persistent}}}, {OpType::GROUPNORM, {{ParamField::Epsilon, FieldPersistence::Persistent}, {ParamField::NumGroups, FieldPersistence::Persistent}}}, @@ -240,7 +251,7 @@ const Schema& op_schema(OpType op_type) { {OpType::CONV1D_CAUSAL, {{ParamField::Dilation, FieldPersistence::Persistent}}}, {OpType::CONV1D_K3, {{ParamField::Stride, FieldPersistence::Persistent}}}, {OpType::CONV1D_K7S3, {{ParamField::Stride, FieldPersistence::Persistent}}}, - {OpType::CONV1D, {{ParamField::Stride, FieldPersistence::Persistent}}}, + {OpType::CONV1D, {{ParamField::Stride, FieldPersistence::Persistent}, {ParamField::NumGroups, FieldPersistence::Persistent}}}, {OpType::SAMPLE, {{ParamField::Temperature, FieldPersistence::Persistent}, {ParamField::TopP, FieldPersistence::Persistent}, {ParamField::MinP, FieldPersistence::Persistent}, {ParamField::RepetitionPenalty, FieldPersistence::Persistent}, {ParamField::TopK, FieldPersistence::Persistent}, {ParamField::RandomSeed, FieldPersistence::Persistent}, {ParamField::BiasIndices, FieldPersistence::Persistent}, {ParamField::BiasValues, FieldPersistence::Persistent}}}, }; diff --git a/python/tests/test_cli_run.py b/python/tests/test_cli_run.py index 5c4c0d7e5..80ff8a471 100644 --- a/python/tests/test_cli_run.py +++ b/python/tests/test_cli_run.py @@ -1,4 +1,5 @@ from argparse import Namespace +import json from pathlib import Path from types import SimpleNamespace @@ -67,3 +68,40 @@ def fake_subprocess_run(cmd): assert cmd[cmd.index("--input-ids-file") + 1] == str(input_ids_file) assert cmd[cmd.index("--max-new-tokens") + 1] == "4" assert cmd[cmd.index("--result-json") + 1] == str(result_json) + + +def test_resolve_bundle_dir_accepts_runtime_plan_path(tmp_path: Path) -> None: + bundle_dir = tmp_path / "bundle" + runtime_plan = bundle_dir / "runtime_plan.json" + runtime_plan.parent.mkdir(parents=True) + runtime_plan.write_text( + json.dumps( + { + "family": "gemma4_e2b", + "components": [ + { + "component": "decoder_step", + "graph": "components/decoder_step.cactus", + "runtime_input_node_ids": [1, 2], + "logical_inputs": ["input_ids", "position_ids"], + "output_node_ids": [3], + "logical_outputs": ["logits"], + } + ], + "routes": [], + "metadata": {"npu_vision_encoder": "vision.mlpackage"}, + } + ), + encoding="utf-8", + ) + + resolved = model_mod.resolve_bundle_dir(str(runtime_plan)) + engine_manifest = bundle_dir / "components" / "manifest.json" + + assert resolved == bundle_dir + assert engine_manifest.exists() + + manifest = json.loads(engine_manifest.read_text(encoding="utf-8")) + assert manifest["family"] == "gemma4_e2b" + assert manifest["npu_vision_encoder"] == "vision.mlpackage" + assert manifest["components"][0]["component"] == "decoder_step" diff --git a/python/tests/test_generator_component_split.py b/python/tests/test_generator_component_split.py new file mode 100644 index 000000000..cbc7a7815 --- /dev/null +++ b/python/tests/test_generator_component_split.py @@ -0,0 +1,115 @@ +import unittest + +from cactus.transpiler.Fusions import models as FModels +from cactus.transpiler.Generator import component_split +from cactus.transpiler.IR import models as IRModels + + +class TestGeneratorComponentSplit(unittest.TestCase): + def test_component_extraction_rewrites_aliases_and_logical_outputs(self) -> None: + input_ids = IRModels.Node( + index=0, + name="input_ids", + node_type="placeholder", + target="input_ids", + args=[], + kwargs={}, + users=(), + tensor_output_meta={"shape": [1, 1], "dtype": "torch.int64"}, + module_stack=None, + value_kind=FModels.ValueKind.USER_INPUT, + ) + weight = IRModels.Node( + index=1, + name="p_embed_weight", + node_type="placeholder", + target="p_embed_weight", + args=[], + kwargs={}, + users=(), + tensor_output_meta={"shape": [10, 4], "dtype": "torch.float16"}, + module_stack=None, + value_kind=FModels.ValueKind.PARAMETER, + ) + cleanup = IRModels.Node( + index=2, + name="index_put", + node_type="call_function", + target="aten.index_put.default", + args=[{"node": "input_ids"}], + kwargs={}, + users=(), + tensor_output_meta={"shape": [1, 1], "dtype": "torch.int64"}, + module_stack=None, + value_kind=FModels.ValueKind.ACTIVATION, + ) + embedding = IRModels.Node( + index=3, + name="embedding", + node_type="call_function", + target="aten.embedding.default", + args=[{"node": "p_embed_weight"}, {"node": "index_put"}], + kwargs={}, + users=(), + tensor_output_meta={"shape": [1, 1, 4], "dtype": "torch.float16"}, + module_stack=None, + value_kind=FModels.ValueKind.ACTIVATION, + ) + scaled = IRModels.Node( + index=4, + name="scaled", + node_type="call_function", + target="aten.mul.Tensor", + args=[{"node": "embedding"}, 2.0], + kwargs={}, + users=(), + tensor_output_meta={"shape": [1, 1, 4], "dtype": "torch.float16"}, + module_stack=None, + value_kind=FModels.ValueKind.ACTIVATION, + attrs={"other": 2.0}, + ) + output = IRModels.Node( + index=5, + name="output", + node_type="output", + target="output", + args=[[{"node": "scaled"}]], + kwargs={}, + users=(), + tensor_output_meta=None, + module_stack=None, + value_kind=FModels.ValueKind.OUTPUT, + ) + graph = IRModels.rebuild_graph((input_ids, weight, cleanup, embedding, scaled, output), empty_graph()) + spec = component_split.ComponentSplitSpec( + name="lm_encoder_step", + graph=graph, + outputs=( + component_split.OutputSpec("scaled", "inputs_embeds"), + component_split.OutputSpec("position_ids", "position_ids"), + ), + placeholders=(component_split.PlaceholderSpec("position_ids", "position_ids", tensor_node="input_ids"),), + ref_aliases={"index_put": "input_ids"}, + ) + + component = component_split.extract_component_graph(spec) + + self.assertNotIn("index_put", component.nodes_map) + self.assertEqual(component.nodes_map["embedding"].args, [{"node": "p_embed_weight"}, {"node": "input_ids"}]) + self.assertEqual(component.nodes_map["scaled"].ir_metadata["logical_output"], "inputs_embeds") + self.assertEqual(component.nodes_map["position_ids"].ir_metadata["logical_input"], "position_ids") + self.assertEqual(component.nodes_map["position_ids"].ir_metadata["logical_output"], "position_ids") + + +def empty_graph() -> IRModels.Graph: + return IRModels.Graph( + source=None, + sources=(), + outputs=(), + nodes=(), + nodes_map={}, + ) + + +if __name__ == "__main__": + unittest.main() diff --git a/python/tests/test_graph.py b/python/tests/test_graph.py index c8b05f88e..cdcc965cd 100644 --- a/python/tests/test_graph.py +++ b/python/tests/test_graph.py @@ -714,5 +714,72 @@ def test_save_load_roundtrip_layernorm(self): np.testing.assert_allclose(out, expected, atol=1e-2) +class TestGraphConvCache(unittest.TestCase): + + def test_conv1d_causal_matches_pytorch_depthwise_layout(self): + try: + import torch + except ImportError: + self.skipTest("torch is required for PyTorch numerical reference") + + batch_size = 1 + hidden_dim = 4 + sequence_length = 5 + kernel_size = 3 + + x_bhs = torch.arange(batch_size * hidden_dim * sequence_length, dtype=torch.float32).reshape( + batch_size, + hidden_dim, + sequence_length, + ) / 10.0 + weight = torch.arange(hidden_dim * kernel_size, dtype=torch.float32).reshape(hidden_dim, 1, kernel_size) / 20.0 + expected = torch.nn.functional.conv1d( + torch.nn.functional.pad(x_bhs, (kernel_size - 1, 0)), + weight, + groups=hidden_dim, + ) + + g = Graph() + x = g.input((batch_size, sequence_length, hidden_dim)) + w = g.input((hidden_dim, 1, kernel_size)) + y = g.conv1d_causal(x, w, kernel_size, 1) + + g.set_input(x, x_bhs.permute(0, 2, 1).numpy().astype(np.float16)) + g.set_input(w, weight.numpy().astype(np.float16)) + g.execute() + + np.testing.assert_allclose( + y.numpy(), + expected.permute(0, 2, 1).numpy().astype(np.float16), + atol=3e-2, + ) + + def test_conv_cache_initialize_then_append_matches_shifted_pytorch_window(self): + try: + import torch + except ImportError: + self.skipTest("torch is required for PyTorch numerical reference") + + window_size = 3 + hidden_dim = 4 + rows = torch.arange(window_size * hidden_dim, dtype=torch.float32).reshape(window_size, hidden_dim) + new_row = torch.tensor([[100.0, 101.0, 102.0, 103.0]], dtype=torch.float32) + expected = torch.cat((rows[1:], new_row), dim=0) + + g = Graph() + rows_input = g.input((window_size, hidden_dim)) + new_input = g.input((1, hidden_dim)) + cache_state = g.conv_cache_state(window_size, hidden_dim) + initialized = g.conv_cache_initialize(rows_input, cache_state) + window = g.conv_cache_append(new_input, cache_state) + + g.set_input(rows_input, rows.numpy().astype(np.float16)) + g.set_input(new_input, new_row.numpy().astype(np.float16)) + g.execute() + + self.assertEqual(initialized.numpy().shape, (0,)) + np.testing.assert_allclose(window.numpy(), expected.numpy().astype(np.float16), atol=1e-2) + + if __name__ == "__main__": unittest.main() diff --git a/python/tests/test_runtime_plan.py b/python/tests/test_runtime_plan.py new file mode 100644 index 000000000..c36322d54 --- /dev/null +++ b/python/tests/test_runtime_plan.py @@ -0,0 +1,96 @@ +import json +import sys +import tempfile +import unittest +from pathlib import Path + +sys.path.insert(0, str(Path(__file__).parent.parent)) + +from cactus.transpiler.ModelProfiles.profiles import GEMMA4_E2B_PROFILE +from cactus.transpiler.Generator import models as GModels +from cactus.transpiler.RuntimePlan import models as RPModels + + +class FakeTensor: + + def __init__(self, node_id): + self.id = node_id + + +class TestRuntimePlan(unittest.TestCase): + + def test_generator_component_manifest_keeps_logical_names(self): + component = GModels.ComponentGraph( + name="decoder_step", + ir_graph=None, + output_path=Path("decoder_step.cactus"), + manifest_path=Path("decoder_step.graph_manifest.json"), + ) + + component.add_runtime_input(FakeTensor(1), "input_ids") + component.add_runtime_input(FakeTensor(2), "position_ids") + component.add_output(FakeTensor(3), "logits") + + manifest = GModels.ComponentGraphManifest.from_component(component).to_dict() + + self.assertEqual(manifest["logical_inputs"], ["input_ids", "position_ids"]) + self.assertEqual(manifest["logical_outputs"], ["logits"]) + self.assertEqual(manifest["graph"], "decoder_step.cactus") + + def test_writes_engine_manifest_from_generator_manifest(self): + with tempfile.TemporaryDirectory() as tmpdir: + bundle_dir = Path(tmpdir) + generator_manifest = bundle_dir / "decoder_step.graph_manifest.json" + generator_manifest.write_text( + json.dumps( + { + "component": "decoder_step", + "graph_path": "decoder_step.cactus", + "runtime_input_node_ids": [1, 2], + "logical_inputs": ["input_ids", "position_ids"], + "output_node_ids": [3], + "logical_outputs": ["logits"], + "weight_bindings": [ + { + "node_id": 4, + "path": "weights/model.layers.0.self_attn.q_proj.weight.bin", + } + ], + "cache_state_node_ids": [ + { + "layer_key": "kv:0", + "key": 5, + "value": 6, + "cache_kind": "kv", + "tensor_indices": [0, 1], + } + ], + } + ), + encoding="utf-8", + ) + + plan = RPModels.runtime_plan_from_generator_manifests( + {"decoder_step": generator_manifest}, + bundle_dir=bundle_dir, + model_profile=GEMMA4_E2B_PROFILE, + ) + engine_manifest_path, runtime_plan_path = plan.write(bundle_dir) + + engine_manifest = json.loads(engine_manifest_path.read_text(encoding="utf-8")) + runtime_plan = json.loads(runtime_plan_path.read_text(encoding="utf-8")) + component = engine_manifest["components"][0] + + self.assertEqual(engine_manifest["family"], "gemma4_e2b") + self.assertEqual(component["component"], "decoder_step") + self.assertEqual(component["graph"], "decoder_step.cactus") + self.assertEqual(component["logical_inputs"], ["input_ids", "position_ids"]) + self.assertEqual(component["logical_outputs"], ["logits"]) + self.assertEqual(component["bound_constant_bindings"][0]["node_id"], 4) + self.assertEqual(component["cache_state_node_ids"][0]["layer_key"], "kv:0") + self.assertTrue(runtime_plan["routes"]) + self.assertEqual(runtime_plan["routes"][0]["edges"][0]["inputs"], ["text_embed"]) + + +if __name__ == "__main__": + unittest.main() diff --git a/python/tests/test_server_live.py b/python/tests/test_server_live.py index 664ceae8b..1d813561f 100644 --- a/python/tests/test_server_live.py +++ b/python/tests/test_server_live.py @@ -42,7 +42,7 @@ def _require_bundle(relative: Path, types: set[str]) -> Path: if not _valid_bundle(candidate): pytest.fail( f"Live-test model is not a prepared v2 bundle: {relative}\n" - "Expected config.txt and components/manifest.json.\n" + "Expected config.txt and either components/manifest.json or runtime_plan.json.\n" f"Prepare it with `cactus download google/gemma-4-E2B-it` from {PROJECT_ROOT}." ) model_type = _read_model_type(candidate) From b65b6670ea717bba8627b0894802b4207d143712 Mon Sep 17 00:00:00 2001 From: Parkirat Date: Wed, 29 Jul 2026 22:36:46 -0400 Subject: [PATCH 68/84] Performance upgrades --- python/cactus/transpiler/Fusions/edges.py | 39 ++ python/cactus/transpiler/Fusions/fusions.py | 64 +++ python/cactus/transpiler/Fusions/nodes.py | 11 + .../cactus/transpiler/Generator/generate.py | 478 +++++++++++++++++- python/cactus/transpiler/Generator/models.py | 4 + python/cactus/transpiler/IR/models.py | 85 ++++ python/cactus/transpiler/IR/simplify_ir.py | 2 +- .../transpiler/ModelProfiles/profiles.py | 1 + .../cactus/transpiler/RuntimePlan/models.py | 15 +- python/cactus/transpiler/main.py | 10 +- 10 files changed, 695 insertions(+), 14 deletions(-) diff --git a/python/cactus/transpiler/Fusions/edges.py b/python/cactus/transpiler/Fusions/edges.py index 5f256b22c..c2c1972f8 100644 --- a/python/cactus/transpiler/Fusions/edges.py +++ b/python/cactus/transpiler/Fusions/edges.py @@ -58,10 +58,13 @@ def _index_edges_by_dest(edges: dict[str, M.FusionEdge]) -> dict[str, tuple[str, EDGES: dict[str, M.FusionEdge] = { "linear_mm_to_bias_add": _edge("linear_mm", "linear_bias_add", 0), + "linear_weight_transpose_to_mm": _edge("linear_weight_transpose", "linear_mm", 1), "rms_square_to_mean": _edge("rms_square", "rms_mean", 0), "rms_mean_to_eps_add": _edge("rms_mean", "rms_eps_add", 0), "rms_eps_add_to_inv": _edge("rms_eps_add", "rms_inv", 0), + "rms_eps_add_to_inv_pow": _edge("rms_eps_add", "rms_inv_pow", 0), "rms_inv_to_scale": _edge("rms_inv", "rms_scale", 1), + "rms_inv_pow_to_scale": _edge("rms_inv_pow", "rms_scale", 1), "rms_scale_to_weight_mul": _edge("rms_scale", "rms_weight_mul", 0), "ln_mean_to_center": _edge("ln_mean", "ln_center", 1), "ln_center_to_square": _edge("ln_center", "ln_square", 0), @@ -76,6 +79,19 @@ def _index_edges_by_dest(edges: dict[str, M.FusionEdge]) -> dict[str, tuple[str, "mlp_activation_to_product": _edge("mlp_activation", "mlp_product", 0), "up_proj_to_product": _edge("up_proj", "mlp_product", 1), "mlp_product_to_down_proj": _edge("mlp_product", "down_proj", 0), + "gate_input_view_to_gate_proj": _edge("gate_input_view", "gate_proj", 0), + "gate_weight_transpose_to_gate_proj": _edge("gate_weight_transpose", "gate_proj", 1), + "gate_proj_to_gate_proj_view": _edge("gate_proj", "gate_proj_view", 0), + "gate_proj_view_to_gelu": _edge("gate_proj_view", "gelu", 0), + "gelu_to_product": _edge("gelu", "mlp_product", 0), + "up_input_view_to_up_proj": _edge("up_input_view", "up_proj", 0), + "up_weight_transpose_to_up_proj": _edge("up_weight_transpose", "up_proj", 1), + "up_proj_to_up_proj_view": _edge("up_proj", "up_proj_view", 0), + "up_proj_view_to_product": _edge("up_proj_view", "mlp_product", 1), + "mlp_product_to_product_view": _edge("mlp_product", "mlp_product_view", 0), + "product_view_to_down_proj": _edge("mlp_product_view", "down_proj", 0), + "down_weight_transpose_to_down_proj": _edge("down_weight_transpose", "down_proj", 1), + "down_proj_to_down_proj_view": _edge("down_proj", "down_proj_view", 0), "up_proj_to_gelu": _edge("up_proj", "gelu", 0), "gelu_to_down_proj": _edge("gelu", "down_proj", 0), "q_proj_to_q_view": _edge("q_proj", "q_view", 0), @@ -223,6 +239,7 @@ def _index_edges_by_dest(edges: dict[str, M.FusionEdge]) -> dict[str, tuple[str, EDGE_GROUPS: dict[str, tuple[str, ...]] = { "linear_bias": ("linear_mm_to_bias_add",), + "linear_transposed": ("linear_weight_transpose_to_mm",), "rms_norm": ( "rms_square_to_mean", "rms_mean_to_eps_add", @@ -230,6 +247,13 @@ def _index_edges_by_dest(edges: dict[str, M.FusionEdge]) -> dict[str, tuple[str, "rms_inv_to_scale", "rms_scale_to_weight_mul", ), + "rms_norm_pow": ( + "rms_square_to_mean", + "rms_mean_to_eps_add", + "rms_eps_add_to_inv_pow", + "rms_inv_pow_to_scale", + "rms_scale_to_weight_mul", + ), "layernorm": ( "ln_mean_to_center", "ln_center_to_square", @@ -251,6 +275,21 @@ def _index_edges_by_dest(edges: dict[str, M.FusionEdge]) -> dict[str, tuple[str, "up_proj_to_gelu", "gelu_to_down_proj", ), + "gemma4_geglu_mlp": ( + "gate_input_view_to_gate_proj", + "gate_weight_transpose_to_gate_proj", + "gate_proj_to_gate_proj_view", + "gate_proj_view_to_gelu", + "gelu_to_product", + "up_input_view_to_up_proj", + "up_weight_transpose_to_up_proj", + "up_proj_to_up_proj_view", + "up_proj_view_to_product", + "mlp_product_to_product_view", + "product_view_to_down_proj", + "down_weight_transpose_to_down_proj", + "down_proj_to_down_proj_view", + ), "attention_core": ( "attn_qk_to_scale", "attn_scale_to_mask_add", diff --git a/python/cactus/transpiler/Fusions/fusions.py b/python/cactus/transpiler/Fusions/fusions.py index be8cc2723..52192da43 100644 --- a/python/cactus/transpiler/Fusions/fusions.py +++ b/python/cactus/transpiler/Fusions/fusions.py @@ -413,6 +413,18 @@ def _index_by_root_op( }, ) +LINEAR_TRANSPOSED_GRAPH = _graph( + "linear_transposed", + "linear_mm", + ("linear_weight_transpose", "linear_mm"), + edge_names=E.EDGE_GROUPS["linear_transposed"], + inputs=( + _input("x", "linear_mm", 0), + _input("weight", "linear_weight_transpose", 0, allowed_value_kinds=(M.ValueKind.PARAMETER, M.ValueKind.BUFFER, M.ValueKind.LIFTED_CONSTANT)), + ), + attr_captures=(M.AttrCapture("pretransposed_rhs", default=True, required=False),), +) + LINEAR_BIAS_GRAPH = _graph( "linear_bias", "linear_bias_add", @@ -439,6 +451,19 @@ def _index_by_root_op( }, ) +RMS_NORM_POW_GRAPH = _graph( + "rms_norm_pow", + "rms_weight_mul", + ("rms_square", "rms_mean", "rms_eps_add", "rms_inv_pow", "rms_scale", "rms_weight_mul"), + edge_names=E.EDGE_GROUPS["rms_norm_pow"], + inputs=(_input("x", "rms_square", 0), _input("weight", "rms_weight_mul", 1)), + shared_inputs=(_shared_input("rms_square", 0, "rms_scale", 0),), + attr_captures=(M.AttrCapture("epsilon", "rms_eps_add", "other", default=1e-6, required=False),), + constraints={ + "node_attr_equals": {"node": "rms_mean", "attr": "dim", "value": -1}, + }, +) + LAYERNORM_NO_BIAS_GRAPH = _graph( "layernorm_no_bias", "ln_weight_mul", @@ -493,6 +518,39 @@ def _index_by_root_op( edge_names=E.EDGE_GROUPS["gelu_mlp"], ) +GEMMA4_GEGLU_MLP_GRAPH = _graph( + "gemma4_geglu_mlp", + "down_proj_view", + ( + "gate_input_view", + "gate_weight_transpose", + "gate_proj", + "gate_proj_view", + "gelu", + "up_input_view", + "up_weight_transpose", + "up_proj", + "up_proj_view", + "mlp_product", + "mlp_product_view", + "down_weight_transpose", + "down_proj", + "down_proj_view", + ), + edge_names=E.EDGE_GROUPS["gemma4_geglu_mlp"], + inputs=( + _input("hidden", "gate_input_view", 0), + _input("gate_weight", "gate_weight_transpose", 0, allowed_value_kinds=(M.ValueKind.PARAMETER, M.ValueKind.BUFFER, M.ValueKind.LIFTED_CONSTANT)), + _input("up_weight", "up_weight_transpose", 0, allowed_value_kinds=(M.ValueKind.PARAMETER, M.ValueKind.BUFFER, M.ValueKind.LIFTED_CONSTANT)), + _input("down_weight", "down_weight_transpose", 0, allowed_value_kinds=(M.ValueKind.PARAMETER, M.ValueKind.BUFFER, M.ValueKind.LIFTED_CONSTANT)), + ), + shared_inputs=(_shared_input("gate_input_view", 0, "up_input_view", 0),), + attr_captures=(M.AttrCapture("product_scale", default=1.0, required=False),), + constraints={ + "node_attr_equals": {"node": "gelu", "attr": "approximate", "value": "tanh"}, + }, +) + ATTENTION_DIRECT_GRAPH = _graph( "scaled_dot_product_attention", "attn_sdp", @@ -1023,12 +1081,15 @@ def _index_by_root_op( **DIRECT_GRAPHS, **DSP_GRAPHS, "linear": LINEAR_GRAPH, + "linear_transposed": LINEAR_TRANSPOSED_GRAPH, "linear_bias": LINEAR_BIAS_GRAPH, "rms_norm": RMS_NORM_GRAPH, + "rms_norm_pow": RMS_NORM_POW_GRAPH, "layernorm_no_bias": LAYERNORM_NO_BIAS_GRAPH, "layernorm": LAYERNORM_GRAPH, "swiglu_mlp": SWIGLU_MLP_GRAPH, "gelu_mlp": GELU_MLP_GRAPH, + "gemma4_geglu_mlp": GEMMA4_GEGLU_MLP_GRAPH, "scaled_dot_product_attention": ATTENTION_DIRECT_GRAPH, "attention_core": ATTENTION_CORE_GRAPH, "attention_masked": ATTENTION_MASKED_GRAPH, @@ -1076,12 +1137,15 @@ def _index_by_root_op( FUSIONS: dict[str, M.FusionDefinition] = { **DIRECT_DEFINITIONS, "linear": _definition("linear", "matmul", LINEAR_GRAPH, fusion_fields=("generic", "linear")), + "linear_transposed": _definition("linear_transposed", "matmul", LINEAR_TRANSPOSED_GRAPH, fusion_fields=("generic", "linear")), "linear_bias": _definition("linear_bias", "linear", LINEAR_BIAS_GRAPH, fusion_fields=("generic", "linear")), "rms_norm": _definition("rms_norm", "rms_norm", RMS_NORM_GRAPH, fusion_fields=("generic", "rmsnorm", "gemma4_rmsnorm", "qwen2_5_rmsnorm")), + "rms_norm_pow": _definition("rms_norm_pow", "rms_norm", RMS_NORM_POW_GRAPH, fusion_fields=("generic", "rmsnorm", "gemma4_rmsnorm", "qwen2_5_rmsnorm")), "layernorm_no_bias": _definition("layernorm_no_bias", "layernorm", LAYERNORM_NO_BIAS_GRAPH, fusion_fields=("generic", "normalization")), "layernorm": _definition("layernorm", "layernorm", LAYERNORM_GRAPH, fusion_fields=("generic", "normalization")), "swiglu_mlp": _definition("swiglu_mlp", "dense_mlp_tq_fused", SWIGLU_MLP_GRAPH, fusion_fields=("generic", "mlp", "gemma4_mlp", "qwen2_5_mlp", "lfm_mlp")), "gelu_mlp": _definition("gelu_mlp", "matmul", GELU_MLP_GRAPH, fusion_fields=("generic", "mlp")), + "gemma4_geglu_mlp": _definition("gemma4_geglu_mlp", "dense_mlp_tq_fused", GEMMA4_GEGLU_MLP_GRAPH, fusion_fields=("gemma4_mlp",)), "scaled_dot_product_attention": _definition("scaled_dot_product_attention", "attention", ATTENTION_DIRECT_GRAPH, fusion_fields=("generic", "attention")), "attention_core": _definition("attention_core", "attention", ATTENTION_CORE_GRAPH, fusion_fields=("generic", "attention", "gemma4_attention", "qwen2_5_attention", "lfm_attention")), "attention_masked": _definition("attention_masked", "attention", ATTENTION_MASKED_GRAPH, fusion_fields=("generic", "attention", "gemma4_attention")), diff --git a/python/cactus/transpiler/Fusions/nodes.py b/python/cactus/transpiler/Fusions/nodes.py index b7cecbc43..e2d423a50 100644 --- a/python/cactus/transpiler/Fusions/nodes.py +++ b/python/cactus/transpiler/Fusions/nodes.py @@ -171,11 +171,13 @@ def _expert_node(name: str, op_key: str) -> M.FusionNode: "equal": _node("equal", "eq"), "embedding": _node("embedding", "embedding"), "linear_mm": _node("linear_mm", "matmul"), + "linear_weight_transpose": _node("linear_weight_transpose", "transpose"), "linear_bias_add": _node("linear_bias_add", "add"), "rms_square": _node("rms_square", "pow", attrs=(_attr("exponent", 2, required=False),)), "rms_mean": _node("rms_mean", "mean", attrs=(_attr("dim", -1, required=False),)), "rms_eps_add": _node("rms_eps_add", "add"), "rms_inv": _node("rms_inv", "sqrt"), + "rms_inv_pow": _node("rms_inv_pow", "pow", attrs=(_attr("exponent", -0.5, required=True),)), "rms_scale": _node("rms_scale", "mul"), "rms_weight_mul": _node("rms_weight_mul", "mul"), "ln_mean": _node("ln_mean", "mean", attrs=(_attr("dim", -1, required=False),)), @@ -204,9 +206,18 @@ def _expert_node(name: str, op_key: str) -> M.FusionNode: "glu": _node("glu", "mul"), "gate_proj": _node("gate_proj", "matmul"), "up_proj": _node("up_proj", "matmul"), + "gate_input_view": _node("gate_input_view", "view"), + "up_input_view": _node("up_input_view", "view"), + "gate_weight_transpose": _node("gate_weight_transpose", "transpose"), + "up_weight_transpose": _node("up_weight_transpose", "transpose"), + "down_weight_transpose": _node("down_weight_transpose", "transpose"), + "gate_proj_view": _node("gate_proj_view", "view"), + "up_proj_view": _node("up_proj_view", "view"), "mlp_activation": _node("mlp_activation", "silu"), "mlp_product": _node("mlp_product", "mul"), + "mlp_product_view": _node("mlp_product_view", "view"), "down_proj": _node("down_proj", "matmul"), + "down_proj_view": _node("down_proj_view", "view"), "q_proj": _node("q_proj", "matmul"), "k_proj": _node("k_proj", "matmul"), "v_proj": _node("v_proj", "matmul"), diff --git a/python/cactus/transpiler/Generator/generate.py b/python/cactus/transpiler/Generator/generate.py index dc227ea72..8a5967b73 100644 --- a/python/cactus/transpiler/Generator/generate.py +++ b/python/cactus/transpiler/Generator/generate.py @@ -1,6 +1,8 @@ from __future__ import annotations import json +import math +import struct from collections.abc import Mapping from dataclasses import replace from pathlib import Path @@ -21,6 +23,7 @@ class UnsupportedLoweringError(NotImplementedError): GraphInput = CModels.LayerMap | IRModels.Graph ComponentInput = GraphInput | Mapping[str, GraphInput] +SIZE_T_MAX = (1 << 64) - 1 def generate( @@ -258,6 +261,112 @@ def add_rules(rules: dict[str, models.LoweringRule], targets: Any, lower: models rules[target] = models.LoweringRule(target, lower) +def cache_attention_generation_plan(graph: IRModels.Graph) -> tuple[frozenset[str], frozenset[str], dict[str, IRModels.CacheAnnotation]]: + skip_names: set[str] = set() + cache_state_names: set[str] = set() + prefill_annotations = prefill_cache_cat_annotations(graph) + + for node in graph.nodes: + if node.target != "cactus.attention" or len(node.parents) < 3: + continue + + key_cat = find_cache_cat_ancestor(node.parents[1], FModels.CacheTensorRole.KEY) + value_cat = find_cache_cat_ancestor(node.parents[2], FModels.CacheTensorRole.VALUE) + + if key_cat is None or value_cat is None: + continue + + key_cache = key_cat.parents[0] + value_cache = value_cat.parents[0] + + if key_cache.cache is None or value_cache.cache is None: + continue + + if key_cache.cache.layer_index != value_cache.cache.layer_index: + continue + + skip_names.update(cache_wrapper_path_names(node.parents[1], key_cat)) + skip_names.update(cache_wrapper_path_names(node.parents[2], value_cat)) + cache_state_names.update((key_cache.name, value_cache.name)) + + for cat_name in prefill_annotations: + cat_node = graph.nodes_map.get(cat_name) + + if cat_node is not None and cat_node.parents: + skip_names.add(cat_node.parents[0].name) + + return frozenset(skip_names), frozenset(cache_state_names), prefill_annotations + + +def prefill_cache_cat_annotations(graph: IRModels.Graph) -> dict[str, IRModels.CacheAnnotation]: + annotations: dict[str, IRModels.CacheAnnotation] = {} + candidates = [node for node in graph.nodes if is_empty_cache_cat(node)] + + for index, node in enumerate(candidates): + shape = concrete_shape(meta_shape(node)) + + if shape is None or len(shape) != 4: + continue + + _, num_kv_heads, sequence_length, head_dim = shape + annotations[node.name] = IRModels.CacheAnnotation( + kind=FModels.CacheKind.KV, + role=FModels.CacheTensorRole.KEY if index % 2 == 0 else FModels.CacheTensorRole.VALUE, + tensor_index=index, + layer_index=index // 2, + shape=tuple(shape), + layout="batch_heads_sequence_head_dim", + sequence_length=int(sequence_length), + num_kv_heads=int(num_kv_heads), + head_dim=int(head_dim), + source="prefill_empty_cache_cat", + ) + + return annotations + + +def is_empty_cache_cat(node: IRModels.Node) -> bool: + if node.target not in {"cactus.cat", "aten.cat.default"} or len(node.parents) < 2: + return False + + first_parent = node.parents[0] + return first_parent.value_kind == FModels.ValueKind.LIFTED_CONSTANT and element_count(meta_shape(first_parent)) == 0 + + +def find_cache_cat_ancestor(node: IRModels.Node, role: str, max_depth: int = 16) -> IRModels.Node | None: + current = node + + for _ in range(max_depth): + if current.target in {"cactus.cat", "aten.cat.default"} and len(current.parents) >= 2: + cache_parent = current.parents[0] + + if cache_parent.cache is not None and cache_parent.cache.kind == FModels.CacheKind.KV and cache_parent.cache.role == role: + return current + + if len(current.parents) != 1: + return None + + current = current.parents[0] + + return None + + +def cache_wrapper_path_names(start: IRModels.Node, stop: IRModels.Node) -> set[str]: + names: set[str] = set() + current = start + + while True: + names.add(current.name) + + if current is stop: + return names + + if len(current.parents) != 1: + return set() + + current = current.parents[0] + + def lower_component( component: models.ComponentGraph, config: models.GeneratorConfig, @@ -266,6 +375,11 @@ def lower_component( component.graph = create_cactus_graph() component.weight_resolver = component_weight_resolver(component, config) context = models.GenerationContext(component=component, config=config, lowerings=lowering_rules) + ( + context.skip_node_names, + context.cache_state_placeholder_names, + context.prefill_cache_cat_annotations, + ) = cache_attention_generation_plan(component.ir_graph) for node in IRModels.topological_sort(component.ir_graph): try: @@ -288,6 +402,9 @@ def lower_component( def lower_node(context: models.GenerationContext, node: IRModels.Node) -> None: + if node.name in context.skip_node_names: + return + if node.is_placeholder: context.bind(node, lower_placeholder(context, node)) return @@ -308,7 +425,7 @@ def lower_node(context: models.GenerationContext, node: IRModels.Node) -> None: def lower_placeholder(context: models.GenerationContext, node: IRModels.Node) -> Any: - if should_lower_cache_placeholder_as_state(node): + if should_lower_cache_placeholder_as_state(context, node): return lower_cache_placeholder(context, node) shape, dynamic_dims = graph_input_shape(node) @@ -333,10 +450,13 @@ def placeholder_precision(context: models.GenerationContext, node: IRModels.Node return cactus_precision(context.graph, tensor_dtype(node)) -def should_lower_cache_placeholder_as_state(node: IRModels.Node) -> bool: +def should_lower_cache_placeholder_as_state(context: models.GenerationContext, node: IRModels.Node) -> bool: if node.cache is None: return False + if node.name in context.cache_state_placeholder_names: + return True + if node.cache.kind == FModels.CacheKind.CONV: return any(child.target in {"cactus.conv_cache_append", "cactus.conv_cache_initialize"} for child in node.children) @@ -535,6 +655,10 @@ def lower_binary(context: models.GenerationContext, node: IRModels.Node) -> Any: if len(inputs) == 2: left, right = align_binary_inputs(context, node, inputs) + + if method == "add" and context.component.name in {"decoder_step", "decoder_prefill_chunk"} and looks_like_gemma_residual_add(node): + return context.graph.add_clipped(left, right) + return getattr(context.graph, method)(left, right) scalar = scalar_attr(node, "other") @@ -546,6 +670,41 @@ def lower_binary(context: models.GenerationContext, node: IRModels.Node) -> Any: raise unsupported_arity(node, len(inputs), "two tensor inputs or one tensor plus scalar attr") +def looks_like_gemma_residual_add(node: IRModels.Node) -> bool: + if len(node.parents) != 2: + return False + + shape = meta_shape(node) + + if not shape or len(shape) > 3: + return False + + left, right = node.parents + return ( + is_gemma_residual_branch(left, right) + or is_gemma_residual_branch(right, left) + ) + + +def is_gemma_residual_branch(residual: IRModels.Node, branch: IRModels.Node) -> bool: + branch_source = strip_precision_passthrough(branch) + + if branch_source.target != "cactus.rms_norm" or not branch_source.parents: + return False + + return strip_precision_passthrough(residual) is not strip_precision_passthrough(branch_source.parents[0]) + + +def strip_precision_passthrough(node: IRModels.Node) -> IRModels.Node: + current = node + passthrough_targets = constants.PASS_THROUGH_TARGETS | constants.TO_COPY_TARGETS + + while current.target in passthrough_targets and len(current.parents) == 1: + current = current.parents[0] + + return current + + def lower_scalar(context: models.GenerationContext, node: IRModels.Node) -> Any: inputs = require_input_count(context, node, 1) method, attr_name = constants.SCALAR_TARGETS[node.target] @@ -888,6 +1047,12 @@ def lower_unfold(context: models.GenerationContext, node: IRModels.Node) -> Any: def lower_cat(context: models.GenerationContext, node: IRModels.Node) -> Any: + if node.name in context.prefill_cache_cat_annotations: + prefill_cache = lower_prefill_cache_cat(context, node) + + if prefill_cache is not None: + return prefill_cache + inputs = context.inputs_for(node) if not inputs: @@ -902,6 +1067,29 @@ def lower_cat(context: models.GenerationContext, node: IRModels.Node) -> Any: return context.graph.cat(inputs, axis=axis) +def lower_prefill_cache_cat(context: models.GenerationContext, node: IRModels.Node) -> Any | None: + annotation = context.prefill_cache_cat_annotations.get(node.name) + + if annotation is None or len(node.parents) < 2: + return None + + new_kv_node = node.parents[1] + original_new_kv = context.require(new_kv_node.name) + new_kv = cache_new_kv_for_native(context, new_kv_node, original_new_kv) + new_kv = cast_to_precision(context, new_kv, context.graph.FP16) + cache_state = context.graph.kv_cache_state( + kv_cache_capacity(context, annotation), + required_cache_int(annotation.num_kv_heads, node, "num_kv_heads"), + required_cache_int(annotation.head_dim, node, "head_dim"), + window_size=int(annotation.window_size or 0), + sink_size=0, + num_slots=1, + ) + record_cache_state_binding(context, node, cache_state, annotation) + context.graph.kv_cache_append(new_kv, cache_state, window_size=int(annotation.window_size or 0), sink_size=0) + return original_new_kv + + def lower_matmul(context: models.GenerationContext, node: IRModels.Node) -> Any: inputs = require_input_count(context, node, 2) @@ -1049,6 +1237,9 @@ def lower_to_copy(context: models.GenerationContext, node: IRModels.Node) -> Any if dtype is None or dtype == "torch.bool": return inputs[0] + if str(dtype) == "torch.float32" and can_skip_float32_copy(node): + return inputs[0] + precision_name = constants.DTYPE_TO_PRECISION.get(str(dtype)) if precision_name is None: @@ -1057,6 +1248,18 @@ def lower_to_copy(context: models.GenerationContext, node: IRModels.Node) -> Any return context.graph.precision_cast(inputs[0], cactus_precision(context.graph, str(dtype))) +def can_skip_float32_copy(node: IRModels.Node) -> bool: + if not node.children: + return False + + safe_consumers = { + "aten.bmm.default", + "cactus.attention", + "cactus.rms_norm", + } + return all(child.target in safe_consumers for child in node.children) + + def lower_pass_through(context: models.GenerationContext, node: IRModels.Node) -> Any: inputs = require_at_least_one_input(context, node) return inputs[0] @@ -1123,6 +1326,32 @@ def lower_constant_pad_nd(context: models.GenerationContext, node: IRModels.Node def lower_constant_input(context: models.GenerationContext, node: IRModels.Node) -> Any: + constant = deterministic_constant_values(node) + + if constant is not None: + values, shape = constant + dtype = cactus_precision(context.graph, tensor_dtype(node)) + tensor = context.graph.input(shape, dtype=dtype) + path = write_constant_tensor(context, node, values, shape, dtype) + node_id = models.tensor_node_id(tensor) + + if node_id is None: + raise UnsupportedLoweringError(f"{node.name}: generated constant does not expose a Cactus node id") + + context.component.add_weight_binding( + models.WeightBinding( + placeholder=node.name, + source_target=node.target, + node_id=node_id, + path=path, + output_name=path, + precision=precision_name(context.graph, dtype), + component=context.component.name, + binding_kind="generated_constant", + ) + ) + return tensor + context.component.warn(f"{node.name}: constant-producing op {node.target} lowered as graph input") shape, dynamic_dims = graph_input_shape(node) dtype = cactus_precision(context.graph, tensor_dtype(node)) @@ -1131,6 +1360,131 @@ def lower_constant_input(context: models.GenerationContext, node: IRModels.Node) return tensor +def deterministic_constant_values(node: IRModels.Node) -> tuple[list[float], tuple[int, ...]] | None: + shape = concrete_shape(meta_shape(node)) + + if node.target.startswith("aten.arange"): + start = numeric_attr(node, "start", "arg_0", default=0) + end = numeric_attr(node, "end", "arg_1") + step = numeric_attr(node, "step", "arg_2", default=1) + + if end is None or step in {None, 0}: + return None + + values = arange_values(float(start), float(end), float(step)) + return values, (len(values),) + + if node.target == "aten.scalar_tensor.default": + value = numeric_attr(node, "value", "arg_0", default=0) + return [float(value or 0)], (1,) + + if node.target in {"aten.full.default", "aten.full_like.default"}: + if shape is None: + raw_size = node.attrs.get("size", node.attrs.get("arg_0")) + shape = tuple(int(dim) for dim in raw_size) if isinstance(raw_size, (list, tuple)) else None + + if shape is None: + return None + + fill_value = numeric_attr(node, "fill_value", "arg_1", default=0) + count = math.prod(shape) if shape else 1 + return [float(fill_value or 0)] * int(count), tuple(shape or (1,)) + + return None + + +def arange_values(start: float, end: float, step: float) -> list[float]: + values: list[float] = [] + current = start + + if step > 0: + while current < end: + values.append(current) + current += step + else: + while current > end: + values.append(current) + current += step + + return values + + +def numeric_attr(node: IRModels.Node, *names: str, default: Any | None = None) -> Any | None: + for name in names: + value = node.attrs.get(name) + + if value is not None: + return value + + return default + + +def write_constant_tensor( + context: models.GenerationContext, + node: IRModels.Node, + values: list[float], + shape: tuple[int, ...], + precision: int, +) -> str: + constants_dir = context.component.output_path.parent / "constants" + constants_dir.mkdir(parents=True, exist_ok=True) + filename = f"{models.sanitize_component_name(context.component.name)}__{models.sanitize_component_name(node.name)}.weights" + path = constants_dir / filename + data = constant_data_bytes(values, precision) + write_cactus_tensor_file(path, shape, precision, data) + return constant_binding_path(context, path) + + +def constant_data_bytes(values: list[float], precision: int) -> bytes: + if precision == 1: + return b"".join(struct.pack(" None: + alignment = 32 + header_size = 84 + ndim = len(shape) + original_n = shape[0] if shape else 0 + + with path.open("wb") as f: + f.write(struct.pack(" str: + try: + return str(path.relative_to(context.config.output_dir.parent)) + except ValueError: + return str(path) + + +def precision_name(graph: Any, precision: int) -> str: + for name in ("INT8", "FP16", "FP32", "CQ1", "CQ2", "CQ3", "CQ4"): + if int(getattr(graph, name)) == int(precision): + return name + + return str(precision) + + def lower_neg(context: models.GenerationContext, node: IRModels.Node) -> Any: inputs = require_input_count(context, node, 1) return context.graph.scalar_multiply(inputs[0], -1.0) @@ -1200,7 +1554,7 @@ def lower_norm(context: models.GenerationContext, node: IRModels.Node) -> Any: if target == "cactus.rms_norm": require_len(node, inputs, 2) - return context.graph.rms_norm(inputs[0], inputs[1], eps=epsilon_attr(node)) + return lower_rms_norm(context, node, inputs[0], inputs[1]) if target in {"cactus.layernorm", "cactus.layer_norm", "aten.native_layer_norm.default", "aten.layer_norm.default"}: require_len(node, inputs, 2) @@ -1218,6 +1572,32 @@ def lower_norm(context: models.GenerationContext, node: IRModels.Node) -> Any: raise UnsupportedLoweringError(f"{node.name}: unsupported norm target {target}") +def lower_rms_norm(context: models.GenerationContext, node: IRModels.Node, x: Any, weight: Any) -> Any: + x = cast_to_precision(context, x, context.graph.FP16) + weight = cast_to_precision(context, weight, context.graph.FP16) + input_shape = concrete_shape(meta_shape(node.parents[0])) if node.parents else concrete_shape(meta_shape(node)) + + if input_shape is None: + raise UnsupportedLoweringError(f"{node.name}: cactus.rms_norm requires concrete input shape") + + if len(input_shape) == 2: + return context.graph.rms_norm(x, weight, eps=epsilon_attr(node)) + + if len(input_shape) == 1: + return context.graph.reshape( + context.graph.rms_norm(context.graph.reshape(x, (1, input_shape[0])), weight, eps=epsilon_attr(node)), + input_shape, + ) + + rows = element_count(input_shape[:-1]) + + if rows is None: + raise UnsupportedLoweringError(f"{node.name}: cactus.rms_norm requires concrete leading dimensions") + + normalized = context.graph.rms_norm(context.graph.reshape(x, (rows, input_shape[-1])), weight, eps=epsilon_attr(node)) + return context.graph.reshape(normalized, input_shape) + + def lower_conv(context: models.GenerationContext, node: IRModels.Node) -> Any: inputs = context.inputs_for(node) target = node.target @@ -1365,9 +1745,13 @@ def require_plain_conv1d_attrs(node: IRModels.Node, method: str) -> None: def lower_attention(context: models.GenerationContext, node: IRModels.Node) -> Any: - inputs = context.inputs_for(node) - if node.target == "cactus.attention" or node.target == "aten.scaled_dot_product_attention.default": + cached = lower_cache_backed_attention(context, node) + + if cached is not None: + return cached + + inputs = context.inputs_for(node) require_len(node, inputs, 3) query, key, value, mask = attention_inputs_for_layout(context, node, inputs) query = cast_to_precision(context, query, context.graph.FP16) @@ -1387,6 +1771,7 @@ def lower_attention(context: models.GenerationContext, node: IRModels.Node) -> A return attention_output_for_layout(context, node, output) if node.target == "cactus.attention_cached": + inputs = context.inputs_for(node) require_len(node, inputs, 5) return context.graph.attention_cached( inputs[0], @@ -1395,7 +1780,7 @@ def lower_attention(context: models.GenerationContext, node: IRModels.Node) -> A inputs[3], inputs[4], scale=float(node.attrs.get("scale", 1.0)), - position_offset=int(node.attrs.get("position_offset", 0)), + position_offset=cached_attention_position_offset(context, node), window_size=int(node.attrs.get("window_size", 0)), v_head_dim=int(node.attrs.get("v_head_dim", 0)), ) @@ -1403,6 +1788,78 @@ def lower_attention(context: models.GenerationContext, node: IRModels.Node) -> A raise UnsupportedLoweringError(f"{node.name}: unsupported attention target {node.target}") +def lower_cache_backed_attention(context: models.GenerationContext, node: IRModels.Node) -> Any | None: + if node.target != "cactus.attention" or len(node.parents) < 3: + return None + + key_cat = find_cache_cat_ancestor(node.parents[1], FModels.CacheTensorRole.KEY) + value_cat = find_cache_cat_ancestor(node.parents[2], FModels.CacheTensorRole.VALUE) + + if key_cat is None or value_cat is None: + return None + + key_cache = key_cat.parents[0] + value_cache = value_cat.parents[0] + key_new_node = key_cat.parents[1] + value_new_node = value_cat.parents[1] + + query = context.require(node.parents[0].name) + key_new = cache_new_kv_for_native(context, key_new_node, context.require(key_new_node.name)) + value_new = cache_new_kv_for_native(context, value_new_node, context.require(value_new_node.name)) + key_cache_state = context.require(key_cache.name) + value_cache_state = context.require(value_cache.name) + + if node.attrs.get("input_layout") == "bhqd_bhds_bhsd": + query = context.graph.permute(query, (0, 2, 1, 3)) + + query = cast_to_precision(context, query, context.graph.FP16) + key_new = cast_to_precision(context, key_new, context.graph.FP16) + value_new = cast_to_precision(context, value_new, context.graph.FP16) + + cache_pair = (key_cat.name, value_cat.name) + + if cache_pair not in context.appended_cache_pairs: + context.graph.kv_cache_append(key_new, key_cache_state, window_size=int(node.attrs.get("window_size", 0)), sink_size=0) + context.graph.kv_cache_append(value_new, value_cache_state, window_size=int(node.attrs.get("window_size", 0)), sink_size=0) + context.appended_cache_pairs.add(cache_pair) + + output = context.graph.attention_cached( + query, + key_new, + value_new, + key_cache_state, + value_cache_state, + scale=float(node.attrs.get("scale", 1.0)), + position_offset=cached_attention_position_offset(context, node), + window_size=int(node.attrs.get("window_size", 0)), + v_head_dim=last_dim(value_new_node), + ) + return attention_output_for_layout(context, node, output) + + +def cached_attention_position_offset(context: models.GenerationContext, node: IRModels.Node) -> int: + if context.component.name == "decoder_step": + return SIZE_T_MAX + + return int(node.attrs.get("position_offset", 0)) + + +def cache_new_kv_for_native(context: models.GenerationContext, node: IRModels.Node, value: Any) -> Any: + if len(meta_shape(node)) == 4: + return context.graph.permute(value, (0, 2, 1, 3)) + + return value + + +def last_dim(node: IRModels.Node) -> int: + shape = concrete_shape(meta_shape(node)) + + if not shape: + return 0 + + return int(shape[-1]) + + def attention_inputs_for_layout(context: models.GenerationContext, node: IRModels.Node, inputs: tuple[Any, ...]) -> tuple[Any, Any, Any, Any | None]: mask = inputs[3] if len(inputs) > 3 else None @@ -1595,7 +2052,14 @@ def lower_moe(context: models.GenerationContext, node: IRModels.Node) -> Any: if node.target == "cactus.dense_mlp_tq_fused": require_len(node, inputs, 4) - return context.graph.dense_mlp_tq_fused(inputs[0], inputs[1], inputs[2], inputs[3], product_scale=float(node.attrs.get("product_scale", 1.0))) + hidden = cast_to_precision(context, inputs[0], context.graph.FP16) + return context.graph.dense_mlp_tq_fused( + hidden, + inputs[1], + inputs[2], + inputs[3], + product_scale=float(node.attrs.get("product_scale", 1.0)), + ) if node.target == "cactus.moe_layer_gated": if len(inputs) < 6: diff --git a/python/cactus/transpiler/Generator/models.py b/python/cactus/transpiler/Generator/models.py index fbfd36098..802bb6c1f 100644 --- a/python/cactus/transpiler/Generator/models.py +++ b/python/cactus/transpiler/Generator/models.py @@ -99,6 +99,10 @@ class GenerationContext: config: GeneratorConfig lowerings: dict[str, LoweringRule] = field(default_factory=dict) values: dict[str, GraphTensor] = field(default_factory=dict) + skip_node_names: frozenset[str] = frozenset() + cache_state_placeholder_names: frozenset[str] = frozenset() + prefill_cache_cat_annotations: dict[str, IRModels.CacheAnnotation] = field(default_factory=dict) + appended_cache_pairs: set[tuple[str, str]] = field(default_factory=set) @property def graph(self) -> CactusGraph: diff --git a/python/cactus/transpiler/IR/models.py b/python/cactus/transpiler/IR/models.py index fe4487187..e2e5a121a 100644 --- a/python/cactus/transpiler/IR/models.py +++ b/python/cactus/transpiler/IR/models.py @@ -81,6 +81,9 @@ def from_map(cls, layer_map: CModels.LayerMap) -> "Graph": def apply_fusions(self, fusion_results: list["FusionResult"] | tuple["FusionResult", ...]) -> "Graph": return apply_fusions_to_graph(self, tuple(fusion_results)) + def remove_noop_nodes(self) -> "Graph": + return remove_noop_nodes_from_graph(self) + def to_layer_map(self) -> CModels.LayerMap: return graph_to_layer_map(self) @@ -413,6 +416,88 @@ def rewrite_node_refs(value: Any, replacement_names: dict[str, str]) -> Any: return value +def remove_noop_nodes_from_graph(graph: Graph) -> Graph: + replacement_names: dict[str, str] = {} + + for node in graph.nodes: + replacement = noop_replacement_name(node) + + if replacement is not None: + replacement_names[node.name] = replacement + + if not replacement_names: + return graph + + replacement_names = resolve_replacement_names(replacement_names) + kept_nodes = tuple(clone_rewritten_node(node, replacement_names) for node in graph.nodes if node.name not in replacement_names) + return prune_dead_nodes(rebuild_graph(kept_nodes, graph, tuple(graph.fusions))) + + +def noop_replacement_name(node: Node) -> str | None: + if len(node.parents) != 1 or node.is_output: + return None + + parent = node.parents[0] + + if node.target == "aten._assert_tensor_metadata.default": + return parent.name + + if node.target in {"aten.clone.default", "aten.contiguous.default", "aten.detach.default"}: + return parent.name + + if node.target in {"aten._to_copy.default", "aten.view.default", "aten.reshape.default", "aten.expand.default", "cactus.view", "cactus.reshape"} and same_tensor_signature(node, parent): + return parent.name + + return None + + +def resolve_replacement_names(replacement_names: dict[str, str]) -> dict[str, str]: + resolved: dict[str, str] = {} + + for name in replacement_names: + current = replacement_names[name] + seen = {name} + + while current in replacement_names and current not in seen: + seen.add(current) + current = replacement_names[current] + + resolved[name] = current + + return resolved + + +def same_tensor_signature(left: Node, right: Node) -> bool: + return same_shape(left, right) and same_effective_dtype(left, right) + + +def same_shape(left: Node, right: Node) -> bool: + return tensor_meta_value(left, "shape") == tensor_meta_value(right, "shape") + + +def same_effective_dtype(left: Node, right: Node) -> bool: + return effective_dtype(tensor_meta_value(left, "dtype")) == effective_dtype(tensor_meta_value(right, "dtype")) + + +def tensor_meta_value(node: Node, key: str) -> Any: + if not isinstance(node.tensor_output_meta, dict): + return None + + return node.tensor_output_meta.get(key) + + +def effective_dtype(dtype: Any) -> str | None: + if dtype is None: + return None + + dtype_name = str(dtype) + + if dtype_name in {"torch.float16", "torch.bfloat16"}: + return "torch.float16" + + return dtype_name + + def rebuild_graph(nodes: tuple[Node, ...], original_graph: Graph, fusion_results: tuple[FusionResult, ...] = ()) -> Graph: nodes_map = {node.name: node for node in nodes} parent_lists: dict[str, list[Node]] = {node.name: [] for node in nodes} diff --git a/python/cactus/transpiler/IR/simplify_ir.py b/python/cactus/transpiler/IR/simplify_ir.py index 2852fb3f9..8d8bdf6a6 100644 --- a/python/cactus/transpiler/IR/simplify_ir.py +++ b/python/cactus/transpiler/IR/simplify_ir.py @@ -40,7 +40,7 @@ def simplify( total_fusions += len(simplified_graph.fusions) graph = simplified_graph - return graph.to_layer_map() + return graph.remove_noop_nodes().to_layer_map() def write_simplified_json( diff --git a/python/cactus/transpiler/ModelProfiles/profiles.py b/python/cactus/transpiler/ModelProfiles/profiles.py index 8be35c27d..447a7492a 100644 --- a/python/cactus/transpiler/ModelProfiles/profiles.py +++ b/python/cactus/transpiler/ModelProfiles/profiles.py @@ -33,6 +33,7 @@ "gemma4_attention", "gemma4_mlp", "linear", + "cache", ), supported_modalties=("audio", "vision", "text"), input_strategy="manual_gemma4_processor", diff --git a/python/cactus/transpiler/RuntimePlan/models.py b/python/cactus/transpiler/RuntimePlan/models.py index bfb09c7d2..8edf0b710 100644 --- a/python/cactus/transpiler/RuntimePlan/models.py +++ b/python/cactus/transpiler/RuntimePlan/models.py @@ -239,7 +239,7 @@ def runtime_plan_from_generator_manifests( ) return RuntimePlan( - family=str(getattr(model_profile, "model_profiles", "") or ""), + family=runtime_family_from_model_profile(model_profile), components=components, routes=runtime_routes_from_model_profile(model_profile), metadata=string_dict(metadata), @@ -273,6 +273,19 @@ def runtime_routes_from_model_profile(model_profile: Any | None) -> tuple[Runtim return tuple(runtime_route_from_profile_route(route) for route in routes.values()) +def runtime_family_from_model_profile(model_profile: Any | None) -> str: + profile_name = str(getattr(model_profile, "model_profiles", "") or "") + normalized_name = profile_name.lower() + + if "gemma4" in normalized_name or "gemma_4" in normalized_name: + return "gemma4" + + if normalized_name in {"lfm_vlm", "lfm2_vl", "lfm-vlm"}: + return "lfm2_vl" + + return profile_name + + def runtime_route_from_profile_route(route: Any) -> RuntimeRoute: return RuntimeRoute( name=str(getattr(route, "name", "")), diff --git a/python/cactus/transpiler/main.py b/python/cactus/transpiler/main.py index cacea140a..887587bda 100644 --- a/python/cactus/transpiler/main.py +++ b/python/cactus/transpiler/main.py @@ -36,7 +36,10 @@ def main() -> None: model_id = "google/gemma-4-E2B" input_modalities = ("text", "vision", "audio") output_dir = constants.CONVERTER_JSON_DIR - bundle_dir = Path(__file__).resolve().parent / "Generator" / "graphs" / "gemma4_e2b_bundle" + repo_root = Path(__file__).resolve().parents[3] + weights_arg = sys.argv[1] if len(sys.argv) > 1 else os.environ.get("CACTUS_TRANSPILER_WEIGHTS_DIR") + weights_dir = Path(weights_arg).expanduser() if weights_arg else repo_root / "weights" / "gemma4-e2b" + bundle_dir = weights_dir output_dir.mkdir(parents=True, exist_ok=True) bundle_dir.mkdir(parents=True, exist_ok=True) @@ -62,11 +65,8 @@ def main() -> None: prefill = LayerMap.model_validate_json(prefill_simplified_path.read_text(encoding="utf-8")) decode = LayerMap.model_validate_json(decode_simplified_path.read_text(encoding="utf-8")) - weights_arg = sys.argv[1] if len(sys.argv) > 1 else os.environ.get("CACTUS_TRANSPILER_WEIGHTS_DIR") - weights_dir = Path(weights_arg).expanduser() if weights_arg else None - if weights_dir is not None: - materialize_runtime_bundle_files(weights_dir, bundle_dir) + materialize_runtime_bundle_files(weights_dir, bundle_dir) result = generate_bundle( {"prefill": prefill, "decode": decode}, From e78aeca4789750e0808d1b1274b9b4c11a9b461b Mon Sep 17 00:00:00 2001 From: Parkirat Date: Wed, 29 Jul 2026 23:54:24 -0400 Subject: [PATCH 69/84] performance upgrades --- .../cactus/transpiler/Generator/generate.py | 149 ++++++- python/cactus/transpiler/Generator/models.py | 19 + python/cactus/transpiler/IR/simplify_ir.py | 131 ++++++ .../cactus/transpiler/RuntimePlan/models.py | 35 ++ .../cactus/transpiler/Validation/__init__.py | 9 + python/cactus/transpiler/Validation/golden.py | 383 ++++++++++++++++++ 6 files changed, 724 insertions(+), 2 deletions(-) create mode 100644 python/cactus/transpiler/Validation/__init__.py create mode 100644 python/cactus/transpiler/Validation/golden.py diff --git a/python/cactus/transpiler/Generator/generate.py b/python/cactus/transpiler/Generator/generate.py index 8a5967b73..3267fae4d 100644 --- a/python/cactus/transpiler/Generator/generate.py +++ b/python/cactus/transpiler/Generator/generate.py @@ -2,6 +2,7 @@ import json import math +import shutil import struct from collections.abc import Mapping from dataclasses import replace @@ -50,6 +51,7 @@ def generate( strict=strict, allow_unsupported_ops=allow_unsupported_ops, ) + prepare_generation_output_dir(generator_config.output_dir) components = component_graphs_from_input(ir_output, generator_config, component_name, model_profile) lowering_rules = build_lowering_rules(lowerings) @@ -97,6 +99,8 @@ def generate_bundle( ) -> models.GenerationResult: bundle_path = Path(bundle_dir) components_dir = bundle_path / "components" + source_weights_dir = Path(weights_dir) if weights_dir is not None else None + materialize_runtime_bundle_files(bundle_path, source_weights_dir, model_profile) result = generate( ir_output, output_dir=components_dir, @@ -120,6 +124,81 @@ def generate_bundle( ################################################# Generator Utils!!!!!!! ################################################# +def materialize_runtime_bundle_files(bundle_dir: Path, weights_dir: Path | None, model_profile: Any | None = None) -> None: + bundle_dir.mkdir(parents=True, exist_ok=True) + + if weights_dir is not None and weights_dir.exists(): + for source in weights_dir.iterdir(): + if source.is_file(): + materialize_bundle_file(source, bundle_dir / source.name) + + for source in model_profile_metadata_files(model_profile): + materialize_bundle_file(source, bundle_dir / source.name, overwrite=False) + + +def model_profile_metadata_files(model_profile: Any | None) -> tuple[Path, ...]: + if model_profile is None: + return () + + try: + from ..Converter import constants as converter_constants + except Exception: + return () + + profile_name = getattr(model_profile, "model_profiles", None) + + if not profile_name: + return () + + metadata_dir = converter_constants.CONVERTER_JSON_DIR / str(profile_name) + filenames = runtime_metadata_filenames(model_profile) + return tuple(metadata_dir / filename for filename in filenames if (metadata_dir / filename).is_file()) + + +def runtime_metadata_filenames(model_profile: Any | None) -> tuple[str, ...]: + profile_files = tuple(str(filename) for filename in getattr(model_profile, "files", ()) or ()) + standard_files = ( + "config.json", + "generation_config.json", + "processor_config.json", + "preprocessor_config.json", + "tokenizer_config.json", + "tokenizer.json", + "special_tokens_map.json", + ) + return tuple(models.unique_strings((*standard_files, *profile_files))) + + +def materialize_bundle_file(source: Path, target: Path, *, overwrite: bool = True) -> None: + if not source.is_file(): + return + + if target.exists() or target.is_symlink(): + try: + if target.samefile(source): + return + except OSError: + pass + + if not overwrite: + return + + target.unlink() + + try: + target.symlink_to(source) + except OSError: + shutil.copy2(source, target) + + +def prepare_generation_output_dir(output_dir: Path) -> None: + output_dir.mkdir(parents=True, exist_ok=True) + constants_dir = output_dir / "constants" + + if constants_dir.exists(): + shutil.rmtree(constants_dir) + + def read_layer_map(input_path: str | Path) -> CModels.LayerMap: path = Path(input_path) return CModels.LayerMap.model_validate_json(path.read_text(encoding="utf-8")) @@ -461,6 +540,9 @@ def should_lower_cache_placeholder_as_state(context: models.GenerationContext, n return any(child.target in {"cactus.conv_cache_append", "cactus.conv_cache_initialize"} for child in node.children) if node.cache.kind == FModels.CacheKind.KV: + if any(child.target in {"cactus.cat", "aten.cat.default"} for child in node.children): + return True + return bool(node.children) and all(child.target in {"cactus.kv_cache_append", "cactus.attention_cached"} for child in node.children) return False @@ -531,8 +613,14 @@ def bind_weight_placeholder(context: models.GenerationContext, node: IRModels.No node_id=node_id, path=record.output_name, output_name=record.output_name, + source_name=record.source_name or record.hf_name or source_target, + value_id=node.name, precision=record.precision, component=record.component, + scale_factor=record.scale_factor, + adapter_family=record.adapter_family, + transform=record.transform, + qdq_restore=record.qdq_restore, ) ) @@ -1053,6 +1141,11 @@ def lower_cat(context: models.GenerationContext, node: IRModels.Node) -> Any: if prefill_cache is not None: return prefill_cache + decode_cache = lower_decode_cache_cat(context, node) + + if decode_cache is not None: + return decode_cache + inputs = context.inputs_for(node) if not inputs: @@ -1067,6 +1160,18 @@ def lower_cat(context: models.GenerationContext, node: IRModels.Node) -> Any: return context.graph.cat(inputs, axis=axis) +def lower_decode_cache_cat(context: models.GenerationContext, node: IRModels.Node) -> Any | None: + if len(node.parents) < 2: + return None + + cache_parent = node.parents[0] + + if cache_parent.cache is None or cache_parent.cache.kind != FModels.CacheKind.KV: + return None + + return context.require(cache_parent.name) + + def lower_prefill_cache_cat(context: models.GenerationContext, node: IRModels.Node) -> Any | None: annotation = context.prefill_cache_cat_annotations.get(node.name) @@ -1245,17 +1350,33 @@ def lower_to_copy(context: models.GenerationContext, node: IRModels.Node) -> Any if precision_name is None: return inputs[0] - return context.graph.precision_cast(inputs[0], cactus_precision(context.graph, str(dtype))) + return cast_to_precision(context, inputs[0], cactus_precision(context.graph, str(dtype))) def can_skip_float32_copy(node: IRModels.Node) -> bool: - if not node.children: + if not node.children or not node.parents: + return False + + source_dtype = tensor_dtype(node.parents[0]) + + if constants.DTYPE_TO_PRECISION.get(str(source_dtype)) != "FP16": return False safe_consumers = { "aten.bmm.default", "cactus.attention", + "cactus.attention_cached", + "cactus.add", + "cactus.dense_mlp_tq_fused", + "cactus.multiply", + "cactus.pow", + "cactus.rope", "cactus.rms_norm", + "cactus.scalar_multiply", + "cactus.silu", + "cactus.slice", + "cactus.transpose", + "cactus.view", } return all(child.target in safe_consumers for child in node.children) @@ -1326,6 +1447,9 @@ def lower_constant_pad_nd(context: models.GenerationContext, node: IRModels.Node def lower_constant_input(context: models.GenerationContext, node: IRModels.Node) -> Any: + if node.target == "aten.full_like.default": + return lower_full_like(context, node) + constant = deterministic_constant_values(node) if constant is not None: @@ -1345,6 +1469,8 @@ def lower_constant_input(context: models.GenerationContext, node: IRModels.Node) node_id=node_id, path=path, output_name=path, + source_name=node.target, + value_id=node.name, precision=precision_name(context.graph, dtype), component=context.component.name, binding_kind="generated_constant", @@ -1360,6 +1486,25 @@ def lower_constant_input(context: models.GenerationContext, node: IRModels.Node) return tensor +def lower_full_like(context: models.GenerationContext, node: IRModels.Node) -> Any: + inputs = require_input_count(context, node, 1) + fill_value = scalar_attr(node, "fill_value") + + if fill_value is None: + fill_value = scalar_attr(node, "arg_1") + + if fill_value is None: + fill_value = 0.0 + + output = context.graph.scalar_multiply(inputs[0], 0.0) + fill_float = float(fill_value) + + if fill_float != 0.0: + output = context.graph.scalar_add(output, fill_float) + + return cast_to_precision(context, output, cactus_precision(context.graph, tensor_dtype(node))) + + def deterministic_constant_values(node: IRModels.Node) -> tuple[list[float], tuple[int, ...]] | None: shape = concrete_shape(meta_shape(node)) diff --git a/python/cactus/transpiler/Generator/models.py b/python/cactus/transpiler/Generator/models.py index 802bb6c1f..3afa71776 100644 --- a/python/cactus/transpiler/Generator/models.py +++ b/python/cactus/transpiler/Generator/models.py @@ -159,8 +159,14 @@ class WeightBinding: node_id: int path: str output_name: str + source_name: str | None = None + value_id: str | None = None precision: str | None = None component: str | None = None + scale_factor: float = 1.0 + adapter_family: str | None = None + transform: str = "none" + qdq_restore: str = "hf_key" binding_kind: str = "mmap_weight" @@ -558,16 +564,29 @@ def component_manifest_to_dict(manifest: ComponentGraphManifest) -> dict[str, An "node_id": binding.node_id, "path": binding.path, "output_name": binding.output_name, + "source_name": binding.source_name, + "value_id": binding.value_id, "precision": binding.precision, "component": binding.component, + "scale_factor": binding.scale_factor, + "adapter_family": binding.adapter_family, + "transform": binding.transform, + "qdq_restore": binding.qdq_restore, "binding_kind": binding.binding_kind, } for binding in manifest.weight_bindings ], "bound_constant_bindings": [ { + "kind": "weight", "node_id": binding.node_id, "path": binding.path, + "source_name": binding.source_name or binding.source_target, + "value_id": binding.value_id or binding.placeholder, + "scale_factor": binding.scale_factor, + "adapter_family": binding.adapter_family, + "transform": binding.transform, + "qdq_restore": binding.qdq_restore, } for binding in manifest.weight_bindings ], diff --git a/python/cactus/transpiler/IR/simplify_ir.py b/python/cactus/transpiler/IR/simplify_ir.py index 8d8bdf6a6..94900ed94 100644 --- a/python/cactus/transpiler/IR/simplify_ir.py +++ b/python/cactus/transpiler/IR/simplify_ir.py @@ -40,6 +40,7 @@ def simplify( total_fusions += len(simplified_graph.fusions) graph = simplified_graph + graph = fuse_gemma4_decode_attention(graph, layer_map.task, fusion_fields) return graph.remove_noop_nodes().to_layer_map() @@ -270,3 +271,133 @@ def unique_nodes(nodes: Iterable[models.Node]) -> tuple[models.Node, ...]: unique.append(node) return tuple(unique) + + +def fuse_gemma4_decode_attention( + graph: models.Graph, + inference_mode: str | None, + fusion_fields: tuple[str, ...] = (), +) -> models.Graph: + if inference_mode != "decode_with_cache": + return graph + + if fusion_fields and "gemma4_attention" not in fusion_fields and "attention" not in fusion_fields: + return graph + + fusion = Fusions.FUSIONS["attention_masked"] + results: list[models.FusionResult] = [] + + for source in graph.nodes: + match_result = gemma4_decode_attention_match(graph, source, fusion) + + if match_result is not None: + results.append(match_result) + + if not results: + return graph + + return graph.apply_fusions(results) + + +def gemma4_decode_attention_match( + graph: models.Graph, + source: models.Node, + fusion: FModels.FusionDefinition, +) -> models.FusionResult | None: + if source.target != "cactus.view" or len(source.parents) != 1: + return None + + value_bmm = source.parents[0] + + if value_bmm.target != "aten.bmm.default" or len(value_bmm.parents) < 2: + return None + + value_cat = cache_cat_ancestor(value_bmm.parents[1], FModels.CacheTensorRole.VALUE) + + if value_cat is None: + return None + + qk_bmm = first_bmm_ancestor(value_bmm.parents[0], value_bmm) + + if qk_bmm is None or len(qk_bmm.parents) < 2: + return None + + key_cat = cache_cat_ancestor(qk_bmm.parents[1], FModels.CacheTensorRole.KEY) + query = first_rank4_single_parent_ancestor(qk_bmm.parents[0]) + + if key_cat is None or query is None: + return None + + return models.FusionResult.from_match( + fusion=fusion, + source=source, + matched_nodes=(source,), + external_inputs=(query, key_cat, value_cat), + attrs={ + "scale": 1.0, + "input_layout": "bhqd_bhds_bhsd", + "output_layout": "bhqd", + "window_size": 0, + }, + ) + + +def cache_cat_ancestor(node: models.Node, role: str, max_depth: int = 18) -> models.Node | None: + current = node + + for _ in range(max_depth): + if current.target in {"cactus.cat", "aten.cat.default"} and len(current.parents) >= 2: + cache_parent = current.parents[0] + + if cache_parent.cache is not None and cache_parent.cache.kind == FModels.CacheKind.KV and cache_parent.cache.role == role: + return current + + if len(current.parents) != 1: + return None + + current = current.parents[0] + + return None + + +def first_bmm_ancestor(node: models.Node, excluded: models.Node) -> models.Node | None: + stack = [node] + seen: set[str] = set() + + while stack: + current = stack.pop() + + if current.name in seen: + continue + + seen.add(current.name) + + if current is not excluded and current.target == "aten.bmm.default": + return current + + stack.extend(current.parents) + + return None + + +def first_rank4_single_parent_ancestor(node: models.Node, max_depth: int = 10) -> models.Node | None: + current = node + + for _ in range(max_depth): + if node_rank(current) == 4: + return current + + if len(current.parents) != 1: + return None + + current = current.parents[0] + + return None + + +def node_rank(node: models.Node) -> int: + if not isinstance(node.tensor_output_meta, dict): + return 0 + + shape = node.tensor_output_meta.get("shape") + return len(shape) if isinstance(shape, list) else 0 diff --git a/python/cactus/transpiler/RuntimePlan/models.py b/python/cactus/transpiler/RuntimePlan/models.py index 8edf0b710..3e4843dc2 100644 --- a/python/cactus/transpiler/RuntimePlan/models.py +++ b/python/cactus/transpiler/RuntimePlan/models.py @@ -17,6 +17,13 @@ class RuntimeTensorBinding: class ConstantBinding: node_id: int path: str + kind: str = "weight" + source_name: str | None = None + value_id: str | None = None + scale_factor: float = 1.0 + adapter_family: str | None = None + transform: str = "none" + qdq_restore: str = "hf_key" @dataclass(slots=True) @@ -179,13 +186,27 @@ def constant_binding_from_dict(data: dict[str, Any]) -> ConstantBinding: return ConstantBinding( node_id=int(data["node_id"]), path=str(data["path"]), + kind=str(data.get("kind", "weight")), + source_name=none_or_str(data.get("source_name")), + value_id=none_or_str(data.get("value_id")), + scale_factor=float(data.get("scale_factor", 1.0) or 1.0), + adapter_family=none_or_str(data.get("adapter_family")), + transform=str(data.get("transform", "none") or "none"), + qdq_restore=str(data.get("qdq_restore", "hf_key") or "hf_key"), ) def constant_binding_to_dict(binding: ConstantBinding) -> dict[str, Any]: return { + "kind": binding.kind, "node_id": binding.node_id, "path": binding.path, + "source_name": binding.source_name, + "value_id": binding.value_id, + "scale_factor": binding.scale_factor, + "adapter_family": binding.adapter_family, + "transform": binding.transform, + "qdq_restore": binding.qdq_restore, } @@ -404,6 +425,13 @@ def constant_bindings_from_generator_manifest(data: dict[str, Any]) -> tuple[Con ConstantBinding( node_id=int(value["node_id"]), path=str(value["path"]), + kind=str(value.get("binding_kind", "weight")), + source_name=none_or_str(value.get("source_name") or value.get("source_target")), + value_id=none_or_str(value.get("value_id") or value.get("placeholder")), + scale_factor=float(value.get("scale_factor", 1.0) or 1.0), + adapter_family=none_or_str(value.get("adapter_family")), + transform=str(value.get("transform", "none") or "none"), + qdq_restore=str(value.get("qdq_restore", "hf_key") or "hf_key"), ) for value in data.get("weight_bindings", ()) if isinstance(value, dict) and "node_id" in value and "path" in value @@ -421,6 +449,13 @@ def string_dict(data: Any) -> dict[str, str]: } +def none_or_str(value: Any) -> str | None: + if value is None: + return None + + return str(value) + + def unique_ints(values: tuple[int, ...]) -> tuple[int, ...]: unique: list[int] = [] diff --git a/python/cactus/transpiler/Validation/__init__.py b/python/cactus/transpiler/Validation/__init__.py new file mode 100644 index 000000000..7cde5b6d5 --- /dev/null +++ b/python/cactus/transpiler/Validation/__init__.py @@ -0,0 +1,9 @@ +from .golden import ArrayComparison +from .golden import TextReference +from .golden import audit_bundle_bindings +from .golden import compare_arrays +from .golden import compare_component_to_reference +from .golden import compare_hf_and_cactus_text +from .golden import run_cactus_component +from .golden import run_cactus_text_smoke +from .golden import run_hf_text_reference diff --git a/python/cactus/transpiler/Validation/golden.py b/python/cactus/transpiler/Validation/golden.py new file mode 100644 index 000000000..51bfd51de --- /dev/null +++ b/python/cactus/transpiler/Validation/golden.py @@ -0,0 +1,383 @@ +from __future__ import annotations + +import json +import os +import subprocess +import sys +import tempfile +from dataclasses import dataclass +from pathlib import Path +from typing import Any, Mapping + +import numpy as np + + +@dataclass(slots=True, frozen=True) +class ArrayComparison: + name: str + shape: tuple[int, ...] + reference_shape: tuple[int, ...] + max_abs_error: float + mean_abs_error: float + topk_match: bool | None = None + passed: bool = False + + +@dataclass(slots=True, frozen=True) +class TextReference: + prompt: str + input_ids: tuple[int, ...] + generated_ids: tuple[int, ...] + text: str + logits: np.ndarray | None = None + + +def run_cactus_component( + bundle_dir: str | Path, + component_name: str, + inputs: Mapping[str, Any], +) -> dict[str, np.ndarray]: + bundle = Path(bundle_dir) + component = load_component_manifest(bundle, component_name) + graph = load_component_graph(bundle, component) + bind_component_weights(graph, bundle, component) + set_component_inputs(graph, component, inputs) + graph.execute() + return read_component_outputs(graph, component) + + +def compare_component_to_reference( + bundle_dir: str | Path, + component_name: str, + inputs: Mapping[str, Any], + reference_outputs: Mapping[str, Any], + *, + atol: float = 1e-2, + rtol: float = 1e-2, + topk: int = 5, +) -> dict[str, ArrayComparison]: + cactus_outputs = run_cactus_component(bundle_dir, component_name, inputs) + return { + name: compare_arrays( + name, + cactus_outputs[name], + reference, + atol=atol, + rtol=rtol, + topk=topk, + ) + for name, reference in reference_outputs.items() + if name in cactus_outputs + } + + +def compare_arrays( + name: str, + actual: Any, + reference: Any, + *, + atol: float = 1e-2, + rtol: float = 1e-2, + topk: int = 5, +) -> ArrayComparison: + actual_array = np.asarray(actual) + reference_array = np.asarray(reference) + + if actual_array.shape != reference_array.shape: + return ArrayComparison( + name=name, + shape=tuple(int(dim) for dim in actual_array.shape), + reference_shape=tuple(int(dim) for dim in reference_array.shape), + max_abs_error=float("inf"), + mean_abs_error=float("inf"), + topk_match=None, + passed=False, + ) + + diff = np.abs(actual_array.astype(np.float32) - reference_array.astype(np.float32)) + max_abs_error = float(diff.max()) if diff.size else 0.0 + mean_abs_error = float(diff.mean()) if diff.size else 0.0 + topk_match = compare_topk(actual_array, reference_array, topk) + passed = bool(np.allclose(actual_array, reference_array, atol=atol, rtol=rtol)) + + return ArrayComparison( + name=name, + shape=tuple(int(dim) for dim in actual_array.shape), + reference_shape=tuple(int(dim) for dim in reference_array.shape), + max_abs_error=max_abs_error, + mean_abs_error=mean_abs_error, + topk_match=topk_match, + passed=passed, + ) + + +def run_cactus_text_smoke( + bundle_dir: str | Path, + prompt: str, + *, + max_new_tokens: int = 1, + python_executable: str | Path | None = None, + env: Mapping[str, str] | None = None, +) -> dict[str, Any]: + with tempfile.NamedTemporaryFile(prefix="cactus_golden_", suffix=".json", delete=False) as result_file: + result_path = Path(result_file.name) + + command = [ + str(python_executable or sys.executable), + "-m", + "cactus", + "run", + str(bundle_dir), + "--prompt", + prompt, + "--max-new-tokens", + str(int(max_new_tokens)), + "--no-cloud-handoff", + "--result-json", + str(result_path), + ] + run_env = dict(os.environ) + if env is not None: + run_env.update({str(key): str(value) for key, value in env.items()}) + + try: + completed = subprocess.run(command, env=run_env, check=False, capture_output=True, text=True) + payload = json.loads(result_path.read_text(encoding="utf-8")) if result_path.exists() else {} + payload["returncode"] = completed.returncode + payload["stdout"] = completed.stdout + payload["stderr"] = completed.stderr + return payload + finally: + result_path.unlink(missing_ok=True) + + +def run_hf_text_reference( + model_id_or_path: str | Path, + prompt: str, + *, + max_new_tokens: int = 1, + apply_chat_template: bool = True, + device: str = "cpu", +) -> TextReference: + import torch + from transformers import AutoModelForCausalLM, AutoProcessor, AutoTokenizer + + model_path = str(model_id_or_path) + try: + tokenizer = AutoTokenizer.from_pretrained(model_path) + except Exception: + processor = AutoProcessor.from_pretrained(model_path) + tokenizer = getattr(processor, "tokenizer", processor) + + model = load_hf_generation_model(model_path).to(device) + model.eval() + + text = prompt + if apply_chat_template and getattr(tokenizer, "chat_template", None): + text = tokenizer.apply_chat_template( + [{"role": "user", "content": prompt}], + tokenize=False, + add_generation_prompt=True, + ) + + encoded = tokenizer(text, return_tensors="pt") + encoded = {key: value.to(device) for key, value in encoded.items()} + + with torch.no_grad(): + outputs = model(**encoded) + generated = model.generate(**encoded, max_new_tokens=int(max_new_tokens), do_sample=False) + + input_ids = tuple(int(value) for value in encoded["input_ids"][0].detach().cpu().tolist()) + generated_ids = tuple(int(value) for value in generated[0].detach().cpu().tolist()) + decoded = tokenizer.decode(generated[0], skip_special_tokens=False) + logits = outputs.logits[:, -1, :].detach().cpu().float().numpy() + return TextReference(prompt=prompt, input_ids=input_ids, generated_ids=generated_ids, text=decoded, logits=logits) + + +def compare_hf_and_cactus_text( + bundle_dir: str | Path, + model_id_or_path: str | Path, + prompt: str, + *, + max_new_tokens: int = 1, + apply_chat_template: bool = True, + device: str = "cpu", +) -> dict[str, Any]: + hf = run_hf_text_reference( + model_id_or_path, + prompt, + max_new_tokens=max_new_tokens, + apply_chat_template=apply_chat_template, + device=device, + ) + cactus = run_cactus_text_smoke(bundle_dir, prompt, max_new_tokens=max_new_tokens) + return { + "hf": { + "input_ids": list(hf.input_ids), + "generated_ids": list(hf.generated_ids), + "text": hf.text, + }, + "cactus": cactus, + "tokenizer_parity_note": "Use hf.input_ids with `cactus run --input-ids` to isolate tokenizer/chat-template differences from graph math differences.", + } + + +def audit_bundle_bindings(bundle_dir: str | Path) -> list[str]: + bundle = Path(bundle_dir) + manifest_path = bundle / "components" / "manifest.json" + + if not manifest_path.exists(): + return [f"Missing component manifest: {manifest_path}"] + + data = json.loads(manifest_path.read_text(encoding="utf-8")) + issues: list[str] = [] + gemma_scaled = 0 + gemma_per_layer = 0 + + for component in data.get("components", ()): + component_name = str(component.get("component") or component.get("name") or component.get("id") or "unknown") + + for index, binding in enumerate(component.get("bound_constant_bindings", ())): + if not isinstance(binding, Mapping): + issues.append(f"{component_name}[{index}]: binding is not an object") + continue + + label = f"{component_name}[{index}] node_id={binding.get('node_id')}" + source_name = binding.get("source_name") + value_id = binding.get("value_id") + scale_factor = float(binding.get("scale_factor", 1.0) or 1.0) + + if not source_name: + issues.append(f"{label}: missing source_name") + + if not value_id: + issues.append(f"{label}: missing value_id") + + if scale_factor != 1.0: + gemma_scaled += 1 + + source_text = str(source_name or binding.get("path") or "") + if "per_layer" in source_text or "embed_tokens_per_layer" in source_text: + gemma_per_layer += 1 + + if gemma_scaled: + issues.append(f"Gemma scaled bindings detected: {gemma_scaled}") + + if "gemma" in bundle.name.lower() and gemma_per_layer == 0: + issues.append("Gemma bundle has no per-layer/PLI binding metadata; regenerate with source_name/value_id metadata before judging correctness") + + return issues + + +def load_hf_generation_model(model_path: str) -> Any: + from transformers import AutoModelForCausalLM + + candidate_classes: list[Any] = [AutoModelForCausalLM] + + try: + from transformers import AutoModelForImageTextToText + candidate_classes.append(AutoModelForImageTextToText) + except Exception: + pass + + try: + from transformers import AutoModelForSeq2SeqLM + candidate_classes.append(AutoModelForSeq2SeqLM) + except Exception: + pass + + last_error: Exception | None = None + + for model_class in candidate_classes: + try: + return model_class.from_pretrained(model_path) + except Exception as exc: + last_error = exc + + raise RuntimeError(f"Unable to load an HF generation model from {model_path}") from last_error + + +def load_component_manifest(bundle_dir: Path, component_name: str) -> dict[str, Any]: + manifest_path = bundle_dir / "components" / "manifest.json" + data = json.loads(manifest_path.read_text(encoding="utf-8")) + + for component in data.get("components", ()): + name = component.get("component") or component.get("name") or component.get("id") + + if name == component_name: + return component + + raise KeyError(f"Component {component_name!r} was not found in {manifest_path}") + + +def load_component_graph(bundle_dir: Path, component: Mapping[str, Any]) -> Any: + graph_path = bundle_dir / str(component.get("graph") or component.get("graph_path")) + graph_cls = cactus_graph_class() + return graph_cls.load(graph_path) + + +def bind_component_weights(graph: Any, bundle_dir: Path, component: Mapping[str, Any]) -> None: + for binding in component.get("bound_constant_bindings", ()): + if not isinstance(binding, Mapping): + continue + + node_id = int(binding["node_id"]) + path = resolve_bundle_path(bundle_dir, str(binding["path"])) + tensor = graph._tensor_from_node(node_id) + graph.bind_mmap_weights(tensor, path) + + +def set_component_inputs(graph: Any, component: Mapping[str, Any], inputs: Mapping[str, Any]) -> None: + node_ids = tuple(int(value) for value in component.get("runtime_input_node_ids", ())) + logical_inputs = tuple(str(value) for value in component.get("logical_inputs", ())) + node_by_name = dict(zip(logical_inputs, node_ids)) + + for name, value in inputs.items(): + if name not in node_by_name: + raise KeyError(f"Input {name!r} is not declared by component {component.get('component')!r}") + + graph.set_input(graph._tensor_from_node(node_by_name[name]), value) + + +def read_component_outputs(graph: Any, component: Mapping[str, Any]) -> dict[str, np.ndarray]: + node_ids = tuple(int(value) for value in component.get("output_node_ids", ())) + logical_outputs = tuple(str(value) for value in component.get("logical_outputs", ())) + outputs: dict[str, np.ndarray] = {} + + for name, node_id in zip(logical_outputs, node_ids): + outputs[name] = graph._tensor_from_node(node_id).numpy() + + return outputs + + +def compare_topk(actual: np.ndarray, reference: np.ndarray, topk: int) -> bool | None: + if topk <= 0 or actual.size == 0: + return None + + actual_flat = actual.reshape(-1, actual.shape[-1]) if actual.ndim > 1 else actual.reshape(1, -1) + reference_flat = reference.reshape(-1, reference.shape[-1]) if reference.ndim > 1 else reference.reshape(1, -1) + + if actual_flat.shape != reference_flat.shape or actual_flat.shape[-1] < topk: + return None + + actual_top = np.argpartition(actual_flat, -topk, axis=-1)[:, -topk:] + reference_top = np.argpartition(reference_flat, -topk, axis=-1)[:, -topk:] + return all(set(map(int, left)) == set(map(int, right)) for left, right in zip(actual_top, reference_top)) + + +def resolve_bundle_path(bundle_dir: Path, path: str) -> Path: + candidate = Path(path) + + if candidate.is_absolute(): + return candidate + + return bundle_dir / candidate + + +def cactus_graph_class() -> Any: + try: + from cactus.bindings.cactus import Graph + except ModuleNotFoundError: + from bindings.cactus import Graph + + return Graph From 909da19b3430b8d77731e2edd348180d9dd8a2fd Mon Sep 17 00:00:00 2001 From: Parkirat Date: Fri, 31 Jul 2026 10:15:38 -0400 Subject: [PATCH 70/84] Performace upgrades --- cactus-engine/src/complete.cpp | 4 + cactus-engine/src/engine.h | 1 + cactus-engine/src/model.cpp | 245 +++++- cactus-graph/cactus_graph.h | 13 +- cactus-graph/src/builder.cpp | 26 + cactus-graph/src/execute.cpp | 7 +- cactus-graph/src/graph_ffi.cpp | 68 ++ cactus-graph/src/io.cpp | 6 +- cactus-graph/src/ops_conv.cpp | 7 +- cactus-graph/src/ops_math.cpp | 24 +- cactus-graph/src/ops_nn.cpp | 7 +- cactus-graph/src/ops_tensor.cpp | 47 ++ cactus-kernels/src/nn.cpp | 5 + python/cactus/bindings/cactus.py | 43 + .../transpiler/Converter/input_processor.py | 22 +- python/cactus/transpiler/Converter/models.py | 50 +- python/cactus/transpiler/Fusions/edges.py | 24 + python/cactus/transpiler/Fusions/fusions.py | 166 +++- python/cactus/transpiler/Fusions/nodes.py | 70 +- .../transpiler/Generator/component_split.py | 187 +++- .../cactus/transpiler/Generator/constants.py | 37 + .../cactus/transpiler/Generator/generate.py | 515 ++++++++++- python/cactus/transpiler/Generator/models.py | 18 + python/cactus/transpiler/IR/extra_matchers.py | 102 +++ python/cactus/transpiler/IR/models.py | 47 +- python/cactus/transpiler/IR/simplify_ir.py | 173 +--- .../cactus/transpiler/IR/special_fusions.py | 797 ++++++++++++++++++ .../transpiler/ModelProfiles/profiles.py | 4 +- .../cactus/transpiler/RuntimePlan/models.py | 25 +- python/tests/test_runtime_plan.py | 2 +- 30 files changed, 2517 insertions(+), 225 deletions(-) create mode 100644 python/cactus/transpiler/IR/special_fusions.py diff --git a/cactus-engine/src/complete.cpp b/cactus-engine/src/complete.cpp index 6d0d99ddb..7baa740b3 100644 --- a/cactus-engine/src/complete.cpp +++ b/cactus-engine/src/complete.cpp @@ -527,6 +527,8 @@ PreparedPrompt prepare_prompt( auto samples_16k = resample_to_16k_fp32(waveform_fp32, 16000); if (!samples_16k.empty()) { auto audio_prep = cactus::audio::preprocess_audio_for_gemma4(samples_16k, handle->model->get_config()); + size_t graph_soft_tokens = handle->model->audio_soft_token_count_for_frames(audio_prep.num_frames); + if (graph_soft_tokens > 0) audio_prep.num_soft_tokens = std::min(audio_prep.num_soft_tokens, graph_soft_tokens); prompt.audio_features.push_back(std::move(audio_prep.features)); size_t u = user_indices.size() - 1; prompt.messages[user_indices[u]].audio_soft_token_count = audio_prep.num_soft_tokens; @@ -541,6 +543,8 @@ PreparedPrompt prepare_prompt( auto samples_16k = resample_to_16k_fp32(wav.samples, wav.sample_rate); if (samples_16k.empty()) continue; auto audio_prep = cactus::audio::preprocess_audio_for_gemma4(samples_16k, handle->model->get_config()); + size_t graph_soft_tokens = handle->model->audio_soft_token_count_for_frames(audio_prep.num_frames); + if (graph_soft_tokens > 0) audio_prep.num_soft_tokens = std::min(audio_prep.num_soft_tokens, graph_soft_tokens); prompt.audio_features.push_back(std::move(audio_prep.features)); prompt.messages[user_indices[u]].audio_soft_token_count = audio_prep.num_soft_tokens; counts[u] = audio_prep.num_soft_tokens; diff --git a/cactus-engine/src/engine.h b/cactus-engine/src/engine.h index 8c9887071..53ec2b0b9 100644 --- a/cactus-engine/src/engine.h +++ b/cactus-engine/src/engine.h @@ -622,6 +622,7 @@ class Model { ~Model(); const Config& get_config() const { return config_; } + size_t audio_soft_token_count_for_frames(size_t frame_count); Tokenizer* get_tokenizer() const { return tokenizer_.get(); } const std::vector& get_debug_nodes() const; diff --git a/cactus-engine/src/model.cpp b/cactus-engine/src/model.cpp index da02085c6..1e97380c3 100644 --- a/cactus-engine/src/model.cpp +++ b/cactus-engine/src/model.cpp @@ -543,6 +543,7 @@ bool Model::init(const std::string& bundle_dir, size_t context_size, "vision_encoder", "audio_encoder", "lm_encoder_media_step", + "lm_encoder_media_chunk", "decoder_prefill_chunk", "decoder_embed_chunk", "lm_encoder", @@ -1182,6 +1183,40 @@ size_t Model::component_output_tokens(const Component& comp, const std::string& return 0; } +size_t Model::audio_soft_token_count_for_frames(size_t frame_count) { + if (!audio_encoder_) return 0; + if (!load_component_graph(*audio_encoder_)) return 0; + + int feature_idx = -1; + for (const auto& name : {"input_features", "audio_features"}) { + int idx = input_index(*audio_encoder_, name); + if (idx >= 0) { + feature_idx = idx; + break; + } + } + if (feature_idx < 0 || audio_encoder_->output_node_ids.empty()) return 0; + + const auto& input_desc = audio_encoder_->graph->get_output_buffer( + static_cast(audio_encoder_->runtime_input_node_ids[feature_idx])); + const auto& output_desc = audio_encoder_->graph->get_output_buffer( + static_cast(audio_encoder_->output_node_ids.front())); + + const size_t frames_per_chunk = input_desc.shape.size() >= 2 ? input_desc.shape[1] : 0; + size_t rows_per_chunk = 0; + if (output_desc.shape.size() >= 2) { + rows_per_chunk = output_desc.shape[output_desc.shape.size() - 2]; + } else if (!output_desc.shape.empty()) { + rows_per_chunk = output_desc.shape[0]; + } + if (frames_per_chunk == 0 || rows_per_chunk == 0) return 0; + + const size_t chunks = frame_count == 0 + ? 1 + : (frame_count + frames_per_chunk - 1) / frames_per_chunk; + return chunks * rows_per_chunk; +} + void Model::execute_prefill_chunk(Component& chunk_comp, Component* enc_comp, size_t encoder_chunk, size_t chunk_tokens, const std::vector& tokens, size_t processed, size_t start_position) { @@ -1210,6 +1245,10 @@ void Model::execute_prefill_chunk(Component& chunk_comp, Component* enc_comp, si copy_component_outputs_to_chunk_inputs(*encoder_, chunk_comp, i); } } + const size_t real_chunk_tokens = tokens.size() > processed ? std::min(chunk_tokens, tokens.size() - processed) : 0; + for (size_t i = 0; i < real_chunk_tokens; ++i) { + write_int_input_at(chunk_comp, "attention_mask", i, 1); + } chunk_comp.graph->execute(); } @@ -1969,8 +2008,10 @@ bool Model::prefill_and_sample_first_token(const std::vector& tokens, ChunkedPrefillResult chunked; if (decoder_prefill_) { chunked = run_chunked_prefill(tokens, cache_total_seq_len_, get_prefill_chunk_size(), true); - if (chunked.logical_tokens == tokens.size() && chunked.padding_tokens > 0 && tokens.size() > 0) { - // Cache already moved into the step component; drop the padded last row and re-run it for logits. + if (chunked.logical_tokens == tokens.size() + && tokens.size() > 0 + && (chunked.padding_tokens > 0 || family_ == "gemma4")) { + // Cache already moved into the step component; drop the last row and re-run it for step-decoder logits. set_cache_current_len(*decoder_, tokens.size() - 1); cache_total_seq_len_ = tokens.size() - 1; run_step(tokens.back(), cache_total_seq_len_, true); @@ -2371,7 +2412,9 @@ bool Model::build_lm_encoder_outputs_dynamic_gemma4( struct OutputInfo { std::string name; int text_idx = -1; + int text_chunk_idx = -1; int media_idx = -1; + int media_chunk_idx = -1; size_t per_token_bytes = 0; Precision precision = Precision::FP16; std::vector shape_template; @@ -2408,6 +2451,181 @@ bool Model::build_lm_encoder_outputs_dynamic_gemma4( store_shape[info.name] = std::move(shape); } + if (lm_encoder_text_chunk_ && lm_encoder_media_chunk_ + && load_component_graph(*lm_encoder_text_chunk_) + && load_component_graph(*lm_encoder_media_chunk_)) { + if (std::getenv("CACTUS_PROFILE_PREFILL")) { + size_t audio_token_count = 0; + size_t image_token_count = 0; + for (uint32_t token : tokens) { + if (audio_tok != 0 && token == audio_tok) ++audio_token_count; + if (image_tok != 0 && token == image_tok) ++image_token_count; + } + std::cerr << "[prefill-profile] gemma4 media tokens: image_tokens=" + << image_token_count << " image_rows=" << image_rows + << " audio_tokens=" << audio_token_count << " audio_rows=" << audio_rows + << std::endl; + } + const size_t text_chunk_tokens = component_chunk_tokens(*lm_encoder_text_chunk_, "input_ids"); + const size_t media_chunk_tokens = component_chunk_tokens(*lm_encoder_media_chunk_, "input_ids"); + const int media_embeds_idx = input_index(*lm_encoder_media_chunk_, "inputs_embeds"); + bool can_chunk = text_chunk_tokens > 0 + && text_chunk_tokens == media_chunk_tokens + && media_embeds_idx >= 0 + && input_index(*lm_encoder_text_chunk_, "input_ids") >= 0 + && input_index(*lm_encoder_text_chunk_, "position_ids") >= 0 + && input_index(*lm_encoder_media_chunk_, "input_ids") >= 0 + && input_index(*lm_encoder_media_chunk_, "position_ids") >= 0; + + for (auto& info : outputs) { + info.text_chunk_idx = output_index(*lm_encoder_text_chunk_, info.name); + info.media_chunk_idx = output_index(*lm_encoder_media_chunk_, info.name); + if (info.text_chunk_idx < 0 || info.media_chunk_idx < 0) { + can_chunk = false; + break; + } + for (Component* chunk_comp : {lm_encoder_text_chunk_, lm_encoder_media_chunk_}) { + int output_idx = chunk_comp == lm_encoder_text_chunk_ ? info.text_chunk_idx : info.media_chunk_idx; + size_t node_id = static_cast(chunk_comp->output_node_ids[output_idx]); + const auto& desc = chunk_comp->graph->get_output_buffer(node_id); + if (desc.precision != info.precision || desc.total_size % media_chunk_tokens != 0) { + can_chunk = false; + break; + } + size_t per_token_elements = desc.total_size / media_chunk_tokens; + size_t per_token_bytes = PrecisionTraits::packed_size_of(desc.precision, per_token_elements); + if (per_token_bytes != info.per_token_bytes) { + can_chunk = false; + break; + } + } + if (!can_chunk) break; + } + + auto write_media_row_at = [&](size_t row, const uint8_t* feature_row, + size_t feature_row_bytes, Precision feature_precision) { + auto& embeds_buf = lm_encoder_media_chunk_->input_buffers[media_embeds_idx]; + size_t node_id = static_cast(lm_encoder_media_chunk_->runtime_input_node_ids[media_embeds_idx]); + const auto& desc = lm_encoder_media_chunk_->graph->get_output_buffer(node_id); + if (media_chunk_tokens == 0 || desc.total_size % media_chunk_tokens != 0) { + throw std::runtime_error("lm_encoder_media_chunk inputs_embeds is not token-aligned"); + } + const size_t row_elements = desc.total_size / media_chunk_tokens; + const size_t dst_offset = PrecisionTraits::byte_offset_of(desc.precision, row * row_elements); + const size_t dst_bytes = PrecisionTraits::packed_size_of(desc.precision, row_elements); + if (dst_offset + dst_bytes > embeds_buf.size()) { + throw std::runtime_error("lm_encoder_media_chunk media row exceeds inputs_embeds buffer"); + } + if (desc.precision == feature_precision) { + size_t to_copy = std::min(feature_row_bytes, dst_bytes); + std::memcpy(embeds_buf.data() + dst_offset, feature_row, to_copy); + if (to_copy < dst_bytes) { + std::memset(embeds_buf.data() + dst_offset + to_copy, 0, dst_bytes - to_copy); + } + return; + } + std::vector row_buf(dst_bytes, 0); + write_typed_buffer(row_buf, desc.precision, feature_row, feature_row_bytes, feature_precision); + std::memcpy(embeds_buf.data() + dst_offset, row_buf.data(), dst_bytes); + }; + + if (can_chunk) { + size_t audio_idx = 0; + size_t image_idx = 0; + for (size_t chunk_start = 0; chunk_start < token_count; chunk_start += media_chunk_tokens) { + const size_t real_tokens = std::min(media_chunk_tokens, token_count - chunk_start); + bool chunk_has_media = false; + bool chunk_has_text = false; + for (size_t i = 0; i < real_tokens; ++i) { + uint32_t token = tokens[chunk_start + i]; + bool is_audio = audio_tok != 0 && token == audio_tok && have_audio_features; + bool is_image = image_tok != 0 && token == image_tok && have_image_features; + if (is_audio || is_image) { + chunk_has_media = true; + } else { + chunk_has_text = true; + } + } + + for (auto& buf : lm_encoder_text_chunk_->input_buffers) { + std::fill(buf.begin(), buf.end(), 0); + } + for (auto& buf : lm_encoder_media_chunk_->input_buffers) { + std::fill(buf.begin(), buf.end(), 0); + } + + if (chunk_has_text) { + for (size_t i = 0; i < media_chunk_tokens; ++i) { + size_t pos = chunk_start + i; + uint32_t token = pos < token_count ? tokens[pos] : static_cast(config_.pad_token_id); + write_int_input_at(*lm_encoder_text_chunk_, "input_ids", i, static_cast(token)); + write_int_input_at(*lm_encoder_text_chunk_, "position_ids", i, static_cast(pos)); + } + + lm_encoder_text_chunk_->graph->execute(); + } + + if (chunk_has_media && chunk_has_text) { + copy_component_outputs_to_inputs(*lm_encoder_text_chunk_, *lm_encoder_media_chunk_); + } + + for (size_t i = 0; i < media_chunk_tokens; ++i) { + size_t pos = chunk_start + i; + uint32_t token = pos < token_count ? tokens[pos] : static_cast(config_.pad_token_id); + uint32_t media_token = token; + + if (pos < token_count && audio_tok != 0 && token == audio_tok && have_audio_features) { + if (audio_idx >= audio_rows) { + throw std::runtime_error("Gemma4 prompt contains more audio tokens than audio feature rows"); + } + const uint8_t* row = audio_it->second.data() + audio_idx * audio_row_bytes; + write_media_row_at(i, row, audio_row_bytes, audio_prec); + media_token = 0; + ++audio_idx; + } else if (pos < token_count && image_tok != 0 && token == image_tok && have_image_features) { + if (image_idx >= image_rows) { + throw std::runtime_error("Gemma4 prompt contains more image tokens than image feature rows"); + } + const uint8_t* row = image_it->second.data() + image_idx * image_row_bytes; + write_media_row_at(i, row, image_row_bytes, image_prec); + media_token = 0; + ++image_idx; + } + + write_int_input_at(*lm_encoder_media_chunk_, "input_ids", i, static_cast(media_token)); + write_int_input_at(*lm_encoder_media_chunk_, "position_ids", i, static_cast(pos)); + } + + Component* output_chunk = lm_encoder_text_chunk_; + if (chunk_has_media) { + lm_encoder_media_chunk_->graph->execute(); + output_chunk = lm_encoder_media_chunk_; + } + + for (const auto& info : outputs) { + int output_idx = output_chunk == lm_encoder_text_chunk_ ? info.text_chunk_idx : info.media_chunk_idx; + size_t node_id = static_cast(output_chunk->output_node_ids[output_idx]); + const auto& desc = output_chunk->graph->get_output_buffer(node_id); + const auto* ptr = static_cast(output_chunk->graph->get_output(node_id)); + size_t per_token_elements = desc.total_size / media_chunk_tokens; + size_t per_token_bytes = PrecisionTraits::packed_size_of(desc.precision, per_token_elements); + for (size_t i = 0; i < real_tokens; ++i) { + size_t src_offset = PrecisionTraits::byte_offset_of(desc.precision, i * per_token_elements); + std::memcpy( + store_bytes[info.name].data() + (chunk_start + i) * info.per_token_bytes, + ptr + src_offset, + per_token_bytes); + } + } + lm_encoder_text_chunk_->graph->release_runtime_buffers(); + lm_encoder_media_chunk_->graph->release_runtime_buffers(); + } + lm_encoder_text_chunk_->graph->release_all_weight_pages(); + lm_encoder_media_chunk_->graph->release_all_weight_pages(); + return true; + } + } + size_t audio_idx = 0; size_t image_idx = 0; for (size_t pos = 0; pos < token_count; ++pos) { @@ -2476,6 +2694,14 @@ bool Model::run_chunk_prefill_path(const std::vector& tokens, const std::vector& image_paths, const std::vector>& audio_features_per_message) { if (cache_total_seq_len_ > 0) return false; + const bool profile_prefill = std::getenv("CACTUS_PROFILE_PREFILL") != nullptr; + auto profile_start = std::chrono::steady_clock::now(); + auto profile_mark = [&](const std::string& label) { + if (!profile_prefill) return; + auto now = std::chrono::steady_clock::now(); + double ms = std::chrono::duration_cast(now - profile_start).count() / 1000.0; + std::cerr << "[prefill-profile] " << label << ": " << ms << " ms" << std::endl; + }; const bool have_images = !image_paths.empty() && vision_encoder_ != nullptr; bool any_audio = false; for (const auto& mel : audio_features_per_message) { if (!mel.empty()) { any_audio = true; break; } } @@ -2585,10 +2811,12 @@ bool Model::run_chunk_prefill_path(const std::vector& tokens, } unload_component_graph(*vision_encoder_); } + profile_mark("media vision encoded"); if (have_audio) { run_audio_encoder_messages(audio_features_per_message); } + profile_mark("media audio encoded"); std::map> store_bytes; std::map store_prec; @@ -2696,6 +2924,7 @@ bool Model::run_chunk_prefill_path(const std::vector& tokens, lm_encoder_->graph->release_all_weight_pages(); unload_component_graph(*lm_encoder_); } + profile_mark("lm encoder outputs built"); auto embeds_shape_it = store_shape.find("inputs_embeds"); if (embeds_shape_it == store_shape.end()) { @@ -2772,8 +3001,13 @@ bool Model::run_chunk_prefill_path(const std::vector& tokens, size_t src_slice_bytes = chunk_seq * src_per_pos; write_typed_buffer(dst_buf, desc.precision, src_ptr, src_slice_bytes, src_prec); } + const size_t real_chunk_tokens = std::min(chunk_seq, valid_seq - chunk_start); + for (size_t i = 0; i < real_chunk_tokens; ++i) { + write_int_input_at(*decoder_prefill_chunk_, "attention_mask", i, 1); + } decoder_prefill_chunk_->graph->execute(); } + profile_mark("decoder chunk prefill executed"); if (whole_chunks_end > 0 && decoder_ != nullptr) { move_cache_states(*decoder_prefill_chunk_, *decoder_); decoder_prefill_chunk_->graph->release_runtime_buffers(); @@ -2794,6 +3028,7 @@ bool Model::run_chunk_prefill_path(const std::vector& tokens, } decoder_->graph->execute(); } + profile_mark("decoder scalar tail executed"); cache_total_seq_len_ += valid_seq; return true; } @@ -2818,8 +3053,12 @@ void Model::prefill_with_media(const std::vector& tokens, return; } + const bool can_chunk_lm_encoder = + lm_encoder_ != nullptr + || (family_ == "gemma4" && lm_encoder_text_chunk_ != nullptr && lm_encoder_media_chunk_ != nullptr); const bool can_chunk_prefill = - lm_encoder_ != nullptr && decoder_prefill_chunk_ != nullptr && + std::getenv("CACTUS_DISABLE_MEDIA_CHUNK_PREFILL") == nullptr && + can_chunk_lm_encoder && decoder_prefill_chunk_ != nullptr && (vision_encoder_ != nullptr || audio_encoder_ != nullptr); if (can_chunk_prefill) { if (run_chunk_prefill_path(tokens, image_paths, audio_features_per_message)) { diff --git a/cactus-graph/cactus_graph.h b/cactus-graph/cactus_graph.h index 7641d2321..6c3f4e2bc 100644 --- a/cactus-graph/cactus_graph.h +++ b/cactus-graph/cactus_graph.h @@ -138,7 +138,8 @@ enum class OpType { PAD, SCALAR_FLOOR_DIVIDE, STRIDED_SLICE, - EXPAND + EXPAND, + MASKED_SELECT_PREFIX }; struct PrecisionTraits { @@ -509,6 +510,7 @@ class CactusGraph { size_t bitwise_and(size_t input1, size_t input2); size_t bitwise_or(size_t input1, size_t input2); size_t where(size_t condition, size_t true_value, size_t false_value); + size_t masked_select_prefix(size_t input, size_t mask); size_t scalar_add(size_t input, float value); @@ -761,6 +763,7 @@ class CactusGraph { const std::vector& output_shape, const OpParams& params = {}); const BufferDesc& get_output_buffer(size_t node_id) const; OpType get_node_op_type(size_t node_id) const; + const std::vector& get_node_inputs(size_t node_id) const; size_t get_node_window_size(size_t node_id) const; size_t get_node_sink_size(size_t node_id) const; size_t get_node_cache_num_slots(size_t node_id) const; @@ -953,6 +956,8 @@ CACTUS_FFI_EXPORT int cactus_graph_bitwise_or(cactus_graph_t graph, cactus_node_ a, cactus_node_t b, cactus_node_t* out); CACTUS_FFI_EXPORT int cactus_graph_where(cactus_graph_t graph, cactus_node_t condition, cactus_node_t true_value, cactus_node_t false_value, cactus_node_t* out); +CACTUS_FFI_EXPORT int cactus_graph_masked_select_prefix(cactus_graph_t graph, cactus_node_t +x, cactus_node_t mask, cactus_node_t* out); CACTUS_FFI_EXPORT int cactus_graph_scalar_add(cactus_graph_t graph, cactus_node_t x, float value, cactus_node_t* out); CACTUS_FFI_EXPORT int cactus_graph_scalar_subtract(cactus_graph_t graph, cactus_node_t x, float value, cactus_node_t* out); @@ -1169,6 +1174,12 @@ CACTUS_FFI_EXPORT int cactus_graph_invalidate_persistent( cactus_graph_t graph, cactus_node_t persistent_node); CACTUS_FFI_EXPORT int cactus_graph_execute(cactus_graph_t graph); +CACTUS_FFI_EXPORT int cactus_graph_retain_outputs(cactus_graph_t graph, +const cactus_node_t* nodes, size_t count); +CACTUS_FFI_EXPORT int cactus_graph_get_node_op_type(cactus_graph_t graph, +cactus_node_t node, int32_t* out_op_type); +CACTUS_FFI_EXPORT int cactus_graph_get_node_inputs(cactus_graph_t graph, +cactus_node_t node, cactus_node_t* out_inputs, size_t* inout_count); CACTUS_FFI_EXPORT int cactus_graph_get_output_ptr(cactus_graph_t graph, cactus_node_t node, void** out_ptr); CACTUS_FFI_EXPORT int cactus_graph_get_output_info(cactus_graph_t graph, diff --git a/cactus-graph/src/builder.cpp b/cactus-graph/src/builder.cpp index ddf2d5b19..10c4c94ec 100644 --- a/cactus-graph/src/builder.cpp +++ b/cactus-graph/src/builder.cpp @@ -59,6 +59,28 @@ size_t CactusGraph::where(size_t condition, size_t true_value, size_t false_valu return add_node(OpType::WHERE, {condition, true_value, false_value}, all_inputs.output_shape, params); } +size_t CactusGraph::masked_select_prefix(size_t input, size_t mask) { + const auto& input_buffer = get_output_buffer(input); + const auto& mask_buffer = get_output_buffer(mask); + if (input_buffer.shape.size() != 2 && input_buffer.shape.size() != 3) { + throw std::runtime_error("masked_select_prefix input must be [T,D] or [1,T,D]"); + } + if (mask_buffer.shape.size() != 1 && mask_buffer.shape.size() != 2) { + throw std::runtime_error("masked_select_prefix mask must be [T] or [1,T]"); + } + + const size_t input_time = input_buffer.shape.size() == 3 ? input_buffer.shape[1] : input_buffer.shape[0]; + const size_t hidden = input_buffer.shape.back(); + const size_t mask_time = mask_buffer.shape.back(); + if (input_time != mask_time) { + throw std::runtime_error("masked_select_prefix input/mask time dimension mismatch"); + } + + OpParams params; + params.output_precision = input_buffer.precision; + return add_node(OpType::MASKED_SELECT_PREFIX, {input, mask}, {input_time, hidden}, params); +} + size_t CactusGraph::abs(size_t input) { const auto& input_buffer = get_output_buffer(input); OpParams params{.output_precision = input_buffer.precision}; @@ -1430,6 +1452,10 @@ OpType CactusGraph::get_node_op_type(size_t node_id) const { return nodes_[node_index_map_.at(node_id)]->op_type; } +const std::vector& CactusGraph::get_node_inputs(size_t node_id) const { + return nodes_[node_index_map_.at(node_id)]->input_ids; +} + size_t CactusGraph::get_node_window_size(size_t node_id) const { return nodes_[node_index_map_.at(node_id)]->params.window_size; } diff --git a/cactus-graph/src/execute.cpp b/cactus-graph/src/execute.cpp index ca79893ae..16127eff5 100644 --- a/cactus-graph/src/execute.cpp +++ b/cactus-graph/src/execute.cpp @@ -60,6 +60,7 @@ DECLARE_COMPUTE(compute_stats_pool_node); DECLARE_COMPUTE(compute_weighted_stats_pool_node); DECLARE_COMPUTE(compute_transpose_node); DECLARE_COMPUTE(compute_gather_node); +DECLARE_COMPUTE(compute_masked_select_prefix_node); DECLARE_COMPUTE(compute_slice_node); DECLARE_COMPUTE(compute_embedding_node); DECLARE_COMPUTE(compute_concat_node); @@ -92,7 +93,7 @@ DECLARE_COMPUTE(compute_expand_node); extern void shrink_thread_local_buffers(); #undef DECLARE_COMPUTE -static constexpr int OP_TYPE_COUNT = static_cast(OpType::EXPAND) + 1; +static constexpr int OP_TYPE_COUNT = static_cast(OpType::MASKED_SELECT_PREFIX) + 1; static_assert(OP_TYPE_COUNT <= 256, "OpType dispatch table overflow"); static ComputeFn dispatch_flat[OP_TYPE_COUNT] = {}; @@ -174,6 +175,7 @@ static bool init_dispatch() { dispatch_flat[static_cast(OpType::GLU)] = compute_glu_node; dispatch_flat[static_cast(OpType::TRANSPOSE)] = compute_transpose_node; dispatch_flat[static_cast(OpType::GATHER)] = compute_gather_node; + dispatch_flat[static_cast(OpType::MASKED_SELECT_PREFIX)] = compute_masked_select_prefix_node; dispatch_flat[static_cast(OpType::SLICE)] = compute_slice_node; dispatch_flat[static_cast(OpType::STRIDED_SLICE)] = compute_strided_slice_node; dispatch_flat[static_cast(OpType::EMBEDDING)] = compute_embedding_node; @@ -278,7 +280,8 @@ static const char* op_type_names[] = { "PAD", "SCALAR_FLOOR_DIVIDE", "STRIDED_SLICE", - "EXPAND" + "EXPAND", + "MASKED_SELECT_PREFIX" }; static const char* get_op_name(OpType op) { diff --git a/cactus-graph/src/graph_ffi.cpp b/cactus-graph/src/graph_ffi.cpp index d1be261e7..138bb7e0f 100644 --- a/cactus-graph/src/graph_ffi.cpp +++ b/cactus-graph/src/graph_ffi.cpp @@ -372,6 +372,19 @@ CACTUS_FFI_EXPORT int cactus_graph_where(cactus_graph_t graph, cactus_node_t con } } +CACTUS_FFI_EXPORT int cactus_graph_masked_select_prefix(cactus_graph_t graph, cactus_node_t x, cactus_node_t mask, cactus_node_t* out) { + if (!graph || !out) return fail_invalid("Invalid args to cactus_graph_masked_select_prefix"); + try { + *out = static_cast(as_graph(graph)->graph.masked_select_prefix( + static_cast(x), + static_cast(mask))); + return 0; + } catch (const std::exception& e) { + last_error_message = e.what(); + return -1; + } +} + int cactus_graph_scalar_add(cactus_graph_t graph, cactus_node_t x, float value, cactus_node_t *out) { if (!graph || !out) { last_error_message = "Invalid args to cactus_graph_scalar_add"; @@ -1716,6 +1729,61 @@ int cactus_graph_execute(cactus_graph_t graph) { } } +int cactus_graph_retain_outputs(cactus_graph_t graph, const cactus_node_t* nodes, size_t count) { + if (!graph || (!nodes && count > 0)) { + last_error_message = "Invalid args to cactus_graph_retain_outputs"; + return -1; + } + try { + std::vector node_ids; + node_ids.reserve(count); + for (size_t i = 0; i < count; ++i) { + node_ids.push_back(static_cast(nodes[i])); + } + as_graph(graph)->graph.retain_outputs(node_ids); + return 0; + } catch (const std::exception& e) { + last_error_message = e.what(); + return -1; + } +} + +int cactus_graph_get_node_op_type(cactus_graph_t graph, cactus_node_t node, int32_t* out_op_type) { + if (!graph || !out_op_type) { + last_error_message = "Invalid args to cactus_graph_get_node_op_type"; + return -1; + } + try { + *out_op_type = static_cast(as_graph(graph)->graph.get_node_op_type(static_cast(node))); + return 0; + } catch (const std::exception& e) { + last_error_message = e.what(); + return -1; + } +} + +int cactus_graph_get_node_inputs(cactus_graph_t graph, cactus_node_t node, cactus_node_t* out_inputs, size_t* inout_count) { + if (!graph || !inout_count) { + last_error_message = "Invalid args to cactus_graph_get_node_inputs"; + return -1; + } + try { + const auto& inputs = as_graph(graph)->graph.get_node_inputs(static_cast(node)); + if (!out_inputs || *inout_count < inputs.size()) { + *inout_count = inputs.size(); + return out_inputs ? -2 : 0; + } + for (size_t i = 0; i < inputs.size(); ++i) { + out_inputs[i] = static_cast(inputs[i]); + } + *inout_count = inputs.size(); + return 0; + } catch (const std::exception& e) { + last_error_message = e.what(); + return -1; + } +} + int cactus_graph_get_output_ptr(cactus_graph_t graph, cactus_node_t node, void** out_ptr) { if (!graph || !out_ptr) { last_error_message = "Invalid args to cactus_graph_get_output_ptr"; diff --git a/cactus-graph/src/io.cpp b/cactus-graph/src/io.cpp index 7a3000553..7a7ff543f 100644 --- a/cactus-graph/src/io.cpp +++ b/cactus-graph/src/io.cpp @@ -248,8 +248,10 @@ namespace { GraphFile::NodeEntry node; node.index = read_u32(in); uint32_t op_type_val = read_u32(in); - if (op_type_val > static_cast(OpType::EXPAND)) { - throw std::runtime_error("Graph file corrupted: invalid op type"); + if (op_type_val > static_cast(OpType::MASKED_SELECT_PREFIX)) { + throw std::runtime_error( + "Graph file corrupted: invalid op type " + std::to_string(op_type_val) + + " > " + std::to_string(static_cast(OpType::MASKED_SELECT_PREFIX))); } node.op_type = static_cast(op_type_val); node.inputs = read_u32_vector(in); diff --git a/cactus-graph/src/ops_conv.cpp b/cactus-graph/src/ops_conv.cpp index 717e3b4c5..65b7a5e00 100644 --- a/cactus-graph/src/ops_conv.cpp +++ b/cactus-graph/src/ops_conv.cpp @@ -197,8 +197,9 @@ void compute_conv1d_node(GraphNode& node, const std::vector(); } else if (W.precision == Precision::INT8) { - W_fp16 = dequantize_int8_weights_to_fp16(W, C_out, C_in * K, "conv1d"); + W_fp16 = dequantize_int8_weights_to_fp16(W, C_out, (is_depthwise ? 1 : C_in) * K, "conv1d"); W_ptr = W_fp16.data(); } else { throw std::runtime_error("Conv1d only supports FP16/INT8 weights"); } - if (node.params.num_groups == C_in && C_out == C_in && W.shape[1] == 1) { + if (is_depthwise) { __fp16* output = Y.data_as<__fp16>(); const __fp16* input = X.data_as<__fp16>(); const size_t L_out = Y.shape[2]; diff --git a/cactus-graph/src/ops_math.cpp b/cactus-graph/src/ops_math.cpp index f1263ff2f..03ea13bc7 100644 --- a/cactus-graph/src/ops_math.cpp +++ b/cactus-graph/src/ops_math.cpp @@ -507,7 +507,8 @@ void compute_where_node(GraphNode& node, const std::vector>& nodes, const std::unordered_map& node_index_map) { const auto& input = get_input(node, 0, nodes, node_index_map); - if (input.precision != Precision::FP16) { + if (input.precision != Precision::FP16 + && !(node.op_type == OpType::CLAMP && input.precision == Precision::FP32)) { throw std::runtime_error("Activation operations only support FP16 precision"); } @@ -548,10 +549,23 @@ void compute_activation_node(GraphNode& node, const std::vector(), - node.output_buffer.data_as<__fp16>(), - node.output_buffer.total_size, - node.params.scalar, node.params.scale); + if (input.precision == Precision::FP32) { + const float lo = node.params.scalar; + const float hi = node.params.scale; + const float* in = input.data_as(); + float* out = node.output_buffer.data_as(); + CactusThreading::parallel_for(node.output_buffer.total_size, CactusThreading::Thresholds::ELEMENT_WISE, + [&](size_t start_idx, size_t end_idx) { + for (size_t i = start_idx; i < end_idx; ++i) { + out[i] = std::min(std::max(in[i], lo), hi); + } + }); + } else { + cactus_clamp_f16(input.data_as<__fp16>(), + node.output_buffer.data_as<__fp16>(), + node.output_buffer.total_size, + node.params.scalar, node.params.scale); + } break; default: break; diff --git a/cactus-graph/src/ops_nn.cpp b/cactus-graph/src/ops_nn.cpp index 23afc2f20..8ab813c05 100644 --- a/cactus-graph/src/ops_nn.cpp +++ b/cactus-graph/src/ops_nn.cpp @@ -635,12 +635,12 @@ void compute_attention_node(GraphNode& node, const std::vectorshape.size() == 4) { if (mask_buffer->shape[0] != batch_size || - mask_buffer->shape[1] != num_q_heads || + (mask_buffer->shape[1] != num_q_heads && mask_buffer->shape[1] != 1) || mask_buffer->shape[2] != seq_len || mask_buffer->shape[3] != kv_seq_len) { - throw std::runtime_error("Attention mask [B, H, T, S] shape mismatch"); + throw std::runtime_error("Attention mask [B, H|1, T, S] shape mismatch"); } - mask_per_head = true; + mask_per_head = mask_buffer->shape[1] != 1; } else { throw std::runtime_error("Attention mask must be rank 3 or 4"); } @@ -930,4 +930,3 @@ void compute_groupnorm_node(GraphNode& node, const std::vector>& nodes, const std::unordered_map& node_index_map) { + const auto& input_buffer = get_input(node, 0, nodes, node_index_map); + const auto& mask_buffer = get_input(node, 1, nodes, node_index_map); + + if (input_buffer.shape.size() != 2 && input_buffer.shape.size() != 3) { + throw std::runtime_error("masked_select_prefix input must be [T,D] or [1,T,D]"); + } + if (mask_buffer.shape.size() != 1 && mask_buffer.shape.size() != 2) { + throw std::runtime_error("masked_select_prefix mask must be [T] or [1,T]"); + } + + const size_t time = input_buffer.shape.size() == 3 ? input_buffer.shape[1] : input_buffer.shape[0]; + const size_t hidden = input_buffer.shape.back(); + if (mask_buffer.shape.back() != time) { + throw std::runtime_error("masked_select_prefix input/mask time dimension mismatch"); + } + + const size_t row_bytes = PrecisionTraits::packed_size_of(input_buffer.precision, hidden); + const char* input = static_cast(input_buffer.get_data()); + char* output = static_cast(node.output_buffer.get_data()); + std::memset(output, 0, node.output_buffer.byte_size); + + auto mask_true = [&](size_t i) -> bool { + switch (mask_buffer.precision) { + case Precision::FP32: + return mask_buffer.data_as()[i] != 0.0f; + case Precision::FP16: + return static_cast(mask_buffer.data_as<__fp16>()[i]) != 0.0f; + case Precision::INT8: + return mask_buffer.data_as()[i] != 0; + default: + throw std::runtime_error("masked_select_prefix mask must be FP32/FP16/INT8"); + } + }; + + size_t out_row = 0; + for (size_t i = 0; i < time; ++i) { + if (!mask_true(i)) continue; + if (out_row >= time) break; + std::memcpy( + output + PrecisionTraits::byte_offset_of(node.output_buffer.precision, out_row * hidden), + input + PrecisionTraits::byte_offset_of(input_buffer.precision, i * hidden), + row_bytes); + ++out_row; + } +} + void compute_slice_node(GraphNode& node, const std::vector>& nodes, const std::unordered_map& node_index_map) { auto* input_node = nodes[node_index_map.at(node.input_ids[0])].get(); auto& input_buffer = input_node->output_buffer; diff --git a/cactus-kernels/src/nn.cpp b/cactus-kernels/src/nn.cpp index c56de09db..a0ae9ef02 100644 --- a/cactus-kernels/src/nn.cpp +++ b/cactus-kernels/src/nn.cpp @@ -439,6 +439,11 @@ void kernel_softmax_f16_single(const __fp16* input, __fp16* output, size_t vocab max_val = std::max(max_val, static_cast(input[i])); } + if (max_val == -std::numeric_limits::infinity()) { + std::fill(output, output + vocab_size, static_cast<__fp16>(0.0f)); + return; + } + const float32x4_t max_broadcast = vdupq_n_f32(max_val); float32x4_t sum_vec[UNROLL_FACTOR * 2]; diff --git a/python/cactus/bindings/cactus.py b/python/cactus/bindings/cactus.py index 9a4a5418d..49ed90c6b 100644 --- a/python/cactus/bindings/cactus.py +++ b/python/cactus/bindings/cactus.py @@ -712,6 +712,24 @@ class cactus_tensor_info_t(ctypes.Structure): _lib.cactus_graph_execute.argtypes = [cactus_graph_t] _lib.cactus_graph_execute.restype = ctypes.c_int +_lib.cactus_graph_retain_outputs.argtypes = [ + cactus_graph_t, ctypes.POINTER(cactus_node_t), ctypes.c_size_t +] +_lib.cactus_graph_retain_outputs.restype = ctypes.c_int + +_lib.cactus_graph_get_node_op_type.argtypes = [ + cactus_graph_t, cactus_node_t, ctypes.POINTER(ctypes.c_int32) +] +_lib.cactus_graph_get_node_op_type.restype = ctypes.c_int + +_lib.cactus_graph_get_node_inputs.argtypes = [ + cactus_graph_t, + cactus_node_t, + ctypes.POINTER(cactus_node_t), + ctypes.POINTER(ctypes.c_size_t), +] +_lib.cactus_graph_get_node_inputs.restype = ctypes.c_int + _lib.cactus_graph_get_output_ptr.argtypes = [cactus_graph_t, cactus_node_t, ctypes.POINTER(ctypes.c_void_p)] _lib.cactus_graph_get_output_ptr.restype = ctypes.c_int @@ -1674,6 +1692,31 @@ def execute(self): if rc != 0: raise RuntimeError(_err("graph_execute failed")) + def retain_outputs(self, node_ids): + node_ids = tuple(int(node_id) for node_id in node_ids) + arr = (cactus_node_t * len(node_ids))(*node_ids) + rc = _lib.cactus_graph_retain_outputs(self.h, arr, len(node_ids)) + if rc != 0: + raise RuntimeError(_err("graph_retain_outputs failed")) + + def node_op_type(self, node_id): + out = ctypes.c_int32() + rc = _lib.cactus_graph_get_node_op_type(self.h, cactus_node_t(int(node_id)), ctypes.byref(out)) + if rc != 0: + raise RuntimeError(_err("graph_get_node_op_type failed")) + return int(out.value) + + def node_inputs(self, node_id): + count = ctypes.c_size_t() + rc = _lib.cactus_graph_get_node_inputs(self.h, cactus_node_t(int(node_id)), None, ctypes.byref(count)) + if rc != 0: + raise RuntimeError(_err("graph_get_node_inputs failed")) + arr = (cactus_node_t * int(count.value))() + rc = _lib.cactus_graph_get_node_inputs(self.h, cactus_node_t(int(node_id)), arr, ctypes.byref(count)) + if rc != 0: + raise RuntimeError(_err("graph_get_node_inputs failed")) + return tuple(int(arr[i]) for i in range(int(count.value))) + def add(self, a, b): return self._binary("cactus_graph_add", a, b) diff --git a/python/cactus/transpiler/Converter/input_processor.py b/python/cactus/transpiler/Converter/input_processor.py index 095bb432c..c73d1c1b2 100644 --- a/python/cactus/transpiler/Converter/input_processor.py +++ b/python/cactus/transpiler/Converter/input_processor.py @@ -87,10 +87,30 @@ def text_tokenizer_processor(model_id: str, configs: dict[str, dict[str, Any]], ) +class ParakeetAudioProcessor: + def __init__(self, configs: dict[str, dict[str, Any]]): + config = configs.get("config.json", {}) + encoder_config = config.get("encoder_config", {}) if isinstance(config.get("encoder_config"), dict) else {} + self.feature_size = int(encoder_config.get("num_mel_bins", 128)) + self.frame_count = int(encoder_config.get("subsampling_factor", 8)) * 32 + + def __call__(self, audio=None, return_tensors: str = "pt", sampling_rate: int = 16000, **_: Any): + import torch + + return { + "input_features": torch.zeros((1, self.frame_count, self.feature_size), dtype=torch.float32), + "attention_mask": torch.ones((1, self.frame_count), dtype=torch.long), + } + + +def parakeet_processor(model_id: str, configs: dict[str, dict[str, Any]], model_profile: str): + return ParakeetAudioProcessor(configs) + + PROCESSOR_MAP = { "google/gemma-4-E2B": gemma4_processor, "openai/whisper-tiny": default_processor, - "nvidia/parakeet-tdt-0.6b-v3": default_processor, + "nvidia/parakeet-tdt-0.6b-v3": parakeet_processor, "LiquidAI/LFM2-VL-3B": default_processor, "Qwen/Qwen2.5-0.5B": default_processor, "LiquidAI/LFM2.5-8B-A1B": text_tokenizer_processor, diff --git a/python/cactus/transpiler/Converter/models.py b/python/cactus/transpiler/Converter/models.py index 96cb36f5e..4ad25f6a3 100644 --- a/python/cactus/transpiler/Converter/models.py +++ b/python/cactus/transpiler/Converter/models.py @@ -353,13 +353,61 @@ def build_processor_kwargs(model_id: str, input_modalities: tuple[str, ...], con if audio_token is not None: prompt.append(audio_token) call_kwargs["audio"] = _load_audio_asset() + call_kwargs["sampling_rate"] = 16000 if "text" in input_modalities: prompt.append("Describe this input.") call_kwargs["text"] = " ".join(prompt) if prompt else "Hello" processed = processor(**call_kwargs) - return {key: value for key, value in dict(processed).items() if isinstance(value, torch.Tensor)} + pad_audio_export_inputs(processed, configs) + kwargs = {key: value for key, value in dict(processed).items() if isinstance(value, torch.Tensor)} + normalize_processor_kwargs_for_export(kwargs, model_profile) + return kwargs + + +def normalize_processor_kwargs_for_export(kwargs: dict[str, torch.Tensor], model_profile: str) -> None: + if model_profile == "whisper" and "labels" in kwargs and "decoder_input_ids" not in kwargs: + kwargs["decoder_input_ids"] = kwargs.pop("labels") + + +def pad_audio_export_inputs(processed: Any, configs: dict[str, dict[str, Any]]) -> None: + target_frames = gemma4_audio_export_frame_count(configs) + if target_frames is None: + return + + if "input_features" in processed: + processed["input_features"] = pad_tensor_dim(processed["input_features"], dim=1, target=target_frames, value=0) + + if "input_features_mask" in processed: + processed["input_features_mask"] = pad_tensor_dim(processed["input_features_mask"], dim=1, target=target_frames, value=0) + + +def gemma4_audio_export_frame_count(configs: dict[str, dict[str, Any]]) -> int | None: + processor_config = configs.get("processor_config.json", {}) + config = configs.get("config.json", {}) + + if processor_config.get("processor_class") != "Gemma4Processor": + return None + + audio_seq_length = processor_config.get("audio_seq_length") + audio_config = config.get("audio_config", {}) + attention_chunk_size = audio_config.get("attention_chunk_size") + + if audio_seq_length is None or attention_chunk_size is None: + return None + + return int(audio_seq_length) * 4 + int(attention_chunk_size) + + +def pad_tensor_dim(tensor: torch.Tensor, dim: int, target: int, value: int | float | bool = 0) -> torch.Tensor: + if dim >= tensor.ndim or tensor.shape[dim] >= target: + return tensor + + shape = list(tensor.shape) + shape[dim] = target - int(tensor.shape[dim]) + padding = torch.full(shape, value, dtype=tensor.dtype, device=tensor.device) + return torch.cat((tensor, padding), dim=dim) #Loads the best HF model class for the requested model profile. diff --git a/python/cactus/transpiler/Fusions/edges.py b/python/cactus/transpiler/Fusions/edges.py index c2c1972f8..b98488c7d 100644 --- a/python/cactus/transpiler/Fusions/edges.py +++ b/python/cactus/transpiler/Fusions/edges.py @@ -234,6 +234,10 @@ def _index_edges_by_dest(edges: dict[str, M.FusionEdge]) -> dict[str, tuple[str, "expert_down_to_weighted": _expert_edge("expert_down", "expert_weighted", 0), "expert_weighted_to_combine": _expert_edge("expert_weighted", "moe_combine", 0), "sample_topk_to_softmax": _edge("sample_topk", "sample_softmax", 0), + "silu_sigmoid_to_product": _edge("silu_sigmoid", "silu_product", 1), + "glu_left_slice_to_product": _edge("glu_left_slice", "glu_product", 0), + "glu_right_slice_to_sigmoid": _edge("glu_right_slice", "glu_sigmoid", 0), + "glu_sigmoid_to_product": _edge("glu_sigmoid", "glu_product", 1), } @@ -254,6 +258,18 @@ def _index_edges_by_dest(edges: dict[str, M.FusionEdge]) -> dict[str, tuple[str, "rms_inv_pow_to_scale", "rms_scale_to_weight_mul", ), + "rms_norm_no_weight": ( + "rms_square_to_mean", + "rms_mean_to_eps_add", + "rms_eps_add_to_inv", + "rms_inv_to_scale", + ), + "rms_norm_pow_no_weight": ( + "rms_square_to_mean", + "rms_mean_to_eps_add", + "rms_eps_add_to_inv_pow", + "rms_inv_pow_to_scale", + ), "layernorm": ( "ln_mean_to_center", "ln_center_to_square", @@ -275,6 +291,14 @@ def _index_edges_by_dest(edges: dict[str, M.FusionEdge]) -> dict[str, tuple[str, "up_proj_to_gelu", "gelu_to_down_proj", ), + "silu_decomposed": ( + "silu_sigmoid_to_product", + ), + "glu_decomposed": ( + "glu_left_slice_to_product", + "glu_right_slice_to_sigmoid", + "glu_sigmoid_to_product", + ), "gemma4_geglu_mlp": ( "gate_input_view_to_gate_proj", "gate_weight_transpose_to_gate_proj", diff --git a/python/cactus/transpiler/Fusions/fusions.py b/python/cactus/transpiler/Fusions/fusions.py index 52192da43..bbdd4497c 100644 --- a/python/cactus/transpiler/Fusions/fusions.py +++ b/python/cactus/transpiler/Fusions/fusions.py @@ -270,6 +270,42 @@ def _index_by_root_op( "scalar_cos": _single_node_graph("scalar_cos", "scalar_cos", ("x",)), "scalar_sin": _single_node_graph("scalar_sin", "scalar_sin", ("x",)), "scalar_log": _single_node_graph("scalar_log", "scalar_log", ("x",)), + "scalar_not_equal": _single_node_graph( + "scalar_not_equal", + "scalar_not_equal", + ("x",), + attr_captures=(M.AttrCapture("value", "scalar_not_equal", "other", required=False),), + ), + "scalar_equal": _single_node_graph( + "scalar_equal", + "scalar_equal", + ("x",), + attr_captures=(M.AttrCapture("value", "scalar_equal", "other", required=False),), + ), + "scalar_less": _single_node_graph( + "scalar_less", + "scalar_less", + ("x",), + attr_captures=(M.AttrCapture("value", "scalar_less", "other", required=False),), + ), + "scalar_less_equal": _single_node_graph( + "scalar_less_equal", + "scalar_less_equal", + ("x",), + attr_captures=(M.AttrCapture("value", "scalar_less_equal", "other", required=False),), + ), + "scalar_greater": _single_node_graph( + "scalar_greater", + "scalar_greater", + ("x",), + attr_captures=(M.AttrCapture("value", "scalar_greater", "other", required=False),), + ), + "scalar_greater_equal": _single_node_graph( + "scalar_greater_equal", + "scalar_greater_equal", + ("x",), + attr_captures=(M.AttrCapture("value", "scalar_greater_equal", "other", required=False),), + ), "mean": _single_node_graph( "mean", "mean", @@ -300,6 +336,12 @@ def _index_by_root_op( ("x",), attr_captures=(M.AttrCapture("axis", "max", "dim", default=-1, required=False),), ), + "cumsum": _single_node_graph( + "cumsum", + "cumsum", + ("x",), + attr_captures=(M.AttrCapture("axis", "cumsum", "dim", default=-1, required=False),), + ), "softmax": _single_node_graph( "softmax", "softmax", @@ -312,6 +354,17 @@ def _index_by_root_op( ("x",), attr_captures=(M.AttrCapture("k", "topk_direct", "k", required=True),), ), + "where": _single_node_graph("where", "where", ("condition", "x", "y")), + "not_equal": _single_node_graph("not_equal", "not_equal", ("a", "b")), + "equal": _single_node_graph("equal", "equal", ("a", "b")), + "less": _single_node_graph("less", "less", ("a", "b")), + "less_equal": _single_node_graph("less_equal", "less_equal", ("a", "b")), + "greater": _single_node_graph("greater", "greater", ("a", "b")), + "greater_equal": _single_node_graph("greater_equal", "greater_equal", ("a", "b")), + "bitwise_and": _single_node_graph("bitwise_and", "bitwise_and", ("a", "b")), + "bitwise_or": _single_node_graph("bitwise_or", "bitwise_or", ("a", "b")), + "logical_not": _single_node_graph("logical_not", "logical_not", ("x",)), + "bitwise_not": _single_node_graph("bitwise_not", "bitwise_not", ("x",)), "view": _single_node_graph( "view", "view", @@ -333,6 +386,12 @@ def _index_by_root_op( M.AttrCapture("end_dim", "flatten", "end_dim", default=-1, required=False), ), ), + "expand": _single_node_graph( + "expand", + "expand", + ("x",), + attr_captures=(M.AttrCapture("shape", "expand", "shape", required=False),), + ), "transpose": _single_node_graph( "transpose", "transpose", @@ -350,6 +409,16 @@ def _index_by_root_op( M.AttrCapture("step", "slice", "step", required=False), ), ), + "unfold": _single_node_graph( + "unfold", + "unfold", + ("x",), + attr_captures=( + M.AttrCapture("dimension", "unfold", "dim", required=False), + M.AttrCapture("size", "unfold", "size", required=False), + M.AttrCapture("step", "unfold", "step", default=1, required=False), + ), + ), "index": _single_node_graph( "index", "index", @@ -360,7 +429,27 @@ def _index_by_root_op( ), ), "cat": _graph("cat", "cat", ("cat",), inputs=(_variadic_input("values", "cat", 0, min_count=2),), attr_captures=(M.AttrCapture("axis", "cat", "dim", default=0, required=False),)), - "gather": _single_node_graph("gather", "gather", ("tensor", "indices")), + "gather": _single_node_graph( + "gather", + "gather", + ("tensor", "indices"), + attr_captures=(M.AttrCapture("axis", "gather", "dim", default=0, required=False),), + ), + "precision_cast": _single_node_graph( + "precision_cast", + "precision_cast", + ("x",), + attr_captures=(M.AttrCapture("dtype", "precision_cast", "dtype", required=False),), + ), + "pad": _single_node_graph( + "pad", + "pad", + ("x",), + attr_captures=( + M.AttrCapture("pad", "pad", "pad", required=False), + M.AttrCapture("value", "pad", "value", default=0.0, required=False), + ), + ), "relu": _single_node_graph("relu", "relu", ("x",)), "silu": _single_node_graph("silu", "silu", ("x",)), "gelu": _single_node_graph("gelu", "gelu", ("x",)), @@ -464,6 +553,26 @@ def _index_by_root_op( }, ) +RMS_NORM_NO_WEIGHT_GRAPH = _graph( + "rms_norm_no_weight", + "rms_scale", + ("rms_square", "rms_mean", "rms_eps_add", "rms_inv", "rms_scale"), + edge_names=E.EDGE_GROUPS["rms_norm_no_weight"], + inputs=(_input("x", "rms_square", 0),), + shared_inputs=(_shared_input("rms_square", 0, "rms_scale", 0),), + attr_captures=(M.AttrCapture("epsilon", "rms_eps_add", "other", default=1e-6, required=False),), +) + +RMS_NORM_POW_NO_WEIGHT_GRAPH = _graph( + "rms_norm_pow_no_weight", + "rms_scale", + ("rms_square", "rms_mean", "rms_eps_add", "rms_inv_pow", "rms_scale"), + edge_names=E.EDGE_GROUPS["rms_norm_pow_no_weight"], + inputs=(_input("x", "rms_square", 0),), + shared_inputs=(_shared_input("rms_square", 0, "rms_scale", 0),), + attr_captures=(M.AttrCapture("epsilon", "rms_eps_add", "other", default=1e-6, required=False),), +) + LAYERNORM_NO_BIAS_GRAPH = _graph( "layernorm_no_bias", "ln_weight_mul", @@ -500,6 +609,31 @@ def _index_by_root_op( }, ) +SILU_DECOMPOSED_GRAPH = _graph( + "silu_decomposed", + "silu_product", + ("silu_sigmoid", "silu_product"), + edge_names=E.EDGE_GROUPS["silu_decomposed"], + inputs=(_input("x", "silu_product", 0),), + shared_inputs=(_shared_input("silu_product", 0, "silu_sigmoid", 0),), +) + +GLU_DECOMPOSED_GRAPH = _graph( + "glu_decomposed", + "glu_product", + ("glu_left_slice", "glu_right_slice", "glu_sigmoid", "glu_product"), + edge_names=E.EDGE_GROUPS["glu_decomposed"], + inputs=(_input("x", "glu_left_slice", 0),), + shared_inputs=(_shared_input("glu_left_slice", 0, "glu_right_slice", 0),), + attr_captures=(M.AttrCapture("axis", "glu_left_slice", "axis", default=-1, required=False),), + constraints={ + "slice_halves": { + "left_node": "glu_left_slice", + "right_node": "glu_right_slice", + }, + }, +) + SWIGLU_MLP_GRAPH = _graph( "swiglu_mlp", "down_proj", @@ -584,6 +718,19 @@ def _index_by_root_op( }, ) +GEMMA4_ATTENTION_LAYOUT_GRAPH = _single_node_graph( + "gemma4_attention_layout", + "view", + ("x",), + attr_captures=( + M.AttrCapture("scale", default=1.0, required=False), + M.AttrCapture("input_layout", default="bhqd_bhds_bhsd", required=False), + M.AttrCapture("output_layout", default="bhqd", required=False), + M.AttrCapture("window_size", default=0, required=False), + ), + metadata={"special_matcher": "gemma4_attention_layout"}, +) + LFM_BMM_MASKED_ATTENTION_GRAPH = _graph( "lfm_bmm_masked_attention", "lfm_attn_output_view", @@ -1085,14 +1232,19 @@ def _index_by_root_op( "linear_bias": LINEAR_BIAS_GRAPH, "rms_norm": RMS_NORM_GRAPH, "rms_norm_pow": RMS_NORM_POW_GRAPH, + "rms_norm_no_weight": RMS_NORM_NO_WEIGHT_GRAPH, + "rms_norm_pow_no_weight": RMS_NORM_POW_NO_WEIGHT_GRAPH, "layernorm_no_bias": LAYERNORM_NO_BIAS_GRAPH, "layernorm": LAYERNORM_GRAPH, + "silu_decomposed": SILU_DECOMPOSED_GRAPH, + "glu_decomposed": GLU_DECOMPOSED_GRAPH, "swiglu_mlp": SWIGLU_MLP_GRAPH, "gelu_mlp": GELU_MLP_GRAPH, "gemma4_geglu_mlp": GEMMA4_GEGLU_MLP_GRAPH, "scaled_dot_product_attention": ATTENTION_DIRECT_GRAPH, "attention_core": ATTENTION_CORE_GRAPH, "attention_masked": ATTENTION_MASKED_GRAPH, + "gemma4_attention_layout": GEMMA4_ATTENTION_LAYOUT_GRAPH, "lfm_bmm_masked_attention": LFM_BMM_MASKED_ATTENTION_GRAPH, "rope": ROPE_GRAPH, "conv": CONV_GRAPH, @@ -1141,14 +1293,26 @@ def _index_by_root_op( "linear_bias": _definition("linear_bias", "linear", LINEAR_BIAS_GRAPH, fusion_fields=("generic", "linear")), "rms_norm": _definition("rms_norm", "rms_norm", RMS_NORM_GRAPH, fusion_fields=("generic", "rmsnorm", "gemma4_rmsnorm", "qwen2_5_rmsnorm")), "rms_norm_pow": _definition("rms_norm_pow", "rms_norm", RMS_NORM_POW_GRAPH, fusion_fields=("generic", "rmsnorm", "gemma4_rmsnorm", "qwen2_5_rmsnorm")), + "rms_norm_no_weight": _definition("rms_norm_no_weight", "rms_norm", RMS_NORM_NO_WEIGHT_GRAPH, fusion_fields=("generic", "rmsnorm", "gemma4_rmsnorm", "qwen2_5_rmsnorm")), + "rms_norm_pow_no_weight": _definition("rms_norm_pow_no_weight", "rms_norm", RMS_NORM_POW_NO_WEIGHT_GRAPH, fusion_fields=("generic", "rmsnorm", "gemma4_rmsnorm", "qwen2_5_rmsnorm")), "layernorm_no_bias": _definition("layernorm_no_bias", "layernorm", LAYERNORM_NO_BIAS_GRAPH, fusion_fields=("generic", "normalization")), "layernorm": _definition("layernorm", "layernorm", LAYERNORM_GRAPH, fusion_fields=("generic", "normalization")), + "silu_decomposed": _definition("silu_decomposed", "silu", SILU_DECOMPOSED_GRAPH, fusion_fields=("generic", "activation", "audio", "mlp")), + "glu_decomposed": _definition("glu_decomposed", "glu", GLU_DECOMPOSED_GRAPH, fusion_fields=("generic", "activation", "audio")), "swiglu_mlp": _definition("swiglu_mlp", "dense_mlp_tq_fused", SWIGLU_MLP_GRAPH, fusion_fields=("generic", "mlp", "gemma4_mlp", "qwen2_5_mlp", "lfm_mlp")), "gelu_mlp": _definition("gelu_mlp", "matmul", GELU_MLP_GRAPH, fusion_fields=("generic", "mlp")), "gemma4_geglu_mlp": _definition("gemma4_geglu_mlp", "dense_mlp_tq_fused", GEMMA4_GEGLU_MLP_GRAPH, fusion_fields=("gemma4_mlp",)), "scaled_dot_product_attention": _definition("scaled_dot_product_attention", "attention", ATTENTION_DIRECT_GRAPH, fusion_fields=("generic", "attention")), "attention_core": _definition("attention_core", "attention", ATTENTION_CORE_GRAPH, fusion_fields=("generic", "attention", "gemma4_attention", "qwen2_5_attention", "lfm_attention")), "attention_masked": _definition("attention_masked", "attention", ATTENTION_MASKED_GRAPH, fusion_fields=("generic", "attention", "gemma4_attention")), + "gemma4_attention_layout": _definition( + "gemma4_attention_layout", + "attention", + GEMMA4_ATTENTION_LAYOUT_GRAPH, + fusion_fields=("attention", "gemma4_attention"), + supported_inference_modes=("prefill_with_cache", "decode_with_cache"), + metadata={"special_matcher": "gemma4_attention_layout"}, + ), "gemma4_bmm_masked_attention": _definition("gemma4_bmm_masked_attention", "attention", LFM_BMM_MASKED_ATTENTION_GRAPH, fusion_fields=("generic", "attention", "gemma4_attention")), "lfm_bmm_masked_attention": _definition("lfm_bmm_masked_attention", "attention", LFM_BMM_MASKED_ATTENTION_GRAPH, fusion_fields=("generic", "attention", "lfm_attention", "lfm_moe")), "rope": _definition("rope", "rope", ROPE_GRAPH, fusion_fields=("generic", "rope", "gemma4_rope", "qwen2_5_rope")), diff --git a/python/cactus/transpiler/Fusions/nodes.py b/python/cactus/transpiler/Fusions/nodes.py index e2d423a50..6acaf39f5 100644 --- a/python/cactus/transpiler/Fusions/nodes.py +++ b/python/cactus/transpiler/Fusions/nodes.py @@ -36,7 +36,7 @@ "sort": ("aten.sort.default",), "view": ("aten.view.default", "aten.reshape.default", "aten.flatten.using_ints", "cactus.view", "cactus.reshape"), "unsqueeze": ("aten.unsqueeze.default",), - "expand": ("aten.expand.default",), + "expand": ("aten.expand.default", "cactus.expand"), "transpose": ("aten.t.default", "aten.transpose.int", "aten.permute.default", "cactus.transpose"), "split": ("aten.split_with_sizes.default", "aten.chunk.default"), "slice": ("aten.slice.Tensor", "aten.narrow.default", "cactus.slice"), @@ -47,25 +47,33 @@ "conv_cache_append_op": ("aten.cat.default", "cactus.conv_cache_append"), "scatter": ("aten.scatter.src", "aten.scatter.value", "aten.scatter_reduce.two"), "histc": ("aten.histc.default",), - "cumsum": ("aten.cumsum.default",), + "cumsum": ("aten.cumsum.default", "cactus.cumsum"), "empty": ("aten.empty_permuted.default",), "grouped_mm": ("transformers.grouped_mm_fallback.default",), "neg": ("aten.neg.default", "operator.neg", ""), - "where": ("aten.where.self",), - "eq": ("aten.eq.Tensor", "aten.eq.Scalar"), - "logical_not": ("aten.logical_not.default",), + "where": ("aten.where.self", "cactus.where"), + "eq": ("aten.eq.Tensor", "aten.eq.Scalar", "cactus.equal", "cactus.scalar_equal"), + "logical_not": ("aten.logical_not.default", "cactus.logical_not"), "any": ("aten.any.dim",), - "ne": ("aten.ne.Tensor", "aten.ne.Scalar"), - "lt": ("aten.lt.Tensor", "aten.lt.Scalar"), - "gt": ("aten.gt.Tensor", "aten.gt.Scalar"), + "ne": ("aten.ne.Tensor", "aten.ne.Scalar", "cactus.not_equal", "cactus.scalar_not_equal"), + "lt": ("aten.lt.Tensor", "aten.lt.Scalar", "cactus.less", "cactus.scalar_less"), + "le": ("aten.le.Tensor", "aten.le.Scalar", "cactus.less_equal", "cactus.scalar_less_equal"), + "gt": ("aten.gt.Tensor", "aten.gt.Scalar", "cactus.greater", "cactus.scalar_greater"), + "ge": ("aten.ge.Tensor", "aten.ge.Scalar", "cactus.greater_equal", "cactus.scalar_greater_equal"), + "bitwise_and": ("aten.bitwise_and.Tensor", "cactus.bitwise_and"), + "bitwise_or": ("aten.bitwise_or.Tensor", "cactus.bitwise_or"), + "bitwise_not": ("aten.bitwise_not.default", "cactus.bitwise_not"), "full": ("aten.full.default", "aten.full_like.default"), "arange": ("aten.arange.default", "aten.arange.start", "aten.arange.start_step"), "scalar_tensor": ("aten.scalar_tensor.default",), + "precision_cast": ("aten._to_copy.default", "cactus.precision_cast"), "copy": ("aten.clone.default", "aten._to_copy.default", "aten.contiguous.default"), "copy_write": ("aten.copy.default",), + "pad": ("aten.constant_pad_nd.default", "cactus.pad"), "slice_scatter": ("aten.slice_scatter.default",), + "unfold": ("aten.unfold.default", "cactus.unfold"), "relu": ("aten.relu.default",), - "silu": ("aten.silu.default", "aten.sigmoid.default"), + "silu": ("aten.silu.default",), "gelu": ("aten.gelu.default",), "tanh": ("aten.tanh.default",), "sigmoid": ("aten.sigmoid.default",), @@ -167,8 +175,26 @@ def _expert_node(name: str, op_key: str) -> M.FusionNode: "histc": _node("histc", "histc"), "cumsum": _node("cumsum", "cumsum"), "empty": _node("empty", "empty"), + "where": _node("where", "where"), "not_equal": _node("not_equal", "ne"), "equal": _node("equal", "eq"), + "less": _node("less", "lt"), + "less_equal": _node("less_equal", "le"), + "greater": _node("greater", "gt"), + "greater_equal": _node("greater_equal", "ge"), + "bitwise_and": _node("bitwise_and", "bitwise_and"), + "bitwise_or": _node("bitwise_or", "bitwise_or"), + "logical_not": _node("logical_not", "logical_not"), + "bitwise_not": _node("bitwise_not", "bitwise_not"), + "scalar_not_equal": _node("scalar_not_equal", "ne"), + "scalar_equal": _node("scalar_equal", "eq"), + "scalar_less": _node("scalar_less", "lt"), + "scalar_less_equal": _node("scalar_less_equal", "le"), + "scalar_greater": _node("scalar_greater", "gt"), + "scalar_greater_equal": _node("scalar_greater_equal", "ge"), + "precision_cast": _node("precision_cast", "precision_cast"), + "pad": _node("pad", "pad"), + "unfold": _node("unfold", "unfold"), "embedding": _node("embedding", "embedding"), "linear_mm": _node("linear_mm", "matmul"), "linear_weight_transpose": _node("linear_weight_transpose", "transpose"), @@ -198,12 +224,18 @@ def _expert_node(name: str, op_key: str) -> M.FusionNode: "batchnorm_div_std": _node("batchnorm_div_std", "div"), "relu": _node("relu", "relu"), "silu": _node("silu", "silu"), + "silu_sigmoid": _node("silu_sigmoid", "sigmoid"), + "silu_product": _node("silu_product", "mul"), "gelu": _node("gelu", "gelu"), "gelu_erf": _node("gelu_erf", "gelu", metadata={"approximation": "erf"}), "sigmoid": _node("sigmoid", "sigmoid"), "tanh": _node("tanh", "tanh"), "clamp": _node("clamp", "clamp"), "glu": _node("glu", "mul"), + "glu_left_slice": _node("glu_left_slice", "slice"), + "glu_right_slice": _node("glu_right_slice", "slice"), + "glu_sigmoid": _node("glu_sigmoid", "sigmoid"), + "glu_product": _node("glu_product", "mul"), "gate_proj": _node("gate_proj", "matmul"), "up_proj": _node("up_proj", "matmul"), "gate_input_view": _node("gate_input_view", "view"), @@ -393,23 +425,43 @@ def _expert_node(name: str, op_key: str) -> M.FusionNode: "scalar_cos", "scalar_sin", "scalar_log", + "scalar_not_equal", + "scalar_equal", + "scalar_less", + "scalar_less_equal", + "scalar_greater", + "scalar_greater_equal", "mean", "sum", "variance", "min", "max", + "cumsum", "softmax", "topk_direct", + "where", "view", "reshape", "flatten", + "expand", "transpose", "slice", + "unfold", "index", "cat", "gather", + "precision_cast", + "pad", "not_equal", "equal", + "less", + "less_equal", + "greater", + "greater_equal", + "bitwise_and", + "bitwise_or", + "logical_not", + "bitwise_not", ), "linear": ("linear_mm", "linear_bias_add"), "normalization": ( diff --git a/python/cactus/transpiler/Generator/component_split.py b/python/cactus/transpiler/Generator/component_split.py index f7b4c14f5..a631a5e64 100644 --- a/python/cactus/transpiler/Generator/component_split.py +++ b/python/cactus/transpiler/Generator/component_split.py @@ -10,6 +10,7 @@ PREFILL_WITH_CACHE_TASK = "prefill_with_cache" DECODE_WITH_CACHE_TASK = "decode_with_cache" +GEMMA4_PREFILL_CHUNK_TOKENS = 128 @dataclass(slots=True, frozen=True) @@ -26,6 +27,7 @@ class PlaceholderSpec: class OutputSpec: node: str logical_name: str + row_limit: int | None = None @dataclass(slots=True, frozen=True) @@ -36,6 +38,7 @@ class ComponentSplitSpec: placeholders: tuple[PlaceholderSpec, ...] = () ref_aliases: Mapping[str, str] = field(default_factory=dict) input_aliases: Mapping[str, str] = field(default_factory=dict) + chunk_tokens: int | None = None @dataclass(slots=True, frozen=True) @@ -83,7 +86,7 @@ def split_gemma4_components(prefill: IRModels.Graph, decode: IRModels.Graph) -> ComponentSplitSpec( name="vision_encoder", graph=prefill, - outputs=(OutputSpec(boundaries.vision_features, "image_features"),), + outputs=(OutputSpec(boundaries.vision_features, "image_features", row_limit=256),), input_aliases={"image_position_ids": "pixel_position_ids"}, ), ComponentSplitSpec( @@ -115,6 +118,22 @@ def split_gemma4_components(prefill: IRModels.Graph, decode: IRModels.Graph) -> boundaries.merged_inputs_embeds: boundaries.prefill_text_inputs_embeds, boundaries.prefill_token_ids: "input_ids", }, + chunk_tokens=GEMMA4_PREFILL_CHUNK_TOKENS, + ), + ComponentSplitSpec( + name="lm_encoder_media_chunk", + graph=prefill, + outputs=( + OutputSpec(inputs_embeds_chunk.name, "inputs_embeds"), + OutputSpec(boundaries.prefill_per_layer_inputs, "per_layer_inputs"), + OutputSpec(position_chunk.name, "position_ids"), + ), + placeholders=(inputs_embeds_chunk, position_chunk), + ref_aliases={ + boundaries.merged_inputs_embeds: inputs_embeds_chunk.name, + boundaries.prefill_token_ids: "input_ids", + }, + chunk_tokens=GEMMA4_PREFILL_CHUNK_TOKENS, ), ComponentSplitSpec( name="lm_encoder_media_step", @@ -132,6 +151,7 @@ def split_gemma4_components(prefill: IRModels.Graph, decode: IRModels.Graph) -> graph=prefill, outputs=(OutputSpec(boundaries.prefill_logits, "logits"),), placeholders=(inputs_embeds_chunk, per_layer_chunk, position_chunk), + chunk_tokens=GEMMA4_PREFILL_CHUNK_TOKENS, ), ComponentSplitSpec( name="decoder_step", @@ -284,6 +304,7 @@ def extract_component_graph(spec: ComponentSplitSpec) -> IRModels.Graph: output_logical_names = { ref_replacements.get(output.node, output.node): output.logical_name for output in spec.outputs + if output.row_limit is None } used_names: set[str] = set() nodes: list[IRModels.Node] = [] @@ -303,8 +324,168 @@ def extract_component_graph(spec: ComponentSplitSpec) -> IRModels.Graph: nodes.append(clone_component_node(spec, node, ref_replacements, output_logical_names, used_names)) - nodes.append(create_output_node(spec.name, output_names, nodes)) - return IRModels.rebuild_graph(tuple(nodes), spec.graph) + final_output_names = [] + + for output, output_name in zip(spec.outputs, output_names): + final_output_name = output_name + + if output.row_limit is not None: + final_output_name = add_row_limited_output_node(nodes, output_name, output.logical_name, output.row_limit, used_names) + + output_logical_names[final_output_name] = output.logical_name + final_output_names.append(final_output_name) + + nodes.append(create_output_node(spec.name, tuple(final_output_names), nodes)) + graph = IRModels.rebuild_graph(tuple(nodes), spec.graph) + + if spec.chunk_tokens is not None: + graph = retarget_chunk_graph_sequence_length(graph, spec.chunk_tokens) + + return graph + + +def add_row_limited_output_node( + nodes: list[IRModels.Node], + source_name: str, + logical_name: str, + row_limit: int, + used_names: set[str], +) -> str: + source = next((node for node in nodes if node.name == source_name), None) + + if source is None: + raise ValueError(f"Component split cannot add row-limited output for missing node {source_name}") + + source_shape = tensor_shape(source) + + if not source_shape: + raise ValueError(f"{source_name}: row-limited output requires tensor shape metadata") + + output_shape = list(source_shape) + output_shape[0] = min(row_limit, int(output_shape[0])) if isinstance(output_shape[0], int) else row_limit + name = unique_component_node_name(f"{source_name}_first_{row_limit}", used_names) + + source_meta = source.tensor_output_meta if isinstance(source.tensor_output_meta, dict) else {} + nodes.append( + IRModels.Node( + index=max((node.index for node in nodes), default=-1) + 1, + name=name, + node_type="call_function", + target="cactus.slice", + args=[{"node": source_name}], + kwargs={}, + users=(), + tensor_output_meta={**source_meta, "shape": output_shape}, + module_stack=source.module_stack, + value_kind=FModels.ValueKind.ACTIVATION, + attrs={"axis": 0, "start": 0, "length": row_limit, "step": 1}, + ir_metadata={"logical_output": logical_name}, + cache=None, + ) + ) + + return name + + +def retarget_chunk_graph_sequence_length(graph: IRModels.Graph, chunk_tokens: int) -> IRModels.Graph: + source_lengths = chunk_source_lengths(graph) + + if not source_lengths: + return graph + + cache_capacity = max(source_lengths) + nodes = tuple(retarget_chunk_node(node, source_lengths, chunk_tokens, cache_capacity) for node in graph.nodes) + return IRModels.rebuild_graph(nodes, graph) + + +def chunk_source_lengths(graph: IRModels.Graph) -> frozenset[int]: + lengths: set[int] = set() + + for node in graph.nodes: + if not node.is_placeholder: + continue + + logical_name = str(node.ir_metadata.get("logical_input") or node.target) + + if logical_name not in {"input_ids", "position_ids", "inputs_embeds", "per_layer_inputs"}: + continue + + shape = tensor_shape(node) + + if len(shape) >= 2 and shape[0] == 1 and isinstance(shape[1], int) and shape[1] > 1: + lengths.add(shape[1]) + + return frozenset(lengths) + + +def retarget_chunk_node( + node: IRModels.Node, + source_lengths: frozenset[int], + chunk_tokens: int, + cache_capacity: int, +) -> IRModels.Node: + metadata = dict(node.ir_metadata) + + if is_empty_cache_cat_node(node): + metadata.setdefault("prefill_cache_capacity", cache_capacity) + + if node.is_placeholder and metadata.get("logical_input") in {"input_ids", "position_ids", "inputs_embeds", "per_layer_inputs"}: + metadata.setdefault("prefill_chunk_tokens", chunk_tokens) + + return IRModels.Node( + index=node.index, + name=node.name, + node_type=node.node_type, + target=node.target, + args=retarget_sequence_value(node.args, source_lengths, chunk_tokens), + kwargs=retarget_sequence_value(node.kwargs, source_lengths, chunk_tokens), + users=node.users, + tensor_output_meta=retarget_sequence_value(node.tensor_output_meta, source_lengths, chunk_tokens), + module_stack=node.module_stack, + value_kind=node.value_kind, + attrs=retarget_sequence_value(dict(node.attrs), source_lengths, chunk_tokens), + ir_metadata=metadata, + cache=node.cache, + ) + + +def retarget_sequence_value(value: Any, source_lengths: frozenset[int], chunk_tokens: int) -> Any: + if isinstance(value, bool): + return value + + if isinstance(value, int): + return chunk_tokens if value in source_lengths else value + + if isinstance(value, list): + return [retarget_sequence_value(item, source_lengths, chunk_tokens) for item in value] + + if isinstance(value, tuple): + return tuple(retarget_sequence_value(item, source_lengths, chunk_tokens) for item in value) + + if isinstance(value, dict): + return {key: retarget_sequence_value(item, source_lengths, chunk_tokens) for key, item in value.items()} + + return value + + +def is_empty_cache_cat_node(node: IRModels.Node) -> bool: + if node.target not in {"cactus.cat", "aten.cat.default"} or len(node.parents) < 2: + return False + + first_parent = node.parents[0] + return first_parent.value_kind == FModels.ValueKind.LIFTED_CONSTANT and element_count(tensor_shape(first_parent)) == 0 + + +def element_count(shape: list[Any]) -> int | None: + product = 1 + + for dim in shape: + if not isinstance(dim, int): + return None + + product *= dim + + return product def collect_required_node_names( diff --git a/python/cactus/transpiler/Generator/constants.py b/python/cactus/transpiler/Generator/constants.py index fc64a644e..d20670e0a 100644 --- a/python/cactus/transpiler/Generator/constants.py +++ b/python/cactus/transpiler/Generator/constants.py @@ -2,6 +2,7 @@ DEFAULT_COMPONENT_NAME = "model" DEFAULT_GRAPH_SUFFIX = ".cactus" DEFAULT_INPUT_PRECISION = "FP16" +GEMMA4_MLP_PRODUCT_SCALE = 1.0 / 64.0 DTYPE_TO_PRECISION: dict[str, str] = { "torch.float16": "FP16", @@ -31,6 +32,23 @@ "cache_state", } +FP16_RUNTIME_INPUTS = { + "input_features", + "inputs_embeds", + "pixel_values", +} + +GEMMA_ADD_CLIPPED_COMPONENTS = { + "audio_encoder", + "vision_encoder", + "lm_encoder_step", + "lm_encoder_text_chunk", + "lm_encoder_media_step", + "lm_encoder_media_chunk", + "decoder_step", + "decoder_prefill_chunk", +} + PASS_THROUGH_TARGETS = { "aten._assert_tensor_metadata.default", "aten.clone.default", @@ -43,6 +61,7 @@ } PAD_TARGETS = { + "cactus.pad", "aten.constant_pad_nd.default", } @@ -138,6 +157,15 @@ "aten.bitwise_not.default": "bitwise_not", } +FP16_UNARY_METHODS = { + "relu", + "silu", + "gelu", + "gelu_erf", + "sigmoid", + "tanh", +} + LOG1P_TARGETS = { "aten.log1p.default", } @@ -170,6 +198,7 @@ } EXPAND_TARGETS = { + "cactus.expand", "aten.expand.default", } @@ -201,10 +230,12 @@ } WHERE_TARGETS = { + "cactus.where", "aten.where.self", } UNFOLD_TARGETS = { + "cactus.unfold", "aten.unfold.default", } @@ -234,6 +265,7 @@ } TO_COPY_TARGETS = { + "cactus.precision_cast", "aten._to_copy.default", } @@ -241,6 +273,11 @@ "aten.neg.default", } +CUMSUM_TARGETS = { + "cactus.cumsum", + "aten.cumsum.default", +} + SOFTMAX_TARGETS = { "cactus.softmax", "aten.softmax.int", diff --git a/python/cactus/transpiler/Generator/generate.py b/python/cactus/transpiler/Generator/generate.py index 3267fae4d..67a7ea35f 100644 --- a/python/cactus/transpiler/Generator/generate.py +++ b/python/cactus/transpiler/Generator/generate.py @@ -9,6 +9,8 @@ from pathlib import Path from typing import Any +import numpy as np + from . import component_split from . import constants from . import models @@ -310,6 +312,7 @@ def build_lowering_rules( add_rules(rules, constants.PAD_TARGETS, lower_constant_pad_nd) add_rules(rules, constants.CONSTANT_INPUT_TARGETS, lower_constant_input) add_rules(rules, constants.NEG_TARGETS, lower_neg) + add_rules(rules, constants.CUMSUM_TARGETS, lower_cumsum) add_rules(rules, constants.SOFTMAX_TARGETS, lower_softmax) add_rules(rules, constants.TOPK_TARGETS, lower_topk) add_rules(rules, constants.GATHER_TARGETS, lower_gather) @@ -388,14 +391,17 @@ def prefill_cache_cat_annotations(graph: IRModels.Graph) -> dict[str, IRModels.C continue _, num_kv_heads, sequence_length, head_dim = shape + cache_capacity = int(node.ir_metadata.get("prefill_cache_capacity", sequence_length)) + layer_index = index // 2 annotations[node.name] = IRModels.CacheAnnotation( kind=FModels.CacheKind.KV, role=FModels.CacheTensorRole.KEY if index % 2 == 0 else FModels.CacheTensorRole.VALUE, tensor_index=index, - layer_index=index // 2, + layer_index=layer_index, shape=tuple(shape), layout="batch_heads_sequence_head_dim", - sequence_length=int(sequence_length), + sequence_length=cache_capacity, + window_size=prefill_cache_window_size(graph, layer_index), num_kv_heads=int(num_kv_heads), head_dim=int(head_dim), source="prefill_empty_cache_cat", @@ -404,6 +410,15 @@ def prefill_cache_cat_annotations(graph: IRModels.Graph) -> dict[str, IRModels.C return annotations +def prefill_cache_window_size(graph: IRModels.Graph, layer_index: int) -> int | None: + model_name = str(getattr(graph, "model_name", "")).lower() + + if "gemma" not in model_name: + return None + + return 0 if layer_index % 5 == 4 else 512 + + def is_empty_cache_cat(node: IRModels.Node) -> bool: if node.target not in {"cactus.cat", "aten.cat.default"} or len(node.parents) < 2: return False @@ -526,6 +541,10 @@ def placeholder_precision(context: models.GenerationContext, node: IRModels.Node if record is not None and record.precision and hasattr(context.graph, record.precision): return int(getattr(context.graph, record.precision)) + logical_name = logical_input_name(context.component.ir_graph, node) + if logical_name in constants.FP16_RUNTIME_INPUTS and str(tensor_dtype(node)) in {"torch.float32", "torch.float"}: + return int(context.graph.FP16) + return cactus_precision(context.graph, tensor_dtype(node)) @@ -744,9 +763,15 @@ def lower_binary(context: models.GenerationContext, node: IRModels.Node) -> Any: if len(inputs) == 2: left, right = align_binary_inputs(context, node, inputs) - if method == "add" and context.component.name in {"decoder_step", "decoder_prefill_chunk"} and looks_like_gemma_residual_add(node): + if method == "add" and context.component.name in constants.GEMMA_ADD_CLIPPED_COMPONENTS and looks_like_gemma_residual_add(node): return context.graph.add_clipped(left, right) + if method == "multiply" and looks_like_gemma4_mlp_product(node): + if is_gemma4_mlp_gate_activation(node.parents[0]): + left = context.graph.scalar_multiply(left, constants.GEMMA4_MLP_PRODUCT_SCALE) + else: + right = context.graph.scalar_multiply(right, constants.GEMMA4_MLP_PRODUCT_SCALE) + return getattr(context.graph, method)(left, right) scalar = scalar_attr(node, "other") @@ -764,7 +789,7 @@ def looks_like_gemma_residual_add(node: IRModels.Node) -> bool: shape = meta_shape(node) - if not shape or len(shape) > 3: + if not shape or len(shape) > 4: return False left, right = node.parents @@ -783,6 +808,61 @@ def is_gemma_residual_branch(residual: IRModels.Node, branch: IRModels.Node) -> return strip_precision_passthrough(residual) is not strip_precision_passthrough(branch_source.parents[0]) +def looks_like_gemma4_mlp_product(node: IRModels.Node) -> bool: + if len(node.parents) != 2: + return False + + text = node_context_text(node) + + if "gemma" not in text.lower() or "language_model" not in text or ".mlp" not in text: + return False + + left, right = node.parents + return ( + is_gemma4_mlp_gate_activation(left) and has_ancestor_text(right, "mlp_up_proj") + ) or ( + is_gemma4_mlp_gate_activation(right) and has_ancestor_text(left, "mlp_up_proj") + ) + + +def is_gemma4_mlp_gate_activation(node: IRModels.Node) -> bool: + source = strip_precision_passthrough(node) + + if source.target not in {"cactus.gelu", "aten.gelu.default"}: + return False + + return has_ancestor_text(source, "mlp_gate_proj") + + +def has_ancestor_text(node: IRModels.Node, pattern: str, max_depth: int = 8) -> bool: + pattern = pattern.lower() + worklist: list[tuple[IRModels.Node, int]] = [(node, 0)] + seen: set[str] = set() + + while worklist: + current, depth = worklist.pop(0) + + if current.name in seen: + continue + + seen.add(current.name) + + if pattern in node_context_text(current).lower(): + return True + + if depth >= max_depth: + continue + + for parent in current.parents: + worklist.append((parent, depth + 1)) + + return False + + +def node_context_text(node: IRModels.Node) -> str: + return f"{node.name} {node.target} {node.module_stack!r}" + + def strip_precision_passthrough(node: IRModels.Node) -> IRModels.Node: current = node passthrough_targets = constants.PASS_THROUGH_TARGETS | constants.TO_COPY_TARGETS @@ -810,6 +890,9 @@ def lower_scalar(context: models.GenerationContext, node: IRModels.Node) -> Any: raise UnsupportedLoweringError(f"{node.name}: scalar lowering {node.target} missing attr {attr_name}") + if method == "scalar_multiply" and float(value) == 1.0: + return inputs[0] + return getattr(context.graph, method)(inputs[0], value) @@ -885,14 +968,28 @@ def resolve_declared_shape_from_actual(declared_shape: tuple[Any, ...], actual_s def lower_unary(context: models.GenerationContext, node: IRModels.Node) -> Any: inputs = require_input_count(context, node, 1) method = constants.UNARY_TARGETS[node.target] - return getattr(context.graph, method)(inputs[0]) + value = fp16_tensor(context, inputs[0]) if method in constants.FP16_UNARY_METHODS else inputs[0] + return getattr(context.graph, method)(value) def lower_log1p(context: models.GenerationContext, node: IRModels.Node) -> Any: inputs = require_input_count(context, node, 1) + + if node.parents and node.parents[0].target in {"aten.exp.default", "cactus.scalar_exp"} and node.parents[0].parents: + x = context.require(node.parents[0].parents[0].name) + return lower_stable_softplus(context, x) + return context.graph.scalar_log(context.graph.scalar_add(inputs[0], 1.0)) +def lower_stable_softplus(context: models.GenerationContext, x: Any) -> Any: + x = cast_to_precision(context, x, context.graph.FP16) + neg_abs = context.graph.scalar_multiply(context.graph.abs(x), -1.0) + log_term = context.graph.scalar_log(context.graph.scalar_add(context.graph.scalar_exp(neg_abs), 1.0)) + positive_term = context.graph.clamp(x, 0.0, math.inf) + return context.graph.add(log_term, positive_term) + + def lower_pow(context: models.GenerationContext, node: IRModels.Node) -> Any: inputs = require_input_count(context, node, 1) exponent = scalar_attr(node, "exponent") @@ -928,7 +1025,7 @@ def lower_reduce(context: models.GenerationContext, node: IRModels.Node) -> Any: def lower_shape(context: models.GenerationContext, node: IRModels.Node) -> Any: inputs = require_input_count(context, node, 1) method = constants.SHAPE_TARGETS[node.target] - shape = shape_attr(node) + shape = decoder_prefill_logits_shape(context, node, shape_attr(node)) return getattr(context.graph, method)(inputs[0], shape) @@ -1058,6 +1155,9 @@ def lower_slice(context: models.GenerationContext, node: IRModels.Node) -> Any: start = int(node.attrs.get("start", 0) or 0) step = node.attrs.get("step", 1) + if is_full_axis_slice(node, axis, start, step): + return inputs[0] + if step is not None and int(step) != 1: target_shape = output_shape(node) length = target_shape[axis] if axis < len(target_shape) else slice_length(node, start) @@ -1067,6 +1167,19 @@ def lower_slice(context: models.GenerationContext, node: IRModels.Node) -> Any: return context.graph.slice(inputs[0], axis, start, length=length) +def is_full_axis_slice(node: IRModels.Node, axis: int, start: int, step: Any) -> bool: + if start != 0: + return False + + if step is not None and int(step) != 1: + return False + + if not node.parents: + return False + + return shape_matches_tensor(node.parents[0], output_shape(node)) + + def lower_index(context: models.GenerationContext, node: IRModels.Node) -> Any: if node.target == "aten.index.Tensor": return lower_index_tensor(context, node) @@ -1089,7 +1202,6 @@ def lower_index_tensor(context: models.GenerationContext, node: IRModels.Node) - source = inputs[0] source_shape = meta_shape(node.parents[0]) target_shape = meta_shape(node) - if len(source_shape) == 3 and len(target_shape) == 2 and concrete_dim(source_shape[0]) == 1: feature_count = concrete_dim(source_shape[1]) feature_dim = concrete_dim(source_shape[2]) @@ -1157,7 +1269,12 @@ def lower_cat(context: models.GenerationContext, node: IRModels.Node) -> Any: return passthrough axis = normalize_dim(int(node.attrs.get("axis", node.attrs.get("dim", 0))), cat_rank(node)) - return context.graph.cat(inputs, axis=axis) + return context.graph.cat(fp16_cat_inputs(context, inputs), axis=axis) + + +def fp16_cat_inputs(context: models.GenerationContext, inputs: tuple[Any, ...]) -> tuple[Any, ...]: + fp16 = int(context.graph.FP16) + return tuple(value if getattr(value, "dtype", None) == fp16 else context.graph.precision_cast(value, fp16) for value in inputs) def lower_decode_cache_cat(context: models.GenerationContext, node: IRModels.Node) -> Any | None: @@ -1180,7 +1297,7 @@ def lower_prefill_cache_cat(context: models.GenerationContext, node: IRModels.No new_kv_node = node.parents[1] original_new_kv = context.require(new_kv_node.name) - new_kv = cache_new_kv_for_native(context, new_kv_node, original_new_kv) + new_kv = cache_new_kv_for_native(context, new_kv_node, original_new_kv, annotation.role or "") new_kv = cast_to_precision(context, new_kv, context.graph.FP16) cache_state = context.graph.kv_cache_state( kv_cache_capacity(context, annotation), @@ -1192,6 +1309,7 @@ def lower_prefill_cache_cat(context: models.GenerationContext, node: IRModels.No ) record_cache_state_binding(context, node, cache_state, annotation) context.graph.kv_cache_append(new_kv, cache_state, window_size=int(annotation.window_size or 0), sink_size=0) + context.prefill_cache_cat_states[node.name] = cache_state return original_new_kv @@ -1203,6 +1321,7 @@ def lower_matmul(context: models.GenerationContext, node: IRModels.Node) -> Any: lhs = matmul_activation_operand(context, inputs[0]) rhs = inputs[1] + lhs = slice_decoder_prefill_logits_lhs(context, node, lhs) if len(node.parents) >= 2 and is_weight_transpose_node(node.parents[1]): original_weight_node = node.parents[1].parents[0] @@ -1228,6 +1347,47 @@ def matmul_activation_operand(context: models.GenerationContext, value: Any) -> return value +def slice_decoder_prefill_logits_lhs(context: models.GenerationContext, node: IRModels.Node, lhs: Any) -> Any: + if not is_decoder_prefill_logits_matmul(context, node): + return lhs + + lhs_shape = meta_shape(node.parents[0]) if node.parents else () + + if len(lhs_shape) != 2 or not isinstance(lhs_shape[0], int) or lhs_shape[0] <= 1: + return lhs + + return context.graph.slice(lhs, 0, int(lhs_shape[0]) - 1, length=1) + + +def is_decoder_prefill_logits_matmul(context: models.GenerationContext, node: IRModels.Node) -> bool: + shape = meta_shape(node) + + return ( + context.component.name == "decoder_prefill_chunk" + and node.target in constants.MATMUL_TARGETS + and len(shape) == 2 + and len(node.parents) >= 2 + and isinstance(shape[0], int) + and shape[0] > 1 + and isinstance(shape[-1], int) + and shape[-1] >= 100000 + ) + + +def decoder_prefill_logits_shape(context: models.GenerationContext, node: IRModels.Node, shape: tuple[int, ...]) -> tuple[int, ...]: + if ( + context.component.name == "decoder_prefill_chunk" + and len(shape) == 3 + and shape[0] == 1 + and isinstance(shape[1], int) + and shape[1] > 1 + and shape[2] >= 100000 + ): + return (1, 1, shape[2]) + + return shape + + def empty_cat_passthrough(node: IRModels.Node, inputs: tuple[Any, ...]) -> Any | None: target_shape = meta_shape(node) non_empty_inputs = [] @@ -1635,6 +1795,16 @@ def lower_neg(context: models.GenerationContext, node: IRModels.Node) -> Any: return context.graph.scalar_multiply(inputs[0], -1.0) +def lower_cumsum(context: models.GenerationContext, node: IRModels.Node) -> Any: + inputs = require_input_count(context, node, 1) + axis = axis_attr(node) + + if axis is None: + raise UnsupportedLoweringError(f"{node.name}: cumsum lowering missing axis/dim attr") + + return context.graph.cumsum(inputs[0], normalize_dim(axis, len(meta_shape(node.parents[0])))) + + def lower_softmax(context: models.GenerationContext, node: IRModels.Node) -> Any: inputs = require_input_count(context, node, 1) axis = axis_attr(node, default=-1) @@ -1642,7 +1812,12 @@ def lower_softmax(context: models.GenerationContext, node: IRModels.Node) -> Any if axis is None: raise UnsupportedLoweringError(f"{node.name}: softmax lowering missing axis/dim attr") - return context.graph.softmax(inputs[0], axis=normalize_dim(axis, len(meta_shape(node.parents[0])))) + return context.graph.softmax(fp16_tensor(context, inputs[0]), axis=normalize_dim(axis, len(meta_shape(node.parents[0])))) + + +def fp16_tensor(context: models.GenerationContext, value: Any) -> Any: + fp16 = int(context.graph.FP16) + return value if getattr(value, "dtype", None) == fp16 else context.graph.precision_cast(value, fp16) def lower_topk(context: models.GenerationContext, node: IRModels.Node) -> Any: @@ -1668,27 +1843,36 @@ def lower_embedding(context: models.GenerationContext, node: IRModels.Node) -> A def lower_clamp(context: models.GenerationContext, node: IRModels.Node) -> Any: inputs = require_input_count(context, node, 1) - lo = node.attrs.get("lo", node.attrs.get("min", node.attrs.get("arg_1"))) - hi = node.attrs.get("hi", node.attrs.get("max", node.attrs.get("arg_2"))) + lo = scalar_weight_bound_value(context, node.parents[1]) if len(node.parents) > 1 else None + hi = scalar_weight_bound_value(context, node.parents[2]) if len(node.parents) > 2 else None + + if lo is None: + lo = inputs[1] if len(inputs) > 1 else node.attrs.get("lo", node.attrs.get("min", node.attrs.get("arg_1"))) + + if hi is None: + hi = inputs[2] if len(inputs) > 2 else node.attrs.get("hi", node.attrs.get("max", node.attrs.get("arg_2"))) if lo is None and hi is None: return inputs[0] result = inputs[0] + if isinstance(lo, (int, float)) and isinstance(hi, (int, float)): + return context.graph.clamp(result, float(lo), float(hi)) + if lo is not None: - keep = context.graph.scalar_greater_equal(result, lo) - result = context.graph.add( - context.graph.multiply(result, keep), - context.graph.scalar_multiply(context.graph.logical_not(keep), lo), - ) + if isinstance(lo, (int, float)): + result = context.graph.clamp(result, float(lo), math.inf) + else: + lo = cast_to_precision(context, lo, getattr(result, "dtype", context.graph.FP16)) + result = context.graph.where(context.graph.greater_equal(result, lo), result, lo) if hi is not None: - keep = context.graph.scalar_less_equal(result, hi) - result = context.graph.add( - context.graph.multiply(result, keep), - context.graph.scalar_multiply(context.graph.logical_not(keep), hi), - ) + if isinstance(hi, (int, float)): + result = context.graph.clamp(result, -math.inf, float(hi)) + else: + hi = cast_to_precision(context, hi, getattr(result, "dtype", context.graph.FP16)) + result = context.graph.where(context.graph.less_equal(result, hi), result, hi) return result @@ -1698,8 +1882,11 @@ def lower_norm(context: models.GenerationContext, node: IRModels.Node) -> Any: target = node.target if target == "cactus.rms_norm": - require_len(node, inputs, 2) - return lower_rms_norm(context, node, inputs[0], inputs[1]) + if len(inputs) not in {1, 2}: + raise unsupported_arity(node, len(inputs), "activation plus optional weight") + + weight = inputs[1] if len(inputs) == 2 else None + return lower_rms_norm(context, node, inputs[0], weight) if target in {"cactus.layernorm", "cactus.layer_norm", "aten.native_layer_norm.default", "aten.layer_norm.default"}: require_len(node, inputs, 2) @@ -1717,14 +1904,25 @@ def lower_norm(context: models.GenerationContext, node: IRModels.Node) -> Any: raise UnsupportedLoweringError(f"{node.name}: unsupported norm target {target}") -def lower_rms_norm(context: models.GenerationContext, node: IRModels.Node, x: Any, weight: Any) -> Any: +def lower_rms_norm(context: models.GenerationContext, node: IRModels.Node, x: Any, weight: Any | None = None) -> Any: x = cast_to_precision(context, x, context.graph.FP16) - weight = cast_to_precision(context, weight, context.graph.FP16) - input_shape = concrete_shape(meta_shape(node.parents[0])) if node.parents else concrete_shape(meta_shape(node)) + declared_input_shape = meta_shape(node.parents[0]) if node.parents else meta_shape(node) + input_shape = concrete_shape(declared_input_shape) + + if input_shape is None: + actual_shape = tuple(getattr(x, "shape", ())) + + if actual_shape and concrete_dim(declared_input_shape[-1] if declared_input_shape else None) is not None: + input_shape = actual_shape if input_shape is None: raise UnsupportedLoweringError(f"{node.name}: cactus.rms_norm requires concrete input shape") + if weight is None: + weight = rms_norm_unit_weight(context, int(input_shape[-1])) + else: + weight = cast_to_precision(context, weight, context.graph.FP16) + if len(input_shape) == 2: return context.graph.rms_norm(x, weight, eps=epsilon_attr(node)) @@ -1743,6 +1941,20 @@ def lower_rms_norm(context: models.GenerationContext, node: IRModels.Node, x: An return context.graph.reshape(normalized, input_shape) +def rms_norm_unit_weight(context: models.GenerationContext, hidden_dim: int) -> Any: + key = f"__embedded_rms_norm_unit_weight_{hidden_dim}" + existing = context.lookup(key) + + if existing is not None: + return existing + + weight = context.graph.input((hidden_dim,), dtype=context.graph.FP16) + context.graph.set_input(weight, np.ones((hidden_dim,), dtype=np.float16)) + context.graph.mark_embedded_input(weight) + context.values[key] = weight + return weight + + def lower_conv(context: models.GenerationContext, node: IRModels.Node) -> Any: inputs = context.inputs_for(node) target = node.target @@ -1902,6 +2114,7 @@ def lower_attention(context: models.GenerationContext, node: IRModels.Node) -> A query = cast_to_precision(context, query, context.graph.FP16) key = cast_to_precision(context, key, context.graph.FP16) value = cast_to_precision(context, value, context.graph.FP16) + mask = cast_to_precision(context, mask, context.graph.FP16) if mask is not None else None output = context.graph.attention( query, key, @@ -1937,8 +2150,8 @@ def lower_cache_backed_attention(context: models.GenerationContext, node: IRMode if node.target != "cactus.attention" or len(node.parents) < 3: return None - key_cat = find_cache_cat_ancestor(node.parents[1], FModels.CacheTensorRole.KEY) - value_cat = find_cache_cat_ancestor(node.parents[2], FModels.CacheTensorRole.VALUE) + key_cat = find_attention_cache_cat_ancestor(context, node.parents[1], FModels.CacheTensorRole.KEY) + value_cat = find_attention_cache_cat_ancestor(context, node.parents[2], FModels.CacheTensorRole.VALUE) if key_cat is None or value_cat is None: return None @@ -1949,10 +2162,17 @@ def lower_cache_backed_attention(context: models.GenerationContext, node: IRMode value_new_node = value_cat.parents[1] query = context.require(node.parents[0].name) - key_new = cache_new_kv_for_native(context, key_new_node, context.require(key_new_node.name)) - value_new = cache_new_kv_for_native(context, value_new_node, context.require(value_new_node.name)) - key_cache_state = context.require(key_cache.name) - value_cache_state = context.require(value_cache.name) + input_layout = str(node.attrs.get("input_layout", "")) + key_new = cache_new_kv_for_native(context, key_new_node, context.require(key_new_node.name), FModels.CacheTensorRole.KEY, input_layout) + value_new = cache_new_kv_for_native(context, value_new_node, context.require(value_new_node.name), FModels.CacheTensorRole.VALUE, input_layout) + key_cache_state = prefill_cache_state_for_cat(context, key_cat) + value_cache_state = prefill_cache_state_for_cat(context, value_cat) + + if key_cache_state is None: + key_cache_state = context.require(key_cache.name) + + if value_cache_state is None: + value_cache_state = context.require(value_cache.name) if node.attrs.get("input_layout") == "bhqd_bhds_bhsd": query = context.graph.permute(query, (0, 2, 1, 3)) @@ -1962,12 +2182,16 @@ def lower_cache_backed_attention(context: models.GenerationContext, node: IRMode value_new = cast_to_precision(context, value_new, context.graph.FP16) cache_pair = (key_cat.name, value_cat.name) + prefill_pair_already_appended = key_cat.name in context.prefill_cache_cat_states and value_cat.name in context.prefill_cache_cat_states - if cache_pair not in context.appended_cache_pairs: + if cache_pair not in context.appended_cache_pairs and not prefill_pair_already_appended: context.graph.kv_cache_append(key_new, key_cache_state, window_size=int(node.attrs.get("window_size", 0)), sink_size=0) context.graph.kv_cache_append(value_new, value_cache_state, window_size=int(node.attrs.get("window_size", 0)), sink_size=0) context.appended_cache_pairs.add(cache_pair) + context.values.setdefault(key_cat.name, key_cache_state) + context.values.setdefault(value_cat.name, value_cache_state) + output = context.graph.attention_cached( query, key_new, @@ -1982,6 +2206,34 @@ def lower_cache_backed_attention(context: models.GenerationContext, node: IRMode return attention_output_for_layout(context, node, output) +def find_attention_cache_cat_ancestor(context: models.GenerationContext, node: IRModels.Node, role: str) -> IRModels.Node | None: + cache_cat = find_cache_cat_ancestor(node, role) + + if cache_cat is not None: + return cache_cat + + return find_prefill_cache_cat_ancestor(context, node) + + +def find_prefill_cache_cat_ancestor(context: models.GenerationContext, node: IRModels.Node, max_depth: int = 16) -> IRModels.Node | None: + current = node + + for _ in range(max_depth): + if current.name in context.prefill_cache_cat_annotations: + return current + + if len(current.parents) != 1: + return None + + current = current.parents[0] + + return None + + +def prefill_cache_state_for_cat(context: models.GenerationContext, node: IRModels.Node) -> Any | None: + return context.prefill_cache_cat_states.get(node.name) + + def cached_attention_position_offset(context: models.GenerationContext, node: IRModels.Node) -> int: if context.component.name == "decoder_step": return SIZE_T_MAX @@ -1989,7 +2241,7 @@ def cached_attention_position_offset(context: models.GenerationContext, node: IR return int(node.attrs.get("position_offset", 0)) -def cache_new_kv_for_native(context: models.GenerationContext, node: IRModels.Node, value: Any) -> Any: +def cache_new_kv_for_native(context: models.GenerationContext, node: IRModels.Node, value: Any, role: str, input_layout: str = "") -> Any: if len(meta_shape(node)) == 4: return context.graph.permute(value, (0, 2, 1, 3)) @@ -2016,11 +2268,24 @@ def attention_inputs_for_layout(context: models.GenerationContext, node: IRModel mask, ) + if node.attrs.get("input_layout") == "bqhd_bhds_bhsd": + return ( + inputs[0], + context.graph.permute(inputs[1], (0, 3, 1, 2)), + context.graph.permute(inputs[2], (0, 2, 1, 3)), + mask, + ) + return inputs[0], inputs[1], inputs[2], mask def attention_output_for_layout(context: models.GenerationContext, node: IRModels.Node, output: Any) -> Any: - if node.attrs.get("output_layout") == "bhqd": + output_layout = node.attrs.get("output_layout") + + if output_layout == "bthd_flat": + return context.graph.reshape(output, output_shape(node)) + + if output_layout == "bhqd": return context.graph.permute(output, (0, 2, 1, 3)) return output @@ -2207,6 +2472,10 @@ def lower_moe(context: models.GenerationContext, node: IRModels.Node) -> Any: ) if node.target == "cactus.moe_layer_gated": + packed = lower_packed_lfm_moe_layer_gated(context, node, inputs) + if packed is not None: + return packed + if len(inputs) < 6: raise UnsupportedLoweringError( f"{node.name}: cactus.moe_layer_gated requires hidden, routing_probs, topk_indices, w1, w3, and w2 inputs; " @@ -2247,6 +2516,119 @@ def lower_moe(context: models.GenerationContext, node: IRModels.Node) -> Any: raise UnsupportedLoweringError(f"{node.name}: unsupported MoE target {node.target}") +def lower_packed_lfm_moe_layer_gated( + context: models.GenerationContext, + node: IRModels.Node, + inputs: tuple[Any, ...], +) -> Any | None: + if len(inputs) != 5 or len(node.parents) != 5: + return None + + hidden_node, router_node, bias_node, gate_up_node, down_node = node.parents + gate_up_shape = concrete_shape(meta_shape(gate_up_node)) + down_shape = concrete_shape(meta_shape(down_node)) + num_experts = required_int_attr(node, "num_experts") + top_k = required_int_attr(node, "num_experts_per_tok") + + if gate_up_shape is None or down_shape is None: + return None + + if len(gate_up_shape) != 3 or len(down_shape) != 3: + return None + + if gate_up_shape[0] != num_experts or down_shape[0] != num_experts: + return None + + packed_intermediate, hidden_dim = gate_up_shape[1], gate_up_shape[2] + if packed_intermediate % 2 != 0: + raise UnsupportedLoweringError(f"{node.name}: packed gated MoE gate/up dimension must be even") + + intermediate_dim = packed_intermediate // 2 + + if down_shape[1:] != (hidden_dim, intermediate_dim): + raise UnsupportedLoweringError( + f"{node.name}: expected LFM packed down_proj shape " + f"[experts, hidden, intermediate], got {down_shape}" + ) + + hidden = moe_hidden_2d(context, node, inputs[0], hidden_dim) + router_weight = inputs[1] + expert_bias = cast_to_precision(context, inputs[2], context.graph.FP16) + gate_up_weight = inputs[3] + down_weight = inputs[4] + + router_logits = context.graph.matmul(hidden, router_weight, pretransposed_rhs=True) + router_logits = context.graph.add(router_logits, expert_bias) + routing_probs = context.graph.softmax(cast_to_precision(context, router_logits, context.graph.FP16), axis=-1) + topk_indices = context.graph.index(context.graph.topk(routing_probs, top_k), 0, axis=0) + w1_weights, w3_weights, w2_weights = split_lfm_packed_moe_weights( + context, + gate_up_weight, + down_weight, + num_experts, + intermediate_dim, + hidden_dim, + ) + + return context.graph.moe_layer_gated( + hidden, + routing_probs, + topk_indices, + w1_weights, + w3_weights, + w2_weights, + num_experts, + top_k, + normalize_routing=bool(node.attrs.get("normalize_routing", True)), + epsilon=float(node.attrs.get("epsilon", 1e-6)), + routed_scaling_factor=float(node.attrs.get("routed_scaling_factor", 1.0)), + ) + + +def moe_hidden_2d(context: models.GenerationContext, node: IRModels.Node, hidden: Any, hidden_dim: int) -> Any: + hidden_shape = concrete_shape(meta_shape(node.parents[0])) if node.parents else None + + if hidden_shape is None: + return cast_to_precision(context, hidden, context.graph.FP16) + + if len(hidden_shape) == 2: + return cast_to_precision(context, hidden, context.graph.FP16) + + if len(hidden_shape) == 3 and hidden_shape[0] == 1 and hidden_shape[2] == hidden_dim: + return cast_to_precision(context, context.graph.reshape(hidden, (hidden_shape[1], hidden_dim)), context.graph.FP16) + + raise UnsupportedLoweringError(f"{node.name}: packed MoE hidden input must be [tokens, hidden] or [1, tokens, hidden]") + + +def split_lfm_packed_moe_weights( + context: models.GenerationContext, + gate_up_weight: Any, + down_weight: Any, + num_experts: int, + intermediate_dim: int, + hidden_dim: int, +) -> tuple[tuple[Any, ...], tuple[Any, ...], tuple[Any, ...]]: + w1_weights = [] + w3_weights = [] + w2_weights = [] + + for expert_index in range(num_experts): + expert_gate_up = context.graph.reshape( + context.graph.slice(gate_up_weight, 0, expert_index, length=1), + (2 * intermediate_dim, hidden_dim), + ) + w1_weights.append(context.graph.slice(expert_gate_up, 0, 0, length=intermediate_dim)) + w3_weights.append(context.graph.slice(expert_gate_up, 0, intermediate_dim, length=intermediate_dim)) + + expert_down = context.graph.reshape( + context.graph.slice(down_weight, 0, expert_index, length=1), + (hidden_dim, intermediate_dim), + ) + w2_weights.append(context.graph.transpose(expert_down)) + + return tuple(w1_weights), tuple(w3_weights), tuple(w2_weights) + + def lower_special_cactus(context: models.GenerationContext, node: IRModels.Node) -> Any: inputs = context.inputs_for(node) target = node.target @@ -2601,6 +2983,67 @@ def scalar_attr(node: IRModels.Node, name: str) -> Any | None: return value +def scalar_weight_bound_value(context: models.GenerationContext, node: IRModels.Node) -> float | None: + resolver = context.component.weight_resolver + + if resolver is None: + return None + + record = resolver.resolve(node.name) + + if record is None or record.output_name is None: + return None + + if record.shape and math.prod(int(dim) for dim in record.shape) != 1: + return None + + value = read_scalar_cactus_weight(resolver.weights_dir / record.output_name) + + if value is None: + return None + + return float(value) * float(record.scale_factor) + + +def read_scalar_cactus_weight(path: Path) -> float | None: + if not path.exists(): + return None + + with path.open("rb") as f: + header = f.read(84) + + if len(header) != 84 or header[:4] != b"CACT": + return None + + fields = struct.unpack("= 2: + return float(struct.unpack("= 4: + return float(struct.unpack(" int: + if alignment <= 0: + alignment = 32 + + remainder = offset % alignment + return offset if remainder == 0 else offset + alignment - remainder + + def attr_value(node: IRModels.Node, name: str, default: Any) -> Any: value = node.attrs.get(name) return default if value is None else value diff --git a/python/cactus/transpiler/Generator/models.py b/python/cactus/transpiler/Generator/models.py index 3afa71776..de7134251 100644 --- a/python/cactus/transpiler/Generator/models.py +++ b/python/cactus/transpiler/Generator/models.py @@ -102,6 +102,7 @@ class GenerationContext: skip_node_names: frozenset[str] = frozenset() cache_state_placeholder_names: frozenset[str] = frozenset() prefill_cache_cat_annotations: dict[str, IRModels.CacheAnnotation] = field(default_factory=dict) + prefill_cache_cat_states: dict[str, GraphTensor] = field(default_factory=dict) appended_cache_pairs: set[tuple[str, str]] = field(default_factory=set) @property @@ -258,9 +259,26 @@ def component_graph_from_ir(cls: type[ComponentGraph], name: str, ir_graph: IRMo ir_graph=ir_graph, output_path=config.component_path(name), manifest_path=generator_component_manifest_path(config, name), + metadata=component_metadata_from_ir(ir_graph), ) +def component_metadata_from_ir(ir_graph: IRModels.Graph) -> dict[str, str]: + metadata: dict[str, str] = {} + + for node in ir_graph.nodes: + chunk_tokens = node.ir_metadata.get("prefill_chunk_tokens") + + if chunk_tokens is None: + continue + + metadata["prefill_strategy"] = "chunked" + metadata["prefill_chunk_tokens"] = str(chunk_tokens) + break + + return metadata + + def save_component_graph(component: ComponentGraph) -> Path: if component.graph is None: raise ValueError(f"Component {component.name} has no Cactus graph to save") diff --git a/python/cactus/transpiler/IR/extra_matchers.py b/python/cactus/transpiler/IR/extra_matchers.py index f0f4d2d50..2762088ef 100644 --- a/python/cactus/transpiler/IR/extra_matchers.py +++ b/python/cactus/transpiler/IR/extra_matchers.py @@ -1537,6 +1537,107 @@ def match_same_layer(source: models.Node, graph: models.Graph, fusion: FModels.F return len(set(layer_keys)) == 1 +def match_slice_halves(source: models.Node, graph: models.Graph, fusion: FModels.FusionGraph, bindings: dict[str, models.Node], spec: dict[str, Any]) -> bool: + """ + Checks that two slice nodes split the same tensor into adjacent equal halves. + + This protects decomposed GLU fusion. GLU is valid when the exported graph + computes `left * sigmoid(right)` where `left` and `right` are the first and + second halves of one source tensor along the same feature axis. + """ + left = bindings.get(spec["left_node"]) + right = bindings.get(spec["right_node"]) + + if left is None or right is None: + return False + + left_source = match_utils.get_parent(left, 0) + right_source = match_utils.get_parent(right, 0) + + if left_source is None or left_source is not right_source: + return False + + source_shape = match_utils.get_tensor_shape(left_source) + left_shape = match_utils.get_tensor_shape(left) + right_shape = match_utils.get_tensor_shape(right) + + if not source_shape or not left_shape or not right_shape: + return bool(spec.get("allow_missing", False)) + + left_axis = normalized_slice_axis(left, len(source_shape)) + right_axis = normalized_slice_axis(right, len(source_shape)) + + if left_axis is match_utils.MISSING or right_axis is match_utils.MISSING or left_axis != right_axis: + return False + + if not slice_step_is_one(left) or not slice_step_is_one(right): + return False + + left_size = match_utils.get_dim(left_shape, left_axis) + right_size = match_utils.get_dim(right_shape, right_axis) + source_size = match_utils.get_dim(source_shape, left_axis) + + if left_size is None or right_size is None: + return bool(spec.get("allow_missing", False)) + + if not match_utils.values_equal(left_size, right_size): + return False + + if not attr_matches(left, "start", 0, default=0): + return False + + if not attr_matches(right, "start", left_size): + return False + + if not optional_attr_matches(left, "end", left_size): + return False + + if not optional_attr_matches(right, "end", add_dims(left_size, right_size)): + return False + + if source_size is not None and not match_utils.values_equal(source_size, add_dims(left_size, right_size)): + return False + + return True + + +def normalized_slice_axis(node: models.Node, rank: int) -> Any: + axis = node.attrs.get("axis", node.attrs.get("dim", match_utils.MISSING)) + + if axis is match_utils.MISSING: + return match_utils.MISSING + + return normalize_dim(axis, rank) + + +def slice_step_is_one(node: models.Node) -> bool: + step = node.attrs.get("step", 1) + return step is None or match_utils.values_equal(step, 1) + + +def attr_matches(node: models.Node, attr: str, expected: Any, *, default: Any = match_utils.MISSING) -> bool: + actual = node.attrs.get(attr, default) + + if actual is match_utils.MISSING: + return False + + return match_utils.values_equal(actual, expected) + + +def optional_attr_matches(node: models.Node, attr: str, expected: Any) -> bool: + actual = node.attrs.get(attr, match_utils.MISSING) + + if actual is match_utils.MISSING or actual is None: + return True + + return match_utils.values_equal(actual, expected) + + +def add_dims(left: Any, right: Any) -> Any: + if isinstance(left, int) and isinstance(right, int): + return left + right + + return left EXTRA_MATCHERS: dict[str, ExtraMatcher] = { @@ -1558,4 +1659,5 @@ def match_same_layer(source: models.Node, graph: models.Graph, fusion: FModels.F "parent_tensor_dim_equal": match_parent_tensor_dim_equal, "input_broadcastable_to_node": match_input_broadcastable_to_node, "same_layer": match_same_layer, + "slice_halves": match_slice_halves, } diff --git a/python/cactus/transpiler/IR/models.py b/python/cactus/transpiler/IR/models.py index e2e5a121a..6a3e22923 100644 --- a/python/cactus/transpiler/IR/models.py +++ b/python/cactus/transpiler/IR/models.py @@ -252,7 +252,7 @@ def generate_graph(cls: type[Graph], layer_map: CModels.LayerMap) -> Graph: node.parents = tuple(parent_lists[node.name]) node.children = tuple(child_lists[node.name]) - annotate_cache_nodes(nodes) + annotate_cache_nodes(nodes, layer_map.model_name) sources = tuple(node for node in nodes if not node.parents) outputs = tuple(node for node in nodes if node.is_output) @@ -336,6 +336,7 @@ def sanitize_node_name(name: str) -> str: def fused_node_from_result(result: FusionResult, fused_name: str) -> Node: + attrs = normalized_fusion_attrs(result) cache = cache_annotation_from_fusion_result(result) ir_metadata = fusion_result_metadata(result, cache) @@ -345,17 +346,44 @@ def fused_node_from_result(result: FusionResult, fused_name: str) -> Node: node_type="call_function", target=result.target, args=[{"node": node.name} for node in result.external_inputs], - kwargs=dict(result.attrs), + kwargs=attrs, users=(), tensor_output_meta=result.source.tensor_output_meta, module_stack=result.source.module_stack, value_kind=FModels.ValueKind.CACHE_OUTPUT if cache is not None else FModels.ValueKind.ACTIVATION, - attrs=dict(result.attrs), + attrs=attrs, ir_metadata=ir_metadata, cache=cache, ) +def normalized_fusion_attrs(result: FusionResult) -> dict[str, Any]: + attrs = dict(result.attrs) + + if result.target in {"cactus.attention", "cactus.attention_cached"} and int(attrs.get("window_size", 0) or 0) == 0: + layer_index = gemma_language_layer_index(result) + + if layer_index is not None: + attrs["window_size"] = 0 if layer_index % 5 == 4 else 512 + + return attrs + + +def gemma_language_layer_index(result: FusionResult) -> int | None: + for node in (result.source, *result.matched_nodes, *result.external_inputs): + text = f"{node.name} {node.target} {node.module_stack!r}" + + if "gemma" not in text.lower() or "language_model" not in text: + continue + + layer_index = layer_index_from_text(text) + + if layer_index is not None: + return layer_index + + return None + + def expanded_consumed_names(graph: Graph, replacement_names: dict[str, str]) -> frozenset[str]: consumed_names = set(replacement_names) changed = True @@ -620,7 +648,7 @@ def topological_sort(graph: Graph) -> tuple[Node, ...]: return tuple(ordered) -def annotate_cache_nodes(nodes: tuple[Node, ...]) -> None: +def annotate_cache_nodes(nodes: tuple[Node, ...], model_name: str = "") -> None: kv_count = 0 conv_count = 0 @@ -650,6 +678,7 @@ def annotate_cache_nodes(nodes: tuple[Node, ...]) -> None: layout="batch_heads_sequence_head_dim", num_kv_heads=known_int(shape[1]), sequence_length=known_int(shape[2]), + window_size=kv_cache_window_size_for_model(model_name, cache_layer_index_from_node(node, kv_count // 2)), head_dim=known_int(shape[3]), source="past_key_values_placeholder", ) @@ -672,6 +701,16 @@ def annotate_cache_nodes(nodes: tuple[Node, ...]) -> None: node.ir_metadata["cache"] = cache_annotation_to_dict(node.cache) +def kv_cache_window_size_for_model(model_name: str, layer_index: int | None) -> int | None: + if layer_index is None: + return None + + if "gemma" not in model_name.lower(): + return None + + return 0 if layer_index % 5 == 4 else 512 + + def cache_annotation_from_metadata(metadata: dict[str, Any]) -> CacheAnnotation | None: raw_cache = metadata.get("cache") diff --git a/python/cactus/transpiler/IR/simplify_ir.py b/python/cactus/transpiler/IR/simplify_ir.py index 94900ed94..34c82aa3a 100644 --- a/python/cactus/transpiler/IR/simplify_ir.py +++ b/python/cactus/transpiler/IR/simplify_ir.py @@ -1,7 +1,7 @@ from collections.abc import Iterable from pathlib import Path -from . import match, match_utils, models +from . import match, match_utils, models, special_fusions from ..Converter import models as CModels from ..Fusions import fusions as Fusions from ..Fusions import models as FModels @@ -40,7 +40,6 @@ def simplify( total_fusions += len(simplified_graph.fusions) graph = simplified_graph - graph = fuse_gemma4_decode_attention(graph, layer_map.task, fusion_fields) return graph.remove_noop_nodes().to_layer_map() @@ -116,19 +115,32 @@ def try_match_from_node( return None for fusion in candidate_fusions_for_node(node, fusion_fields): - bindings = bind_matching_fusion( - node, - graph, - fusion, - inference_mode=inference_mode, - input_modalities=input_modalities, - fusion_fields=fusion_fields, - ) - - if bindings is None: - continue + if special_fusions.has_special_matcher(fusion): + result = special_fusions.match_special_fusion( + node, + graph, + fusion, + inference_mode=inference_mode, + input_modalities=input_modalities, + fusion_fields=fusion_fields, + ) + + if result is None: + continue + else: + bindings = bind_matching_fusion( + node, + graph, + fusion, + inference_mode=inference_mode, + input_modalities=input_modalities, + fusion_fields=fusion_fields, + ) + + if bindings is None: + continue - result = fusion_result_from_bindings(fusion, node, bindings) + result = fusion_result_from_bindings(fusion, node, bindings) if is_noop_fusion(node, result): continue @@ -245,10 +257,11 @@ def fusion_enabled_for_fields(fusion: FModels.FusionDefinition, fusion_fields: t return not fusion_specific_fields.isdisjoint(selected_fields - {"generic"}) -def fusion_priority(fusion: FModels.FusionDefinition) -> tuple[int, int, int, int, int]: +def fusion_priority(fusion: FModels.FusionDefinition) -> tuple[int, int, int, int, int, int]: required_attrs = fusion.metadata.get("required_attrs", {}) return ( + int(special_fusions.has_special_matcher(fusion)), len(fusion.graph.nodes), len(fusion.graph.edges), len(fusion.graph.constraints), @@ -271,133 +284,3 @@ def unique_nodes(nodes: Iterable[models.Node]) -> tuple[models.Node, ...]: unique.append(node) return tuple(unique) - - -def fuse_gemma4_decode_attention( - graph: models.Graph, - inference_mode: str | None, - fusion_fields: tuple[str, ...] = (), -) -> models.Graph: - if inference_mode != "decode_with_cache": - return graph - - if fusion_fields and "gemma4_attention" not in fusion_fields and "attention" not in fusion_fields: - return graph - - fusion = Fusions.FUSIONS["attention_masked"] - results: list[models.FusionResult] = [] - - for source in graph.nodes: - match_result = gemma4_decode_attention_match(graph, source, fusion) - - if match_result is not None: - results.append(match_result) - - if not results: - return graph - - return graph.apply_fusions(results) - - -def gemma4_decode_attention_match( - graph: models.Graph, - source: models.Node, - fusion: FModels.FusionDefinition, -) -> models.FusionResult | None: - if source.target != "cactus.view" or len(source.parents) != 1: - return None - - value_bmm = source.parents[0] - - if value_bmm.target != "aten.bmm.default" or len(value_bmm.parents) < 2: - return None - - value_cat = cache_cat_ancestor(value_bmm.parents[1], FModels.CacheTensorRole.VALUE) - - if value_cat is None: - return None - - qk_bmm = first_bmm_ancestor(value_bmm.parents[0], value_bmm) - - if qk_bmm is None or len(qk_bmm.parents) < 2: - return None - - key_cat = cache_cat_ancestor(qk_bmm.parents[1], FModels.CacheTensorRole.KEY) - query = first_rank4_single_parent_ancestor(qk_bmm.parents[0]) - - if key_cat is None or query is None: - return None - - return models.FusionResult.from_match( - fusion=fusion, - source=source, - matched_nodes=(source,), - external_inputs=(query, key_cat, value_cat), - attrs={ - "scale": 1.0, - "input_layout": "bhqd_bhds_bhsd", - "output_layout": "bhqd", - "window_size": 0, - }, - ) - - -def cache_cat_ancestor(node: models.Node, role: str, max_depth: int = 18) -> models.Node | None: - current = node - - for _ in range(max_depth): - if current.target in {"cactus.cat", "aten.cat.default"} and len(current.parents) >= 2: - cache_parent = current.parents[0] - - if cache_parent.cache is not None and cache_parent.cache.kind == FModels.CacheKind.KV and cache_parent.cache.role == role: - return current - - if len(current.parents) != 1: - return None - - current = current.parents[0] - - return None - - -def first_bmm_ancestor(node: models.Node, excluded: models.Node) -> models.Node | None: - stack = [node] - seen: set[str] = set() - - while stack: - current = stack.pop() - - if current.name in seen: - continue - - seen.add(current.name) - - if current is not excluded and current.target == "aten.bmm.default": - return current - - stack.extend(current.parents) - - return None - - -def first_rank4_single_parent_ancestor(node: models.Node, max_depth: int = 10) -> models.Node | None: - current = node - - for _ in range(max_depth): - if node_rank(current) == 4: - return current - - if len(current.parents) != 1: - return None - - current = current.parents[0] - - return None - - -def node_rank(node: models.Node) -> int: - if not isinstance(node.tensor_output_meta, dict): - return 0 - - shape = node.tensor_output_meta.get("shape") - return len(shape) if isinstance(shape, list) else 0 diff --git a/python/cactus/transpiler/IR/special_fusions.py b/python/cactus/transpiler/IR/special_fusions.py new file mode 100644 index 000000000..2aedeb0c3 --- /dev/null +++ b/python/cactus/transpiler/IR/special_fusions.py @@ -0,0 +1,797 @@ +from collections.abc import Callable + +from . import models +from ..Fusions import models as FModels + + +SpecialFusionMatcher = Callable[ + [ + models.Node, + models.Graph, + FModels.FusionDefinition, + str | None, + tuple[str, ...], + tuple[str, ...], + ], + models.FusionResult | None, +] + + +SPECIAL_MATCHER_KEY = "special_matcher" +GEMMA4_ATTENTION_LAYOUT = "gemma4_attention_layout" +HISTORY_POSITION_OFFSET = (1 << 64) - 1 + + +def has_special_matcher(fusion: FModels.FusionDefinition) -> bool: + return special_matcher_name(fusion) in SPECIAL_MATCHERS + + +def match_special_fusion( + source: models.Node, + graph: models.Graph, + fusion: FModels.FusionDefinition, + *, + inference_mode: str | None = None, + input_modalities: tuple[str, ...] = (), + fusion_fields: tuple[str, ...] = (), +) -> models.FusionResult | None: + matcher = SPECIAL_MATCHERS.get(special_matcher_name(fusion)) + + if matcher is None: + return None + + return matcher(source, graph, fusion, inference_mode, input_modalities, fusion_fields) + + +def special_matcher_name(fusion: FModels.FusionDefinition) -> str | None: + value = fusion.metadata.get(SPECIAL_MATCHER_KEY) + + if value is None: + value = fusion.graph.metadata.get(SPECIAL_MATCHER_KEY) + + return str(value) if value is not None else None + + +def match_gemma4_attention_layout( + source: models.Node, + graph: models.Graph, + fusion: FModels.FusionDefinition, + inference_mode: str | None, + input_modalities: tuple[str, ...], + fusion_fields: tuple[str, ...], +) -> models.FusionResult | None: + if inference_mode not in {"prefill_with_cache", "decode_with_cache"}: + return None + + if fusion_fields and "gemma4_attention" not in fusion_fields and "attention" not in fusion_fields: + return None + + if inference_mode == "prefill_with_cache": + prefill_match = gemma4_prefill_attention_match(graph, source, fusion) + if prefill_match is not None: + return prefill_match + + return gemma4_vision_attention_match(graph, source, fusion) + + return gemma4_decode_attention_match(graph, source, fusion) + + +def gemma4_prefill_attention_match( + graph: models.Graph, + source: models.Node, + fusion: FModels.FusionDefinition, +) -> models.FusionResult | None: + value_bmm, layout_nodes = gemma4_prefill_value_bmm_from_output(source) + + if value_bmm is None or len(value_bmm.parents) < 2: + return None + + softmax = first_target_ancestor(value_bmm.parents[0], {"cactus.softmax", "aten.softmax.int", "aten._softmax.default"}) + + if softmax is None or len(softmax.parents) != 1: + return None + + qk_bmm = gemma4_prefill_qk_bmm_from_softmax(softmax, value_bmm) + + if qk_bmm is None or len(qk_bmm.parents) < 2: + return None + + query = first_rank4_single_parent_ancestor(qk_bmm.parents[0]) + key_cat = prefill_cache_cat_ancestor(qk_bmm.parents[1]) + value_cat = prefill_cache_cat_ancestor(value_bmm.parents[1]) + + if query is None or key_cat is None or value_cat is None: + return None + + if not gemma4_prefill_attention_shapes_match(query, key_cat, value_cat, qk_bmm, value_bmm): + return None + + external_inputs = (query, key_cat, value_cat) + + return models.FusionResult.from_match( + fusion=fusion, + source=source, + matched_nodes=layout_nodes, + external_inputs=external_inputs, + attrs={ + "scale": 1.0, + "is_causal": True, + "input_layout": "bhqd_bhds_bhsd", + "q_layout": "bhqd", + "k_layout": "bhsd", + "v_layout": "bhsd", + "output_layout": gemma4_prefill_output_layout(source), + "position_offset": HISTORY_POSITION_OFFSET, + "window_size": gemma4_layer_window_size(source), + "dropped_mask_builder": True, + }, + ) + + +def gemma4_prefill_value_bmm_from_output(source: models.Node, max_depth: int = 10) -> tuple[models.Node | None, tuple[models.Node, ...]]: + current = source + matched_nodes: list[models.Node] = [source] + + for _ in range(max_depth): + if current.target == "aten.bmm.default": + return current, tuple(matched_nodes) + + if len(current.parents) != 1: + return None, () + + parent = current.parents[0] + + if parent.target == "aten.bmm.default": + matched_nodes.append(parent) + return parent, tuple(matched_nodes) + + if parent.target not in { + "cactus.view", + "aten.view.default", + "aten.reshape.default", + "cactus.transpose", + "aten.permute.default", + "aten.transpose.int", + "cactus.precision_cast", + "aten._to_copy.default", + "aten.clone.default", + "aten.contiguous.default", + }: + return None, () + + matched_nodes.append(parent) + current = parent + + return None, () + + +def gemma4_prefill_qk_bmm_from_softmax(softmax: models.Node, value_bmm: models.Node) -> models.Node | None: + logits = softmax.parents[0] + + if logits.target in {"cactus.add", "aten.add.Tensor"} and logits.parents: + return first_bmm_ancestor(logits.parents[0], value_bmm) + + return first_bmm_ancestor(logits, value_bmm) + + +def prefill_cache_cat_ancestor(node: models.Node, max_depth: int = 24) -> models.Node | None: + current = node + + for _ in range(max_depth): + if current.target in {"cactus.cat", "aten.cat.default"} and len(current.parents) >= 2: + first_parent_shape = models.tensor_shape(current.parents[0]) + + if current.parents[0].cache is not None: + return current + + if element_count(first_parent_shape) == 0: + return current + + if len(current.parents) != 1: + return None + + current = current.parents[0] + + return None + + +def gemma4_prefill_attention_shapes_match( + query: models.Node, + key: models.Node, + value: models.Node, + qk_bmm: models.Node, + value_bmm: models.Node, +) -> bool: + query_shape = models.tensor_shape(query) + key_shape = models.tensor_shape(key) + value_shape = models.tensor_shape(value) + qk_shape = models.tensor_shape(qk_bmm) + value_shape_out = models.tensor_shape(value_bmm) + + if len(query_shape) != 4 or len(key_shape) != 4 or len(value_shape) != 4: + return False + + if len(qk_shape) != 3 or len(value_shape_out) != 3: + return False + + return ( + query_shape[0] == key_shape[0] == value_shape[0] + and query_shape[1] == qk_shape[0] + and qk_shape[0] == value_shape_out[0] + and query_shape[2] == qk_shape[1] + and key_shape[2] == qk_shape[2] + and value_shape[2] == qk_shape[2] + and query_shape[3] == key_shape[3] + ) + + +def gemma4_prefill_output_layout(source: models.Node) -> str: + rank = node_rank(source) + + if rank in {2, 3}: + return "bthd_flat" + + return "bhqd" + + +def gemma4_layer_window_size(source: models.Node) -> int: + layer_index = models.layer_index_from_text(f"{source.name} {source.target} {source.module_stack!r}") + + if layer_index is not None and layer_index % 5 == 4: + return 0 + + return 512 + + +def gemma4_vision_attention_match( + graph: models.Graph, + source: models.Node, + fusion: FModels.FusionDefinition, +) -> models.FusionResult | None: + value_bmm, output_layout_nodes = gemma4_vision_value_bmm_from_output(source) + + if value_bmm is None or len(value_bmm.parents) < 2: + return None + + softmax = first_target_ancestor(value_bmm.parents[0], {"cactus.softmax", "aten.softmax.int", "aten._softmax.default"}) + + if softmax is None: + return None + + qk_bmm, mask = gemma4_vision_qk_and_mask_from_softmax(softmax, value_bmm) + + if qk_bmm is None or len(qk_bmm.parents) < 2: + return None + + query = first_rank4_single_parent_ancestor(qk_bmm.parents[0]) + key = first_rank4_single_parent_ancestor(qk_bmm.parents[1]) + value = first_rank4_single_parent_ancestor(value_bmm.parents[1]) + + if query is None or key is None or value is None: + return None + + if not gemma4_vision_attention_shapes_match(source, query, key, value): + return None + + external_inputs = (query, key, value) if mask is None else (query, key, value, mask) + matched_nodes = attention_internal_nodes(source, external_inputs) + + if value_bmm.name not in {node.name for node in matched_nodes}: + matched_nodes = tuple(sorted((*matched_nodes, *output_layout_nodes, value_bmm), key=lambda node: (node.index, node.name))) + + return models.FusionResult.from_match( + fusion=fusion, + source=source, + matched_nodes=matched_nodes, + external_inputs=external_inputs, + attrs={ + "scale": 1.0, + "is_causal": False, + "input_layout": "bhqd_bhds_bhsd", + "output_layout": "bthd_flat", + "window_size": 0, + "additive_mask": mask is not None, + "dropped_softmax_nan_guard": True, + }, + ) + + +def gemma4_vision_value_bmm_from_output(source: models.Node, max_depth: int = 8) -> tuple[models.Node | None, tuple[models.Node, ...]]: + if source.target not in {"cactus.view", "aten.view.default", "aten.reshape.default"} or node_rank(source) != 3: + return None, () + + current = source + matched_nodes: list[models.Node] = [source] + + for _ in range(max_depth): + if len(current.parents) != 1: + return None, () + + parent = current.parents[0] + + if parent.target == "aten.bmm.default": + matched_nodes.append(parent) + return parent, tuple(matched_nodes) + + if parent.target not in { + "cactus.view", + "aten.view.default", + "aten.reshape.default", + "cactus.transpose", + "aten.permute.default", + "aten.transpose.int", + "cactus.precision_cast", + "aten._to_copy.default", + }: + return None, () + + matched_nodes.append(parent) + current = parent + + return None, () + + +def gemma4_vision_qk_and_mask_from_softmax( + softmax: models.Node, + value_bmm: models.Node, +) -> tuple[models.Node | None, models.Node | None]: + if not softmax.parents: + return None, None + + logits = softmax.parents[0] + + if logits.target in {"cactus.add", "aten.add.Tensor"} and len(logits.parents) >= 2: + qk_parent: models.Node | None = None + mask: models.Node | None = None + + for parent in logits.parents: + candidate = first_bmm_ancestor(parent, value_bmm) + + if candidate is not None: + qk_parent = candidate + else: + mask = parent + + return qk_parent, mask + + return first_bmm_ancestor(logits, value_bmm), None + + +def gemma4_vision_attention_shapes_match( + source: models.Node, + query: models.Node, + key: models.Node, + value: models.Node, +) -> bool: + source_shape = models.tensor_shape(source) + query_shape = models.tensor_shape(query) + key_shape = models.tensor_shape(key) + value_shape = models.tensor_shape(value) + + if len(source_shape) != 3 or len(query_shape) != 4 or len(key_shape) != 4 or len(value_shape) != 4: + return False + + batch, heads, tokens, head_dim = query_shape + + if not all(isinstance(dim, int) for dim in (batch, heads, tokens, head_dim)): + return False + + return ( + key_shape == (batch, heads, head_dim, tokens) + and value_shape[0] == batch + and value_shape[1] == heads + and value_shape[2] == tokens + and value_shape[3] == head_dim + and source_shape[0] == batch + and source_shape[1] == tokens + and source_shape[2] == heads * head_dim + ) + + +def attention_internal_nodes(source: models.Node, external_inputs: tuple[models.Node, ...]) -> tuple[models.Node, ...]: + external_names = {node.name for node in external_inputs} + matched: dict[str, models.Node] = {} + stack = [source] + + while stack: + node = stack.pop() + + if node.name in external_names or node.name in matched: + continue + + matched[node.name] = node + + for parent in node.parents: + if parent.name not in external_names: + stack.append(parent) + + return tuple(sorted(matched.values(), key=lambda node: (node.index, node.name))) + + +def first_target_ancestor(node: models.Node, targets: set[str], max_depth: int = 12) -> models.Node | None: + stack: list[tuple[models.Node, int]] = [(node, 0)] + seen: set[int] = set() + + while stack: + current, depth = stack.pop() + + if id(current) in seen or depth > max_depth: + continue + + seen.add(id(current)) + + if current.target in targets: + return current + + stack.extend((parent, depth + 1) for parent in reversed(current.parents)) + + return None + + +def element_count(shape: tuple[object, ...]) -> int | None: + count = 1 + + for dim in shape: + if not isinstance(dim, int): + return None + + count *= dim + + return count + + +def gemma4_decode_attention_match( + graph: models.Graph, + source: models.Node, + fusion: FModels.FusionDefinition, +) -> models.FusionResult | None: + flat_match = gemma4_decode_attention_flat_output_match(graph, source, fusion) + + if flat_match is not None: + return flat_match + + return gemma4_decode_attention_core_output_match(graph, source, fusion) + + +def gemma4_decode_attention_flat_output_match( + graph: models.Graph, + source: models.Node, + fusion: FModels.FusionDefinition, +) -> models.FusionResult | None: + if source.target not in {"cactus.view", "aten.view.default", "aten.reshape.default"}: + return None + + if node_rank(source) != 3: + return None + + existing_attention, existing_layout_nodes = existing_attention_layout_chain(source) + + if existing_attention is not None: + attrs = dict(existing_attention.attrs) + attrs["output_layout"] = "bthd_flat" + external_inputs = tuple(existing_attention.parents) + + if len(existing_attention.parents) >= 3: + key_cat = cache_cat_ancestor(existing_attention.parents[1], FModels.CacheTensorRole.KEY) + value_cat = cache_cat_ancestor(existing_attention.parents[2], FModels.CacheTensorRole.VALUE) + + if key_cat is not None and value_cat is not None: + query, query_layout = gemma4_attention_query_input(existing_attention.parents[0]) + external_inputs = (query, key_cat, value_cat) + attrs["input_layout"] = query_layout + attrs["window_size"] = gemma4_attention_window_size(key_cat) + else: + query, query_layout = gemma4_attention_query_input(existing_attention.parents[0]) + key = gemma4_attention_passthrough_input(existing_attention.parents[1]) + value = gemma4_attention_passthrough_input(existing_attention.parents[2]) + external_inputs = (query, key, value, *existing_attention.parents[3:]) + attrs["input_layout"] = query_layout + + return models.FusionResult.from_match( + fusion=fusion, + source=source, + matched_nodes=(*existing_layout_nodes, existing_attention), + external_inputs=external_inputs, + attrs=attrs, + ) + + inner_view, layout_nodes = post_attention_layout_chain(source) + + if inner_view is None: + return None + + result = gemma4_decode_attention_core_output_match(graph, inner_view, fusion) + + if result is None: + return None + + attrs = dict(result.attrs) + attrs["output_layout"] = "bthd_flat" + attrs["window_size"] = gemma4_attention_window_size(result.external_inputs[1]) + + return models.FusionResult.from_match( + fusion=fusion, + source=source, + matched_nodes=(*layout_nodes, *result.matched_nodes), + external_inputs=result.external_inputs, + attrs=attrs, + ) + + +def existing_attention_layout_chain(source: models.Node, max_depth: int = 8) -> tuple[models.Node | None, tuple[models.Node, ...]]: + current = source + matched_nodes: list[models.Node] = [source] + saw_transpose = False + + for _ in range(max_depth): + if len(current.parents) != 1: + return None, () + + parent = current.parents[0] + + if parent.target in {"cactus.transpose", "aten.permute.default", "aten.transpose.int"}: + saw_transpose = True + matched_nodes.append(parent) + current = parent + continue + + if parent.target == "aten._to_copy.default": + matched_nodes.append(parent) + current = parent + continue + + if parent.target == "cactus.attention" and saw_transpose: + return parent, tuple(matched_nodes) + + return None, () + + return None, () + + +def post_attention_layout_chain(source: models.Node, max_depth: int = 8) -> tuple[models.Node | None, tuple[models.Node, ...]]: + current = source + matched_nodes: list[models.Node] = [source] + saw_transpose = False + + for _ in range(max_depth): + if len(current.parents) != 1: + return None, () + + parent = current.parents[0] + + if parent.target in {"cactus.transpose", "aten.permute.default", "aten.transpose.int"}: + saw_transpose = True + matched_nodes.append(parent) + current = parent + continue + + if parent.target == "aten._to_copy.default": + matched_nodes.append(parent) + current = parent + continue + + if parent.target in {"cactus.view", "aten.view.default", "aten.reshape.default"} and node_rank(parent) == 4: + if saw_transpose and len(parent.parents) == 1 and parent.parents[0].target == "aten.bmm.default": + matched_nodes.append(parent) + return parent, tuple(matched_nodes) + + return None, () + + return None, () + + +def gemma4_decode_attention_core_output_match( + graph: models.Graph, + source: models.Node, + fusion: FModels.FusionDefinition, +) -> models.FusionResult | None: + if source.target != "cactus.view" or len(source.parents) != 1: + return None + + value_bmm = source.parents[0] + + if value_bmm.target != "aten.bmm.default" or len(value_bmm.parents) < 2: + return None + + value_cat = cache_cat_ancestor(value_bmm.parents[1], FModels.CacheTensorRole.VALUE) + + if value_cat is None: + return None + + qk_bmm = first_bmm_ancestor(value_bmm.parents[0], value_bmm) + + if qk_bmm is None or len(qk_bmm.parents) < 2: + return None + + key_cat = cache_cat_ancestor(qk_bmm.parents[1], FModels.CacheTensorRole.KEY) + query = first_rank4_single_parent_ancestor(qk_bmm.parents[0]) + + if key_cat is None or query is None: + return None + + query, query_layout = gemma4_attention_query_input(query) + + return models.FusionResult.from_match( + fusion=fusion, + source=source, + matched_nodes=(source,), + external_inputs=(query, key_cat, value_cat), + attrs={ + "scale": 1.0, + "input_layout": query_layout, + "output_layout": "bhqd", + "window_size": gemma4_attention_window_size(key_cat), + }, + ) + + +def gemma4_attention_window_size(key_cat: models.Node) -> int: + for node in ancestor_nodes(key_cat): + if node.cache is not None and node.cache.window_size is not None: + return int(node.cache.window_size) + + layer_index = models.layer_index_from_text(f"{node.name} {node.target} {node.module_stack!r}") + + if layer_index is not None: + return 0 if layer_index % 5 == 4 else 512 + + shape = models.tensor_shape(node) + sequence_length = shape[2] if len(shape) >= 3 else None + + if sequence_length == 511: + return 512 + + return 0 + + +def ancestor_nodes(node: models.Node, max_depth: int = 16) -> tuple[models.Node, ...]: + found: list[models.Node] = [] + worklist: list[tuple[models.Node, int]] = [(node, 0)] + seen: set[str] = set() + + while worklist: + current, depth = worklist.pop(0) + + if current.name in seen: + continue + + seen.add(current.name) + found.append(current) + + if depth >= max_depth: + continue + + for parent in current.parents: + worklist.append((parent, depth + 1)) + + return tuple(found) + + +def gemma4_attention_query_input(node: models.Node) -> tuple[models.Node, str]: + current = node + layout = "bhqd_bhds_bhsd" + + for _ in range(6): + if len(current.parents) != 1: + return current, layout + + if current.target == "cactus.scalar_multiply" and float(current.attrs.get("value", 1.0)) == 1.0: + current = current.parents[0] + continue + + if current.target == "aten._to_copy.default" and str(current.attrs.get("dtype")) == "torch.float32": + current = current.parents[0] + continue + + if is_bqhd_to_bhqd_transpose(current): + current = current.parents[0] + layout = "bqhd_bhds_bhsd" + continue + + return current, layout + + return current, layout + + +def gemma4_attention_passthrough_input(node: models.Node) -> models.Node: + current = node + + for _ in range(4): + if len(current.parents) != 1: + return current + + if current.target == "cactus.scalar_multiply" and float(current.attrs.get("value", 1.0)) == 1.0: + current = current.parents[0] + continue + + if current.target == "aten._to_copy.default" and str(current.attrs.get("dtype")) == "torch.float32": + current = current.parents[0] + continue + + return current + + return current + + +def is_bqhd_to_bhqd_transpose(node: models.Node) -> bool: + if node.target not in {"cactus.transpose", "aten.permute.default"} or len(node.parents) != 1: + return False + + permutation = node.attrs.get("permutation") + + if permutation != [0, 2, 1, 3]: + return False + + parent_shape = models.tensor_shape(node.parents[0]) + node_shape = models.tensor_shape(node) + + if len(parent_shape) != 4 or len(node_shape) != 4: + return False + + return parent_shape[0] == node_shape[0] and parent_shape[1] == node_shape[2] and parent_shape[2] == node_shape[1] and parent_shape[3] == node_shape[3] + + +def cache_cat_ancestor(node: models.Node, role: str, max_depth: int = 18) -> models.Node | None: + current = node + + for _ in range(max_depth): + if current.target in {"cactus.cat", "aten.cat.default"} and len(current.parents) >= 2: + cache_parent = current.parents[0] + + if cache_parent.cache is not None and cache_parent.cache.kind == FModels.CacheKind.KV and cache_parent.cache.role == role: + return current + + if len(current.parents) != 1: + return None + + current = current.parents[0] + + return None + + +def first_bmm_ancestor(node: models.Node, excluded: models.Node) -> models.Node | None: + stack = [node] + seen: set[str] = set() + + while stack: + current = stack.pop() + + if current.name in seen: + continue + + seen.add(current.name) + + if current is not excluded and current.target == "aten.bmm.default": + return current + + stack.extend(current.parents) + + return None + + +def first_rank4_single_parent_ancestor(node: models.Node, max_depth: int = 10) -> models.Node | None: + current = node + + for _ in range(max_depth): + if node_rank(current) == 4: + return current + + if len(current.parents) != 1: + return None + + current = current.parents[0] + + return None + + +def node_rank(node: models.Node) -> int: + if not isinstance(node.tensor_output_meta, dict): + return 0 + + shape = node.tensor_output_meta.get("shape") + return len(shape) if isinstance(shape, list) else 0 + + +SPECIAL_MATCHERS: dict[str, SpecialFusionMatcher] = { + GEMMA4_ATTENTION_LAYOUT: match_gemma4_attention_layout, +} diff --git a/python/cactus/transpiler/ModelProfiles/profiles.py b/python/cactus/transpiler/ModelProfiles/profiles.py index 447a7492a..aea2fd50e 100644 --- a/python/cactus/transpiler/ModelProfiles/profiles.py +++ b/python/cactus/transpiler/ModelProfiles/profiles.py @@ -29,9 +29,7 @@ "audio", "gemma4_token_merge", "gemma4_rmsnorm", - "gemma4_rope", "gemma4_attention", - "gemma4_mlp", "linear", "cache", ), @@ -68,7 +66,7 @@ inference_type=PARAKEET_INFERENCE_PATTERNS, cache_type=(), cache_policy=(), - files=("config.json", "preprocessor_config.json", "tokenizer_config.json"), + files=("config.json", "tokenizer_config.json"), fusion_fields=( "generic", "audio", diff --git a/python/cactus/transpiler/RuntimePlan/models.py b/python/cactus/transpiler/RuntimePlan/models.py index 3e4843dc2..2170074fe 100644 --- a/python/cactus/transpiler/RuntimePlan/models.py +++ b/python/cactus/transpiler/RuntimePlan/models.py @@ -258,15 +258,38 @@ def runtime_plan_from_generator_manifests( ) for name, manifest_path in component_manifest_paths.items() ) + plan_metadata = runtime_plan_metadata_from_components(components) + plan_metadata.update(string_dict(metadata)) return RuntimePlan( family=runtime_family_from_model_profile(model_profile), components=components, routes=runtime_routes_from_model_profile(model_profile), - metadata=string_dict(metadata), + metadata=plan_metadata, ) +def runtime_plan_metadata_from_components(components: tuple[RuntimeComponent, ...]) -> dict[str, str]: + metadata: dict[str, str] = {} + component_by_name = {component.component: component for component in components} + decoder_prefill = component_by_name.get("decoder_prefill_chunk") + + if decoder_prefill is None: + return metadata + + component_metadata = string_dict(decoder_prefill.metadata) + chunk_tokens = component_metadata.get("prefill_chunk_tokens") + + if chunk_tokens is None: + return metadata + + if "lm_encoder_text_chunk" in component_by_name: + metadata["prefill_strategy"] = "chunked" + metadata["prefill_chunk_tokens"] = chunk_tokens + + return metadata + + def runtime_plan_from_generation_result( generation_result: Any, *, diff --git a/python/tests/test_runtime_plan.py b/python/tests/test_runtime_plan.py index c36322d54..e2a7040cd 100644 --- a/python/tests/test_runtime_plan.py +++ b/python/tests/test_runtime_plan.py @@ -81,7 +81,7 @@ def test_writes_engine_manifest_from_generator_manifest(self): runtime_plan = json.loads(runtime_plan_path.read_text(encoding="utf-8")) component = engine_manifest["components"][0] - self.assertEqual(engine_manifest["family"], "gemma4_e2b") + self.assertEqual(engine_manifest["family"], "gemma4") self.assertEqual(component["component"], "decoder_step") self.assertEqual(component["graph"], "decoder_step.cactus") self.assertEqual(component["logical_inputs"], ["input_ids", "position_ids"]) From 4c6733b58e3c35c4b9bd0e7e763194d2bbd4d52c Mon Sep 17 00:00:00 2001 From: Parkirat Date: Fri, 31 Jul 2026 12:09:25 -0400 Subject: [PATCH 71/84] Code reformatting --- python/cactus/transpiler/CONTRIBUTING.md | 10 + .../transpiler/Converter/cache_utils.py | 139 +- .../transpiler/Converter/input_processor.py | 14 +- python/cactus/transpiler/Converter/models.py | 11 +- .../transpiler/Fusions/fusion_attention.py | 126 + .../transpiler/Fusions/fusion_builders.py | 224 ++ .../transpiler/Fusions/fusion_cache_conv.py | 239 ++ .../transpiler/Fusions/fusion_direct.py | 154 + .../cactus/transpiler/Fusions/fusion_moe.py | 124 + .../transpiler/Fusions/fusion_neural.py | 213 ++ .../transpiler/Fusions/fusion_special.py | 103 + python/cactus/transpiler/Fusions/fusions.py | 1247 +------ .../transpiler/Generator/component_split.py | 669 +--- .../Generator/component_split_boundaries.py | 234 ++ .../Generator/component_split_builder.py | 374 ++ .../Generator/component_split_plans.py | 235 ++ .../Generator/component_split_types.py | 66 + .../Generator/component_split_utils.py | 30 + .../cactus/transpiler/Generator/constants.py | 5 + python/cactus/transpiler/Generator/errors.py | 2 + .../cactus/transpiler/Generator/generate.py | 3065 +---------------- .../Generator/lowering_basic_ops.py | 795 +++++ .../transpiler/Generator/lowering_cache.py | 580 ++++ .../Generator/lowering_constant_ops.py | 253 ++ .../transpiler/Generator/lowering_core.py | 315 ++ .../transpiler/Generator/lowering_nn_ops.py | 240 ++ .../Generator/lowering_special_ops.py | 358 ++ .../transpiler/Generator/lowering_utils.py | 534 +++ .../transpiler/Generator/lowering_weights.py | 293 ++ .../cactus/transpiler/Generator/lowerings.py | 10 + python/cactus/transpiler/Generator/models.py | 26 +- .../transpiler/IR/extra_matcher_cache.py | 235 ++ .../transpiler/IR/extra_matcher_common.py | 346 ++ .../transpiler/IR/extra_matcher_generic.py | 208 ++ .../transpiler/IR/extra_matcher_linear.py | 399 +++ .../cactus/transpiler/IR/extra_matcher_moe.py | 434 +++ python/cactus/transpiler/IR/extra_matchers.py | 1681 +-------- python/cactus/transpiler/IR/models.py | 3 + python/cactus/transpiler/IR/simplify_ir.py | 2 - .../transpiler/ModelProfiles/profiles.py | 21 + python/cactus/transpiler/main.py | 23 +- 41 files changed, 7511 insertions(+), 6529 deletions(-) create mode 100644 python/cactus/transpiler/CONTRIBUTING.md create mode 100644 python/cactus/transpiler/Fusions/fusion_attention.py create mode 100644 python/cactus/transpiler/Fusions/fusion_builders.py create mode 100644 python/cactus/transpiler/Fusions/fusion_cache_conv.py create mode 100644 python/cactus/transpiler/Fusions/fusion_direct.py create mode 100644 python/cactus/transpiler/Fusions/fusion_moe.py create mode 100644 python/cactus/transpiler/Fusions/fusion_neural.py create mode 100644 python/cactus/transpiler/Fusions/fusion_special.py create mode 100644 python/cactus/transpiler/Generator/component_split_boundaries.py create mode 100644 python/cactus/transpiler/Generator/component_split_builder.py create mode 100644 python/cactus/transpiler/Generator/component_split_plans.py create mode 100644 python/cactus/transpiler/Generator/component_split_types.py create mode 100644 python/cactus/transpiler/Generator/component_split_utils.py create mode 100644 python/cactus/transpiler/Generator/errors.py create mode 100644 python/cactus/transpiler/Generator/lowering_basic_ops.py create mode 100644 python/cactus/transpiler/Generator/lowering_cache.py create mode 100644 python/cactus/transpiler/Generator/lowering_constant_ops.py create mode 100644 python/cactus/transpiler/Generator/lowering_core.py create mode 100644 python/cactus/transpiler/Generator/lowering_nn_ops.py create mode 100644 python/cactus/transpiler/Generator/lowering_special_ops.py create mode 100644 python/cactus/transpiler/Generator/lowering_utils.py create mode 100644 python/cactus/transpiler/Generator/lowering_weights.py create mode 100644 python/cactus/transpiler/Generator/lowerings.py create mode 100644 python/cactus/transpiler/IR/extra_matcher_cache.py create mode 100644 python/cactus/transpiler/IR/extra_matcher_common.py create mode 100644 python/cactus/transpiler/IR/extra_matcher_generic.py create mode 100644 python/cactus/transpiler/IR/extra_matcher_linear.py create mode 100644 python/cactus/transpiler/IR/extra_matcher_moe.py diff --git a/python/cactus/transpiler/CONTRIBUTING.md b/python/cactus/transpiler/CONTRIBUTING.md new file mode 100644 index 000000000..166d631e8 --- /dev/null +++ b/python/cactus/transpiler/CONTRIBUTING.md @@ -0,0 +1,10 @@ +First Triumvirate: + +Crasus (v1.0.0 - Millennium): +- How to update converter!!!!! + +Pompey (v1.0.0 - Millennium): +- How to update IR simplification!!!!! + +Caesar (v1.0.0 - Millennium) +- How to update Generation!!!!! \ No newline at end of file diff --git a/python/cactus/transpiler/Converter/cache_utils.py b/python/cactus/transpiler/Converter/cache_utils.py index 0676cb8f5..e651eb44e 100644 --- a/python/cactus/transpiler/Converter/cache_utils.py +++ b/python/cactus/transpiler/Converter/cache_utils.py @@ -3,14 +3,14 @@ from typing import Any import torch -#TODO: Clean out and simplify code - @dataclass(slots=True) class CacheLayerSpec: index: int layer_type: str key_shape: tuple[int, ...] | None = None value_shape: tuple[int, ...] | None = None + cross_key_shape: tuple[int, ...] | None = None + cross_value_shape: tuple[int, ...] | None = None conv_state_shape: tuple[int, ...] | None = None recurrent_state_shape: tuple[int, ...] | None = None sliding_window: int | None = None @@ -58,6 +58,10 @@ def forward( input_features=None, attention_mask=None, input_features_mask=None, + decoder_input_ids=None, + decoder_attention_mask=None, + decoder_inputs_embeds=None, + decoder_position_ids=None, position_ids=None, image_position_ids=None, video_position_ids=None, @@ -74,6 +78,10 @@ def forward( "input_features": input_features, "attention_mask": attention_mask, "input_features_mask": input_features_mask, + "decoder_input_ids": decoder_input_ids, + "decoder_attention_mask": decoder_attention_mask, + "decoder_inputs_embeds": decoder_inputs_embeds, + "decoder_position_ids": decoder_position_ids, "position_ids": position_ids, "image_position_ids": image_position_ids, "video_position_ids": video_position_ids, @@ -171,7 +179,7 @@ def linear_out_features(module: Any) -> int | None: #How: reads layer_types from config, otherwise derives full/sliding attention from window settings. #Why: cache_spec_from_model uses this to know how many cache tensor pairs to create and which are sliding-window. def cache_layer_types(text_config: Any) -> tuple[str, ...]: - num_hidden_layers = int(getattr(text_config, "num_hidden_layers", 1)) + num_hidden_layers = cache_layer_count(text_config) layer_types = getattr(text_config, "layer_types", None) if layer_types is None: @@ -188,21 +196,74 @@ def cache_layer_types(text_config: Any) -> tuple[str, ...]: return layer_types +def cache_layer_count(text_config: Any) -> int: + for attr in ("num_hidden_layers", "decoder_layers", "num_layers", "n_layer"): + value = getattr(text_config, attr, None) + + if value is not None: + return int(value) + + return 1 + + +def cache_num_attention_heads(text_config: Any) -> int: + for attr in ("num_attention_heads", "decoder_attention_heads", "n_head"): + value = getattr(text_config, attr, None) + + if value is not None: + return int(value) + + return 1 + + +def cache_hidden_size(text_config: Any, num_attention_heads: int) -> int: + for attr in ("hidden_size", "d_model", "n_embd"): + value = getattr(text_config, attr, None) + + if value is not None: + return int(value) + + return num_attention_heads + + +def encoder_decoder_source_length(text_config: Any) -> int | None: + for attr in ("max_source_positions", "encoder_seq_length", "max_encoder_position_embeddings"): + value = getattr(text_config, attr, None) + + if value is not None: + return int(value) + + return None + + +def encoder_decoder_target_length(text_config: Any) -> int | None: + for attr in ("max_target_positions", "decoder_seq_length", "max_position_embeddings", "n_positions"): + value = getattr(text_config, attr, None) + + if value is not None: + return int(value) + + return None + + #How: combines config defaults with actual attention module inspection to create one CacheLayerSpec per cache layer. #Why: CacheSpec.from_model calls this before decode/cache export so past_key_values tensors have correct shapes. def cache_spec_from_model(cls: type[CacheSpec], model: torch.nn.Module, batch_size: int, past_sequence_length: int) -> CacheSpec: config = getattr(model, "config", None) text_config = get_text_config(config) - num_key_value_heads = int(getattr(text_config, "num_key_value_heads", getattr(text_config, "num_attention_heads", 1))) - num_attention_heads = int(getattr(text_config, "num_attention_heads", num_key_value_heads)) - hidden_size = int(getattr(text_config, "hidden_size", num_attention_heads)) + num_attention_heads = cache_num_attention_heads(text_config) + num_key_value_heads = int(getattr(text_config, "num_key_value_heads", num_attention_heads)) + hidden_size = cache_hidden_size(text_config, num_attention_heads) head_dim = int(getattr(text_config, "head_dim", max(1, hidden_size // max(1, num_attention_heads)))) sliding_window = getattr(text_config, "sliding_window", None) or getattr(text_config, "attention_chunk_size", None) + is_encoder_decoder = bool(getattr(config, "is_encoder_decoder", False) or getattr(text_config, "is_encoder_decoder", False)) + cross_sequence_length = encoder_decoder_source_length(text_config) if is_encoder_decoder else None + target_sequence_length = encoder_decoder_target_length(text_config) if is_encoder_decoder else None decoder_layers = model_decoder_layers(model) layers: list[CacheLayerSpec] = [] for index, layer_type in enumerate(cache_layer_types(text_config)): - cache_sequence_length = past_sequence_length + cache_sequence_length = int(target_sequence_length or past_sequence_length) layer_sliding_window = None layer_num_key_value_heads = num_key_value_heads layer_head_dim = head_dim @@ -234,12 +295,19 @@ def cache_spec_from_model(cls: type[CacheSpec], model: torch.nn.Module, batch_si continue cache_shape = (batch_size, layer_num_key_value_heads, cache_sequence_length, layer_head_dim) + cross_cache_shape = ( + (batch_size, layer_num_key_value_heads, int(cross_sequence_length), layer_head_dim) + if cross_sequence_length is not None + else None + ) layers.append( CacheLayerSpec( index=index, layer_type=str(layer_type), key_shape=cache_shape, value_shape=cache_shape, + cross_key_shape=cross_cache_shape, + cross_value_shape=cross_cache_shape, sliding_window=layer_sliding_window, ) ) @@ -266,6 +334,10 @@ def cache_spec_empty_tensors( tensors.append(torch.zeros(layer.key_shape, dtype=dtype, device=device)) if layer.value_shape is not None: tensors.append(torch.zeros(layer.value_shape, dtype=dtype, device=device)) + if layer.cross_key_shape is not None: + tensors.append(torch.zeros(layer.cross_key_shape, dtype=dtype, device=device)) + if layer.cross_value_shape is not None: + tensors.append(torch.zeros(layer.cross_value_shape, dtype=dtype, device=device)) if layer.conv_state_shape is not None: tensors.append(torch.zeros(layer.conv_state_shape, dtype=dtype, device=device)) if layer.recurrent_state_shape is not None: @@ -277,18 +349,21 @@ def cache_spec_empty_tensors( #How: creates DynamicCache and assigns each flat key/value tensor into the matching cache layer. #Why: CacheSpec.to_dynamic_cache uses this so CacheExportWrapper.forward can call the original HF model normally. def cache_spec_to_dynamic_cache(cache_spec: CacheSpec, flat_tensors: tuple[torch.Tensor, ...]): - from transformers.cache_utils import DynamicCache + from transformers.cache_utils import DynamicCache, EncoderDecoderCache expected_tensor_count = cache_spec_tensor_count(cache_spec) if len(flat_tensors) != expected_tensor_count: raise ValueError(f"Expected {expected_tensor_count} cache tensors, got {len(flat_tensors)}") - cache = DynamicCache(config=cache_spec.config) + has_cross_cache = any(layer.cross_key_shape is not None or layer.cross_value_shape is not None for layer in cache_spec.layers) + cache = EncoderDecoderCache(DynamicCache(config=cache_spec.config), DynamicCache(config=cache_spec.config)) if has_cross_cache else DynamicCache(config=cache_spec.config) + self_cache = cache.self_attention_cache if has_cross_cache else cache + cross_cache = cache.cross_attention_cache if has_cross_cache else None flat_index = 0 for layer_spec in cache_spec.layers: - layer = cache.layers[layer_spec.index] + layer = self_cache.layers[layer_spec.index] if layer_spec.key_shape is not None and layer_spec.value_shape is not None: key = flat_tensors[flat_index] @@ -303,6 +378,18 @@ def cache_spec_to_dynamic_cache(cache_spec: CacheSpec, flat_tensors: tuple[torch if hasattr(layer, "cumulative_length"): layer.cumulative_length = cache_spec.past_sequence_length + if cross_cache is not None and layer_spec.cross_key_shape is not None and layer_spec.cross_value_shape is not None: + cross_key = flat_tensors[flat_index] + cross_value = flat_tensors[flat_index + 1] + flat_index += 2 + cross_layer = cross_cache.layers[layer_spec.index] + cross_layer.keys = cross_key + cross_layer.values = cross_value + cross_layer.dtype = cross_key.dtype + cross_layer.device = cross_key.device + cross_layer.is_initialized = True + cache.is_updated[layer_spec.index] = True + if layer_spec.conv_state_shape is not None: conv_state = flat_tensors[flat_index].clone() flat_index += 1 @@ -332,6 +419,8 @@ def cache_spec_tensor_count(cache_spec: CacheSpec) -> int: for layer in cache_spec.layers: count += int(layer.key_shape is not None) count += int(layer.value_shape is not None) + count += int(layer.cross_key_shape is not None) + count += int(layer.cross_value_shape is not None) count += int(layer.conv_state_shape is not None) count += int(layer.recurrent_state_shape is not None) @@ -366,6 +455,36 @@ def flatten_dynamic_cache(cache: Any, cache_spec: CacheSpec | None = None) -> tu return () flat: list[torch.Tensor] = [] + self_cache = getattr(cache, "self_attention_cache", None) + cross_cache = getattr(cache, "cross_attention_cache", None) + + if self_cache is not None and cross_cache is not None and cache_spec is not None: + for layer_spec in cache_spec.layers: + self_layer = self_cache.layers[layer_spec.index] + cross_layer = cross_cache.layers[layer_spec.index] + + if layer_spec.key_shape is not None: + key = getattr(self_layer, "keys", None) + if key is not None: + flat.append(key) + + if layer_spec.value_shape is not None: + value = getattr(self_layer, "values", None) + if value is not None: + flat.append(value) + + if layer_spec.cross_key_shape is not None: + cross_key = getattr(cross_layer, "keys", None) + if cross_key is not None: + flat.append(cross_key) + + if layer_spec.cross_value_shape is not None: + cross_value = getattr(cross_layer, "values", None) + if cross_value is not None: + flat.append(cross_value) + + return tuple(flat) + layers = getattr(cache, "layers", None) if layers is None and isinstance(cache, (tuple, list)): diff --git a/python/cactus/transpiler/Converter/input_processor.py b/python/cactus/transpiler/Converter/input_processor.py index c73d1c1b2..fcaaa281a 100644 --- a/python/cactus/transpiler/Converter/input_processor.py +++ b/python/cactus/transpiler/Converter/input_processor.py @@ -67,12 +67,18 @@ def gemma4_processor(model_id: str, configs: dict[str, dict[str, Any]], model_pr def default_processor(model_id: str, configs: dict[str, dict[str, Any]], model_profile: str): from transformers import AutoProcessor - processor_kwargs: dict[str, Any] = {"trust_remote_code": True, "token": constants.token} + processor_kwargs: dict[str, Any] = {"trust_remote_code": True} + + if constants.token is not None: + processor_kwargs["token"] = constants.token try: - return AutoProcessor.from_pretrained(model_id, **processor_kwargs) - except Exception as e: - raise RuntimeError(f"Unable to load processor for {model_id}") + return AutoProcessor.from_pretrained(model_id, local_files_only=True, **processor_kwargs) + except Exception as local_error: + try: + return AutoProcessor.from_pretrained(model_id, **processor_kwargs) + except Exception as remote_error: + raise RuntimeError(f"Unable to load processor for {model_id}") from remote_error def text_tokenizer_processor(model_id: str, configs: dict[str, dict[str, Any]], model_profile: str): diff --git a/python/cactus/transpiler/Converter/models.py b/python/cactus/transpiler/Converter/models.py index 4ad25f6a3..bf6125392 100644 --- a/python/cactus/transpiler/Converter/models.py +++ b/python/cactus/transpiler/Converter/models.py @@ -1,5 +1,4 @@ import json -import os from dataclasses import dataclass from pathlib import Path from typing import Any @@ -216,8 +215,7 @@ def load_configs(mp: MP_Models.ModelProfile, model_id: str | None) -> dict[str, if not local_path.exists(): try: local_path = Path(hf_hub_download(repo_id=model_id, filename=filename, local_dir=output_dir, token=constants.token)) - except Exception as e: - print(f"Error downloading {filename} from {model_id}: {e}") + except Exception: continue with local_path.open("r", encoding="utf-8") as f: @@ -228,7 +226,6 @@ def load_configs(mp: MP_Models.ModelProfile, model_id: str | None) -> dict[str, def build_input(mp: MP_Models.ModelProfile, input_modalities: tuple[str, ...], input_cls: Any, model_id: str | None = None, inference_mode: str = "prefill_no_cache") -> Any | None: if not all(modality in mp.supported_modalties for modality in input_modalities): - print("Requesting unsupported modalities") return None configs = load_configs(mp, model_id) @@ -468,12 +465,6 @@ def create_model(mp: MP_Models.ModelProfile, input_modalities: tuple[str, ...], if input_.inference_mode == constants.DECODE_WITH_CACHE_MODE and constants.DYNAMIC_CACHE_POLICY in mp.cache_policy: input_ = build_decode_with_cache_input(model=loaded_model, input_=input_, input_cls=Input, cache_spec_cls=CU.CacheSpec, model_dtype_fn=CU.model_dtype, drop_multimodal=constants.DROP_MULTIMODAL_ON_DECODE_POLICY in mp.cache_policy) - - #TODO: Add for models that do support multimodal decoders - # elif input_.inference_mode == constants.DECODE_WITH_CACHE_MODE: - # input_ = build_decode_with_cache_input() - - return Model(name=model_id, model_profile=mp, input=input_, model=loaded_model) diff --git a/python/cactus/transpiler/Fusions/fusion_attention.py b/python/cactus/transpiler/Fusions/fusion_attention.py new file mode 100644 index 000000000..8a33504bc --- /dev/null +++ b/python/cactus/transpiler/Fusions/fusion_attention.py @@ -0,0 +1,126 @@ +from __future__ import annotations + +from . import edges as E +from . import models as M +from .fusion_builders import _graph, _input, _shared_input, _single_node_graph + + +ATTENTION_DIRECT_GRAPH = _graph( + "scaled_dot_product_attention", + "attn_sdp", + ("attn_sdp",), + inputs=(_input("query", "attn_sdp", 0), _input("key", "attn_sdp", 1), _input("value", "attn_sdp", 2), _input("mask", "attn_sdp", 3, optional=True)), + attr_captures=(M.AttrCapture("scale", "attn_sdp", "scale", required=False), M.AttrCapture("is_causal", "attn_sdp", "is_causal", default=True, required=False)), +) + + +ATTENTION_CORE_GRAPH = _graph( + "attention_core", + "attn_value", + ("attn_qk", "attn_scale", "attn_softmax", "attn_value"), + edge_names=("attn_qk_to_scale", "attn_scale_to_softmax", "attn_softmax_to_value"), + inputs=(_input("query", "attn_qk", 0), _input("key", "attn_qk", 1), _input("value", "attn_value", 1)), + attr_captures=(M.AttrCapture("scale", "attn_scale", "other", required=False), M.AttrCapture("is_causal", default=True, required=False), M.AttrCapture("window_size", default=0, required=False)), + constraints={ + "node_attr_equals": {"node": "attn_softmax", "attr": "dim", "value": -1}, + "parent_tensor_dim_equal": {"left_node": "attn_qk", "left_parent_index": 0, "left_dim": -1, "right_node": "attn_qk", "right_parent_index": 1, "right_dim": -1, "allow_missing": True}, + }, +) + + +ATTENTION_MASKED_GRAPH = _graph( + "attention_masked", + "attn_value", + ("attn_qk", "attn_scale", "attn_mask_add", "attn_softmax", "attn_value"), + edge_names=("attn_qk_to_scale", "attn_scale_to_mask_add", "attn_mask_add_to_softmax", "attn_softmax_to_value"), + inputs=(_input("query", "attn_qk", 0), _input("key", "attn_qk", 1), _input("mask", "attn_mask_add", 1), _input("value", "attn_value", 1)), + attr_captures=(M.AttrCapture("scale", "attn_scale", "other", required=False), M.AttrCapture("use_mask", default=True), M.AttrCapture("additive_mask", default=True, required=False)), + constraints={ + "input_broadcastable_to_node": {"role": "mask", "node": "attn_qk", "allow_missing": True}, + }, +) + + +GEMMA4_ATTENTION_LAYOUT_GRAPH = _single_node_graph( + "gemma4_attention_layout", + "view", + ("x",), + attr_captures=( + M.AttrCapture("scale", default=1.0, required=False), + M.AttrCapture("input_layout", default="bhqd_bhds_bhsd", required=False), + M.AttrCapture("output_layout", default="bhqd", required=False), + M.AttrCapture("window_size", default=0, required=False), + ), + metadata={"special_matcher": "gemma4_attention_layout"}, +) + + +LFM_BMM_MASKED_ATTENTION_GRAPH = _graph( + "lfm_bmm_masked_attention", + "lfm_attn_output_view", + ( + "lfm_attn_query_expand", + "lfm_attn_key_expand", + "lfm_attn_value_expand", + "lfm_attn_query_flat", + "lfm_attn_key_flat", + "lfm_attn_value_flat", + "lfm_attn_qk", + "lfm_attn_logits_view", + "lfm_attn_mask_where", + "lfm_attn_mask_zero", + "lfm_attn_mask_fill", + "lfm_attn_mask_add", + "lfm_attn_softmax", + "lfm_attn_null_eq", + "lfm_attn_valid_not", + "lfm_attn_any_valid", + "lfm_attn_invalid_not", + "lfm_attn_zero_like", + "lfm_attn_probs_where", + "lfm_attn_probs_expand", + "lfm_attn_probs_flat", + "lfm_attn_output_bmm", + "lfm_attn_output_view", + ), + edge_names=E.EDGE_GROUPS["lfm_bmm_masked_attention"], + inputs=( + _input("query", "lfm_attn_query_expand", 0), + _input("key", "lfm_attn_key_expand", 0), + _input("value", "lfm_attn_value_expand", 0), + _input("mask_predicate", "lfm_attn_mask_where", 0, metadata={"drop_after_fusion": True}), + ), + attr_captures=( + M.AttrCapture("scale", default=1.0, required=False), + M.AttrCapture("is_causal", default=True, required=False), + M.AttrCapture("position_offset", default=0, required=False), + M.AttrCapture("window_size", default=0, required=False), + M.AttrCapture("input_layout", default="bhqd_bhds_bhsd", required=False), + M.AttrCapture("output_layout", default="bhqd", required=False), + M.AttrCapture("dropped_mask_builder", default=True, required=False), + ), + constraints={ + "node_attr_equals": ( + {"node": "lfm_attn_softmax", "attr": "axis", "value": -1, "allow_missing": True}, + {"node": "lfm_attn_null_eq", "attr": "other", "value": None}, + {"node": "lfm_attn_any_valid", "attr": "arg_1", "value": -1}, + ), + }, +) + + +ROPE_GRAPH = _graph( + "rope", + "rope_add", + ("rope_slice_even", "rope_slice_odd", "rope_neg", "rope_rotate_cat", "rope_cos_mul", "rope_sin_mul", "rope_add"), + edge_names=E.EDGE_GROUPS["rope"], + inputs=(_input("x", "rope_slice_even", 0), _input("cos", "rope_cos_mul", 1), _input("sin", "rope_sin_mul", 1)), + shared_inputs=( + _shared_input("rope_slice_even", 0, "rope_slice_odd", 0), + _shared_input("rope_slice_even", 0, "rope_cos_mul", 0), + ), + attr_captures=(M.AttrCapture("theta", default=None, required=False), M.AttrCapture("position_offset", default=0, required=False)), + constraints={ + "rope_tables_compatible": {"x_role": "x", "cos_role": "cos", "sin_role": "sin"}, + }, +) diff --git a/python/cactus/transpiler/Fusions/fusion_builders.py b/python/cactus/transpiler/Fusions/fusion_builders.py new file mode 100644 index 000000000..7ac376222 --- /dev/null +++ b/python/cactus/transpiler/Fusions/fusion_builders.py @@ -0,0 +1,224 @@ +from __future__ import annotations + +from . import edges as E +from . import models as M +from . import nodes as N + + +def _nodes(*names: str) -> dict[str, M.FusionNode]: + return {name: N.NODES[name] for name in names} + + +def _edges(*names: str) -> tuple[M.FusionEdge, ...]: + return tuple(E.EDGES[name] for name in names) + + +def _input(role: str, node: str, parent_index: int | None = None, **kwargs) -> M.FusionInput: + return M.FusionInput(role=role, source=M.NodeRef(node=node, parent_index=parent_index), **kwargs) + + +def _variadic_input( + role: str, + node: str, + start_parent_index: int, + *, + min_count: int = 1, + max_count: int | None = None, + end_parent_index: int | None = None, +) -> M.FusionInput: + return M.FusionInput( + role=role, + source=M.NodeRef(node=node, parent_index=start_parent_index), + variadic=True, + min_count=min_count, + max_count=max_count, + end_parent_index=end_parent_index, + ) + + +def _shared_input( + left_node: str, + left_parent_index: int, + right_node: str, + right_parent_index: int, +) -> tuple[M.NodeRef, M.NodeRef]: + return ( + M.NodeRef(left_node, left_parent_index), + M.NodeRef(right_node, right_parent_index), + ) + + +def _output(node: str, role: str = "out", output_index: int | None = None) -> M.FusionOutput: + return M.FusionOutput(role=role, node=node, output_index=output_index) + + +def _cache_input( + role: str, + node: str, + parent_index: int, + *, + cache_kind: str = M.CacheKind.KV, + tensor_role: str | None = None, + optional: bool = False, +) -> M.CacheInput: + return M.CacheInput( + role, + M.NodeRef(node, parent_index), + cache_kind=cache_kind, + tensor_role=tensor_role, + optional=optional, + ) + + +def _cache_output( + role: str, + node: str, + *, + cache_kind: str = M.CacheKind.KV, + tensor_role: str | None = None, +) -> M.CacheOutput: + return M.CacheOutput(role, node, cache_kind=cache_kind, tensor_role=tensor_role) + + +def _cache_mutation( + name: str, + *, + cache_kind: str = M.CacheKind.KV, + read_roles: tuple[str, ...] = (), + write_roles: tuple[str, ...] = (), +) -> M.CacheMutation: + return M.CacheMutation( + name, + cache_kind=cache_kind, + read_roles=read_roles, + write_roles=write_roles, + ) + + +def _required_attrs(**attrs) -> dict: + return {"required_attrs": attrs} + + +def _graph( + name: str, + root: str, + node_names: tuple[str, ...], + *, + edge_names: tuple[str, ...] = (), + inputs: tuple[M.FusionInput, ...] = (), + shared_inputs: tuple[tuple[M.NodeRef, M.NodeRef], ...] = (), + outputs: tuple[M.FusionOutput, ...] | None = None, + attr_captures: tuple[M.AttrCapture, ...] = (), + repeated_subgraphs: tuple[M.RepeatedSubgraph, ...] = (), + cache_inputs: tuple[M.CacheInput, ...] = (), + cache_outputs: tuple[M.CacheOutput, ...] = (), + cache_mutations: tuple[M.CacheMutation, ...] = (), + constraints: dict[str, M.ConstraintValue] | None = None, + variants: tuple[str, ...] = (), + metadata: dict | None = None, + allow_root_external_children: bool = True, +) -> M.FusionGraph: + return M.FusionGraph( + name=name, + root=root, + nodes=_nodes(*node_names), + edges=_edges(*edge_names), + inputs=inputs, + shared_inputs=shared_inputs, + outputs=outputs or (_output(root),), + attr_captures=attr_captures, + repeated_subgraphs=repeated_subgraphs, + cache_inputs=cache_inputs, + cache_outputs=cache_outputs, + cache_mutations=cache_mutations, + constraints=constraints or {}, + variants=variants, + metadata=metadata or {}, + allow_root_external_children=allow_root_external_children, + ) + + +def _single_node_graph( + name: str, + root: str, + input_roles: tuple[str, ...] = (), + *, + attr_captures: tuple[M.AttrCapture, ...] = (), + metadata: dict | None = None, +) -> M.FusionGraph: + return _graph( + name, + root, + (root,), + inputs=tuple(_input(role, root, index) for index, role in enumerate(input_roles)), + attr_captures=attr_captures, + metadata=metadata, + ) + + +def _definition( + name: str, + cactus_op: str, + graph: M.FusionGraph, + *, + fusion_fields: tuple[str, ...] = ("generic",), + supported_inference_modes: tuple[str, ...] = (), + supported_modalities: tuple[str, ...] = (), + metadata: dict | None = None, +) -> M.FusionDefinition: + return M.FusionDefinition( + name=name, + target=f"cactus.{cactus_op}", + graph=graph, + fusion_fields=fusion_fields, + supported_inference_modes=supported_inference_modes, + supported_modalities=supported_modalities, + cactus_op=cactus_op, + metadata=metadata or {}, + ) + + +def _index_by_target(fusions: dict[str, M.FusionDefinition]) -> dict[str, M.FusionDefinition]: + return {fusion.target: fusion for fusion in fusions.values()} + + +def _index_by_cactus_op( + fusions: dict[str, M.FusionDefinition], +) -> dict[str, tuple[M.FusionDefinition, ...]]: + index: dict[str, tuple[M.FusionDefinition, ...]] = {} + + for fusion in fusions.values(): + if fusion.cactus_op is None: + continue + + index[fusion.cactus_op] = (*index.get(fusion.cactus_op, ()), fusion) + + return index + + +def _index_by_field( + fusions: dict[str, M.FusionDefinition], +) -> dict[str, tuple[M.FusionDefinition, ...]]: + index: dict[str, tuple[M.FusionDefinition, ...]] = {} + + for fusion in fusions.values(): + for field in fusion.fusion_fields: + index[field] = (*index.get(field, ()), fusion) + + return index + + +def _index_by_root_op( + fusions: dict[str, M.FusionDefinition], +) -> dict[str, tuple[M.FusionDefinition, ...]]: + index: dict[str, tuple[M.FusionDefinition, ...]] = {} + + for fusion in fusions.values(): + root_node = fusion.graph.nodes.get(fusion.graph.root) + if root_node is None: + continue + + for op in root_node.ops: + index[op] = (*index.get(op, ()), fusion) + + return index diff --git a/python/cactus/transpiler/Fusions/fusion_cache_conv.py b/python/cactus/transpiler/Fusions/fusion_cache_conv.py new file mode 100644 index 000000000..d85d007b4 --- /dev/null +++ b/python/cactus/transpiler/Fusions/fusion_cache_conv.py @@ -0,0 +1,239 @@ +from __future__ import annotations + +from . import edges as E +from . import models as M +from .fusion_builders import _cache_input, _cache_mutation, _cache_output, _graph, _input + + +CONV_GRAPH = _graph( + "conv", + "conv", + ("conv",), + inputs=(_input("x", "conv", 0), _input("weight", "conv", 1), _input("bias", "conv", 2, optional=True)), + attr_captures=(M.AttrCapture("stride", "conv", "stride", default=1, required=False), M.AttrCapture("padding", "conv", "padding", default=0, required=False), M.AttrCapture("dilation", "conv", "dilation", default=1, required=False), M.AttrCapture("groups", "conv", "groups", default=1, required=False)), +) + + +LFM_SHORT_CONV_PREFILL_GRAPH = _graph( + "lfm_short_conv_prefill", + "slice", + ("conv", "slice"), + edge_names=E.EDGE_GROUPS["short_conv_prefill"], + inputs=(_input("x", "conv", 0), _input("weight", "conv", 1), _input("bias", "conv", 2, optional=True)), + attr_captures=( + M.AttrCapture("kernel_size", default=3, required=False), + M.AttrCapture("dilation", "conv", "dilation", default=1, required=False), + M.AttrCapture("causal", default=True, required=False), + M.AttrCapture("layout", default="batch_hidden_sequence", required=False), + ), + constraints={ + "short_conv_prefill_structure": {"conv_node": "conv", "slice_node": "slice", "x_role": "x", "weight_role": "weight"}, + }, +) + + +LFM_SHORT_CONV_DECODE_GRAPH = _graph( + "lfm_short_conv_decode", + "sum", + ("multiply", "sum"), + edge_names=E.EDGE_GROUPS["short_conv_decode"], + inputs=( + _input("cache_window", "multiply", 0, allowed_value_kinds=(M.ValueKind.CACHE_OUTPUT, M.ValueKind.CACHE_STATE, M.ValueKind.ACTIVATION)), + _input("weight", "multiply", 1, allowed_value_kinds=(M.ValueKind.PARAMETER, M.ValueKind.BUFFER, M.ValueKind.ACTIVATION)), + ), + attr_captures=( + M.AttrCapture("window_size", default=3, required=False), + M.AttrCapture("layout", default="batch_hidden_sequence", required=False), + ), + constraints={ + "short_conv_decode_structure": {"cache_window_role": "cache_window", "multiply_node": "multiply", "sum_node": "sum", "weight_role": "weight"}, + }, +) + + +CONV_BIAS_GRAPH = _graph( + "conv_bias", + "conv_bias_add", + ("conv", "conv_bias_add"), + edge_names=("conv_to_bias_add",), + inputs=(_input("x", "conv", 0), _input("weight", "conv", 1), _input("bias", "conv_bias_add", 1)), + attr_captures=(M.AttrCapture("stride", "conv", "stride", default=1, required=False), M.AttrCapture("padding", "conv", "padding", default=0, required=False), M.AttrCapture("dilation", "conv", "dilation", default=1, required=False), M.AttrCapture("groups", "conv", "groups", default=1, required=False)), +) + + +KV_CACHE_APPEND_GRAPH = _graph( + "kv_cache_append", + "kv_cache_append", + ("kv_cache_append",), + inputs=( + _input("new_kv", "kv_cache_append", 1), + _input("cache_state", "kv_cache_append", 0, allowed_value_kinds=(M.ValueKind.CACHE_INPUT, M.ValueKind.CACHE_STATE, M.ValueKind.CACHE_OUTPUT)), + ), + cache_inputs=(_cache_input("cache_state", "kv_cache_append", 0),), + cache_outputs=(_cache_output("updated_cache_state", "kv_cache_append"),), + cache_mutations=(_cache_mutation("append_kv", read_roles=("cache_state",), write_roles=("updated_cache_state",)),), + constraints={ + "cache_output_consumers": {"node": "kv_cache_append", "allow_output_child": True}, + }, +) + + +KV_CACHE_INITIAL_APPEND_GRAPH = _graph( + "kv_cache_initial_append", + "kv_cache_append", + ("kv_cache_append",), + inputs=( + _input("new_kv", "kv_cache_append", 1), + _input("empty_cache", "kv_cache_append", 0, metadata={"drop_after_fusion": True}), + ), + cache_outputs=(_cache_output("initialized_cache_state", "kv_cache_append"),), + cache_mutations=(_cache_mutation("initialize_append_kv", write_roles=("initialized_cache_state",)),), + constraints={ + "empty_cache_initializer": {"empty_role": "empty_cache", "new_value_role": "new_kv", "output_node": "kv_cache_append"}, + "cache_output_consumers": {"node": "kv_cache_append", "allow_output_child": True}, + }, +) + + +ATTENTION_CACHED_GRAPH = _graph( + "attention_cached", + "attn_value", + ("attn_qk", "attn_scale", "attn_softmax", "attn_value"), + edge_names=("attn_qk_to_scale", "attn_scale_to_softmax", "attn_softmax_to_value"), + inputs=(_input("query", "attn_qk", 0), _input("key_new", "attn_qk", 1), _input("value_new", "attn_value", 1)), + cache_inputs=( + _cache_input("key_cache_state", "attn_qk", 1, tensor_role=M.CacheTensorRole.KEY, optional=True), + _cache_input("value_cache_state", "attn_value", 1, tensor_role=M.CacheTensorRole.VALUE, optional=True), + ), + attr_captures=(M.AttrCapture("scale", "attn_scale", "other", required=False), M.AttrCapture("position_offset", default=0, required=False), M.AttrCapture("window_size", default=0, required=False)), + constraints={ + "same_layer": {"nodes": ("attn_qk", "attn_value"), "allow_missing": True}, + }, +) + + +KV_CACHE_STATE_GRAPH = _graph( + "kv_cache_state", + "kv_cache_state", + ("kv_cache_state",), + outputs=(M.FusionOutput("cache_state", "kv_cache_state"),), + attr_captures=(M.AttrCapture("max_seq_len", default=None, required=False), M.AttrCapture("num_kv_heads", default=None, required=False), M.AttrCapture("head_dim", default=None, required=False), M.AttrCapture("window_size", default=0, required=False), M.AttrCapture("sink_size", default=0, required=False), M.AttrCapture("num_slots", default=1, required=False)), +) + + +CONV_CACHE_STATE_GRAPH = _graph( + "conv_cache_state", + "conv_cache_state", + ("conv_cache_state",), + outputs=(M.FusionOutput("cache_state", "conv_cache_state"),), + attr_captures=(M.AttrCapture("window_size", default=None, required=False), M.AttrCapture("hidden_dim", default=None, required=False)), +) + + +CONV_CACHE_APPEND_GRAPH = _graph( + "conv_cache_append", + "conv_cache_append", + ("conv_cache_append",), + inputs=( + _input("new_data", "conv_cache_append", 0), + _input("cache_state", "conv_cache_append", 1, allowed_value_kinds=(M.ValueKind.CACHE_INPUT, M.ValueKind.CACHE_STATE, M.ValueKind.CACHE_OUTPUT)), + ), + cache_inputs=(_cache_input("cache_state", "conv_cache_append", 1, cache_kind=M.CacheKind.CONV, tensor_role=M.CacheTensorRole.STATE),), + cache_outputs=(_cache_output("updated_cache_state", "conv_cache_append", cache_kind=M.CacheKind.CONV, tensor_role=M.CacheTensorRole.STATE),), + cache_mutations=(_cache_mutation("append_conv_state", cache_kind=M.CacheKind.CONV, read_roles=("cache_state",), write_roles=("updated_cache_state",)),), +) + + +LFM_CONV_CACHE_ROLL_APPEND_GRAPH = _graph( + "lfm_conv_cache_roll_append", + "lfm_conv_cache_output_copy", + ( + "lfm_conv_cache_clone", + "lfm_conv_cache_roll_arange", + "lfm_conv_cache_roll_add", + "lfm_conv_cache_roll_mod", + "lfm_conv_cache_gather", + "lfm_conv_cache_scatter_slice_dim0", + "lfm_conv_cache_scatter_slice_dim1", + "lfm_conv_cache_value_slice_dim0", + "lfm_conv_cache_value_slice_dim1", + "lfm_conv_cache_value_slice_last", + "lfm_conv_cache_value_copy", + "lfm_conv_cache_scatter_dim2", + "lfm_conv_cache_scatter_dim1", + "lfm_conv_cache_scatter_dim0", + "lfm_conv_cache_output_copy", + ), + edge_names=E.EDGE_GROUPS["lfm_conv_cache_roll_append"], + inputs=( + _input("new_data", "lfm_conv_cache_value_copy", 1), + _input("cache_state", "lfm_conv_cache_clone", 0, allowed_value_kinds=(M.ValueKind.CACHE_INPUT, M.ValueKind.CACHE_STATE, M.ValueKind.CACHE_OUTPUT)), + ), + cache_inputs=( + _cache_input("cache_state", "lfm_conv_cache_clone", 0, cache_kind=M.CacheKind.CONV, tensor_role=M.CacheTensorRole.STATE), + ), + cache_outputs=( + _cache_output("updated_cache_state", "lfm_conv_cache_output_copy", cache_kind=M.CacheKind.CONV, tensor_role=M.CacheTensorRole.STATE), + ), + cache_mutations=( + _cache_mutation("roll_append_conv_state", cache_kind=M.CacheKind.CONV, read_roles=("cache_state",), write_roles=("updated_cache_state",)), + ), + attr_captures=( + M.AttrCapture("window_size", default=3, required=False), + ), + constraints={ + "cache_roll_append_structure": { + "cache_role": "cache_state", + "new_data_role": "new_data", + "clone_node": "lfm_conv_cache_clone", + "gather_node": "lfm_conv_cache_gather", + "arange_node": "lfm_conv_cache_roll_arange", + "add_node": "lfm_conv_cache_roll_add", + "mod_node": "lfm_conv_cache_roll_mod", + "scatter_slice_dim0_node": "lfm_conv_cache_scatter_slice_dim0", + "scatter_slice_dim1_node": "lfm_conv_cache_scatter_slice_dim1", + "value_slice_dim0_node": "lfm_conv_cache_value_slice_dim0", + "value_slice_dim1_node": "lfm_conv_cache_value_slice_dim1", + "value_slice_last_node": "lfm_conv_cache_value_slice_last", + "scatter_dim2_node": "lfm_conv_cache_scatter_dim2", + "scatter_dim1_node": "lfm_conv_cache_scatter_dim1", + "scatter_dim0_node": "lfm_conv_cache_scatter_dim0", + }, + }, +) + + +CONV_CACHE_INITIALIZE_GRAPH = _graph( + "conv_cache_initialize", + "conv_cache_initialize", + ("conv_cache_initialize",), + inputs=( + _input("rows", "conv_cache_initialize", 1), + _input("empty_cache", "conv_cache_initialize", 0, metadata={"drop_after_fusion": True}), + ), + cache_outputs=(_cache_output("initialized_cache_state", "conv_cache_initialize", cache_kind=M.CacheKind.CONV, tensor_role=M.CacheTensorRole.STATE),), + cache_mutations=(_cache_mutation("initialize_conv_state", cache_kind=M.CacheKind.CONV, write_roles=("initialized_cache_state",)),), + constraints={ + "conv_cache_initialize_structure": {"rows_role": "rows", "empty_role": "empty_cache", "output_node": "conv_cache_initialize"}, + }, +) + + +RECURRENT_CACHE_STATE_GRAPH = _graph( + "recurrent_cache_state", + "recurrent_cache_state", + ("recurrent_cache_state",), + outputs=(M.FusionOutput("cache_state", "recurrent_cache_state"),), + attr_captures=(M.AttrCapture("shape", default=None, required=False), M.AttrCapture("precision", default=None, required=False)), +) + + +RECURRENT_CACHE_WRITE_GRAPH = _graph( + "recurrent_cache_write", + "recurrent_cache_write", + ("recurrent_cache_write",), + inputs=(_input("new_value", "recurrent_cache_write", 0), _input("cache_input", "recurrent_cache_write", 1)), + cache_inputs=(_cache_input("cache_input", "recurrent_cache_write", 1, cache_kind=M.CacheKind.RECURRENT, tensor_role=M.CacheTensorRole.STATE),), + cache_outputs=(_cache_output("updated_cache", "recurrent_cache_write", cache_kind=M.CacheKind.RECURRENT, tensor_role=M.CacheTensorRole.STATE),), + cache_mutations=(_cache_mutation("write_recurrent_state", cache_kind=M.CacheKind.RECURRENT, read_roles=("cache_input",), write_roles=("updated_cache",)),), +) diff --git a/python/cactus/transpiler/Fusions/fusion_direct.py b/python/cactus/transpiler/Fusions/fusion_direct.py new file mode 100644 index 000000000..b41a780dd --- /dev/null +++ b/python/cactus/transpiler/Fusions/fusion_direct.py @@ -0,0 +1,154 @@ +from __future__ import annotations + +from typing import Any + +from . import models as M +from .fusion_builders import _graph, _single_node_graph, _variadic_input + + +def _capture(name: str, node: str, source_attr: str, *, default: Any = None, required: bool = False) -> M.AttrCapture: + return M.AttrCapture(name, node, source_attr, default=default, required=required) + + +def _same_name_graphs(specs: dict[str, tuple[str, ...]]) -> dict[str, M.FusionGraph]: + return {name: _single_node_graph(name, name, inputs) for name, inputs in specs.items()} + + +def _scalar_graph(name: str) -> M.FusionGraph: + return _single_node_graph(name, name, ("x",), attr_captures=(_capture("value", name, "other"),)) + + +PLAIN_DIRECT_SPECS = { + "add": ("a", "b"), + "subtract": ("a", "b"), + "multiply": ("a", "b"), + "divide": ("a", "b"), + "abs": ("x",), + "sqrt": ("x",), + "scalar_exp": ("x",), + "scalar_sqrt": ("x",), + "scalar_cos": ("x",), + "scalar_sin": ("x",), + "scalar_log": ("x",), + "where": ("condition", "x", "y"), + "not_equal": ("a", "b"), + "equal": ("a", "b"), + "less": ("a", "b"), + "less_equal": ("a", "b"), + "greater": ("a", "b"), + "greater_equal": ("a", "b"), + "bitwise_and": ("a", "b"), + "bitwise_or": ("a", "b"), + "logical_not": ("x",), + "bitwise_not": ("x",), + "relu": ("x",), + "silu": ("x",), + "gelu": ("x",), + "gelu_erf": ("x",), + "sigmoid": ("x",), + "tanh": ("x",), +} + +SCALAR_OPS = ( + "scalar_add", + "scalar_subtract", + "scalar_multiply", + "scalar_divide", + "scalar_floor_divide", + "scalar_not_equal", + "scalar_equal", + "scalar_less", + "scalar_less_equal", + "scalar_greater", + "scalar_greater_equal", +) + +DIRECT_GRAPHS: dict[str, M.FusionGraph] = { + **_same_name_graphs(PLAIN_DIRECT_SPECS), + **{name: _scalar_graph(name) for name in SCALAR_OPS}, + "pow": _single_node_graph("pow", "pow", ("x",), attr_captures=(_capture("exponent", "pow", "exponent"),)), + **{ + name: _single_node_graph(name, name, ("x",), attr_captures=(_capture("axis", name, "dim", default=-1),)) + for name in ("mean", "sum", "variance", "min", "max", "cumsum", "softmax") + }, + "topk": _single_node_graph("topk", "topk_direct", ("x",), attr_captures=(M.AttrCapture("k", "topk_direct", "k"),)), + **{ + name: _single_node_graph(name, name, ("x",), attr_captures=(_capture("shape", name, "shape"),)) + for name in ("view", "reshape", "expand") + }, + "flatten": _single_node_graph( + "flatten", + "flatten", + ("x",), + attr_captures=( + _capture("start_dim", "flatten", "start_dim", default=0), + _capture("end_dim", "flatten", "end_dim", default=-1), + ), + ), + "transpose": _single_node_graph("transpose", "transpose", ("x",), attr_captures=(_capture("permutation", "transpose", "permutation"),)), + "slice": _single_node_graph( + "slice", + "slice", + ("x",), + attr_captures=tuple(_capture(name, "slice", attr) for name, attr in (("axis", "dim"), ("start", "start"), ("end", "end"), ("step", "step"))), + ), + "unfold": _single_node_graph( + "unfold", + "unfold", + ("x",), + attr_captures=( + _capture("dimension", "unfold", "dim"), + _capture("size", "unfold", "size"), + _capture("step", "unfold", "step", default=1), + ), + ), + "index": _single_node_graph( + "index", + "index", + ("x",), + attr_captures=( + _capture("index_value", "index", "index"), + _capture("axis", "index", "dim", default=0), + ), + ), + "cat": _graph("cat", "cat", ("cat",), inputs=(_variadic_input("values", "cat", 0, min_count=2),), attr_captures=(_capture("axis", "cat", "dim", default=0),)), + "gather": _single_node_graph("gather", "gather", ("tensor", "indices"), attr_captures=(_capture("axis", "gather", "dim", default=0),)), + "precision_cast": _single_node_graph("precision_cast", "precision_cast", ("x",), attr_captures=(_capture("dtype", "precision_cast", "dtype"),)), + "pad": _single_node_graph( + "pad", + "pad", + ("x",), + attr_captures=( + _capture("pad", "pad", "pad"), + _capture("value", "pad", "value", default=0.0), + ), + ), + "clamp": _single_node_graph( + "clamp", + "clamp", + ("x",), + attr_captures=( + _capture("lo", "clamp", "min"), + _capture("hi", "clamp", "max"), + ), + ), + "groupnorm": _single_node_graph( + "groupnorm", + "groupnorm", + ("x", "weight", "bias"), + attr_captures=( + _capture("num_groups", "groupnorm", "num_groups"), + _capture("epsilon", "groupnorm", "eps", default=1e-5), + ), + ), + "batchnorm": _single_node_graph( + "batchnorm", + "batchnorm", + ("x", "weight", "bias", "running_mean", "running_var"), + attr_captures=( + _capture("axis", "batchnorm", "axis", default=1), + _capture("epsilon", "batchnorm", "eps", default=1e-5), + ), + ), + "embedding_from_tensor": _single_node_graph("embedding_from_tensor", "embedding", ("embedding_tensor", "indices")), +} diff --git a/python/cactus/transpiler/Fusions/fusion_moe.py b/python/cactus/transpiler/Fusions/fusion_moe.py new file mode 100644 index 000000000..41d88ba89 --- /dev/null +++ b/python/cactus/transpiler/Fusions/fusion_moe.py @@ -0,0 +1,124 @@ +from __future__ import annotations + +from . import edges as E +from . import models as M +from .fusion_builders import _graph, _input, _shared_input, _variadic_input + + +EXPERT_BRANCH_GRAPH = _graph( + "moe_expert_branch", + "expert_weighted", + ("expert_gate", "expert_up", "expert_activation", "expert_product", "expert_down", "expert_weighted"), + edge_names=("expert_gate_to_activation", "expert_activation_to_product", "expert_up_to_product", "expert_product_to_down", "expert_down_to_weighted"), + inputs=(_input("hidden", "expert_gate", 0), _input("gate_weight", "expert_gate", 1), _input("up_weight", "expert_up", 1), _input("down_weight", "expert_down", 1), _input("routing_weight", "expert_weighted", 1)), + shared_inputs=(_shared_input("expert_gate", 0, "expert_up", 0),), +) + + +MOE_GATED_GRAPH = _graph( + "moe_layer_gated", + "moe_combine", + ("moe_combine",), + inputs=(_variadic_input("expert_outputs", "moe_combine", 0),), + repeated_subgraphs=(M.RepeatedSubgraph("experts", EXPERT_BRANCH_GRAPH, min_count=1),), + attr_captures=(M.AttrCapture("num_experts", default=None, required=False), M.AttrCapture("num_experts_per_tok", "topk", "k", required=False), M.AttrCapture("normalize_routing", default=True, required=False), M.AttrCapture("epsilon", default=1e-6, required=False), M.AttrCapture("routed_scaling_factor", default=1.0, required=False)), + constraints={ + "moe_expert_branch_routing": {"repeated_subgraph": "experts", "combine_node": "moe_combine", "routing_weight_role": "routing_weight", "max_depth": 64}, + "moe_routing_weights_combine": {"repeated_subgraph": "experts", "combine_node": "moe_combine", "weighted_node": "expert_weighted", "routing_weight_role": "routing_weight", "max_depth": 64}, + }, +) + + +LFM_GROUPED_MOE_GRAPH = _graph( + "lfm_grouped_moe", + "moe_grouped_combine", + ( + "moe_hidden_view", + "moe_router_weight_transpose", + "moe_router_logits", + "moe_router_sigmoid", + "moe_router_bias", + "moe_topk", + "moe_topk_indices", + "moe_routing_gather", + "moe_routing_sum", + "moe_routing_eps", + "moe_routing_norm", + "moe_routing_scale", + "moe_token_arange", + "moe_token_unsqueeze", + "moe_token_expand", + "moe_token_clone", + "moe_token_flat", + "moe_routing_flat", + "moe_topk_flat", + "moe_sort", + "moe_sort_indices", + "moe_unsort_empty", + "moe_unsort_arange", + "moe_unsort_index_put", + "moe_sorted_expert_index", + "moe_sorted_expert_copy", + "moe_histc", + "moe_offsets", + "moe_sorted_routing_index", + "moe_token_positions_index", + "moe_sorted_hidden_index", + "moe_gate_up_weight_transpose", + "moe_grouped_gate_up", + "moe_gate_up_split", + "moe_gate_tensor", + "moe_up_tensor", + "moe_gate_copy", + "moe_gate_sigmoid", + "moe_gate_activation", + "moe_gate_activation_copy", + "moe_expert_product", + "moe_down_weight_transpose", + "moe_grouped_down", + "moe_routing_unsqueeze", + "moe_grouped_weighted", + "moe_unsort_index", + "moe_output_view", + "moe_grouped_combine", + ), + edge_names=E.EDGE_GROUPS["lfm_grouped_moe"], + inputs=( + _input("hidden", "moe_hidden_view", 0), + _input("router_weight", "moe_router_weight_transpose", 0, allowed_value_kinds=(M.ValueKind.PARAMETER, M.ValueKind.BUFFER)), + _input("router_bias", "moe_router_bias", 1, allowed_value_kinds=(M.ValueKind.PARAMETER, M.ValueKind.BUFFER)), + _input("gate_up_weight", "moe_gate_up_weight_transpose", 0, allowed_value_kinds=(M.ValueKind.PARAMETER, M.ValueKind.BUFFER)), + _input("down_weight", "moe_down_weight_transpose", 0, allowed_value_kinds=(M.ValueKind.PARAMETER, M.ValueKind.BUFFER)), + ), + attr_captures=( + M.AttrCapture("num_experts", "moe_histc", "bins", required=False), + M.AttrCapture("num_experts_per_tok", "moe_topk", "k", required=True), + M.AttrCapture("epsilon", "moe_routing_eps", "other", default=1e-6, required=False), + M.AttrCapture("routed_scaling_factor", "moe_routing_scale", "other", default=1.0, required=False), + ), + constraints={ + "grouped_moe_structure": { + "topk_node": "moe_topk", + "topk_indices_node": "moe_topk_indices", + "router_bias_role": "router_bias", + "gate_up_weight_role": "gate_up_weight", + "down_weight_role": "down_weight", + "gate_up_grouped_node": "moe_grouped_gate_up", + "down_grouped_node": "moe_grouped_down", + "split_node": "moe_gate_up_split", + "offsets_node": "moe_offsets", + "combine_node": "moe_grouped_combine", + }, + }, +) + + +LFM_GROUPED_MOE_NO_TOKEN_CLONE_GRAPH = _graph( + "lfm_grouped_moe_no_token_clone", + "moe_grouped_combine", + tuple(node_name for node_name in LFM_GROUPED_MOE_GRAPH.nodes if node_name != "moe_token_clone"), + edge_names=E.EDGE_GROUPS["lfm_grouped_moe_no_token_clone"], + inputs=LFM_GROUPED_MOE_GRAPH.inputs, + attr_captures=LFM_GROUPED_MOE_GRAPH.attr_captures, + constraints=LFM_GROUPED_MOE_GRAPH.constraints, +) diff --git a/python/cactus/transpiler/Fusions/fusion_neural.py b/python/cactus/transpiler/Fusions/fusion_neural.py new file mode 100644 index 000000000..b0cff1651 --- /dev/null +++ b/python/cactus/transpiler/Fusions/fusion_neural.py @@ -0,0 +1,213 @@ +from __future__ import annotations + +from . import edges as E +from . import models as M +from .fusion_builders import _graph, _input, _shared_input + + +LINEAR_GRAPH = _graph( + "linear", + "linear_mm", + ("linear_mm",), + inputs=(_input("x", "linear_mm", 0), _input("weight", "linear_mm", 1)), + attr_captures=(M.AttrCapture("pretransposed_rhs", default=False, required=False),), + constraints={ + "linear_weight_layout": {"input_role": "x", "weight_role": "weight", "output_node": "linear_mm"}, + }, +) + + +LINEAR_TRANSPOSED_GRAPH = _graph( + "linear_transposed", + "linear_mm", + ("linear_weight_transpose", "linear_mm"), + edge_names=E.EDGE_GROUPS["linear_transposed"], + inputs=( + _input("x", "linear_mm", 0), + _input("weight", "linear_weight_transpose", 0, allowed_value_kinds=(M.ValueKind.PARAMETER, M.ValueKind.BUFFER, M.ValueKind.LIFTED_CONSTANT)), + ), + attr_captures=(M.AttrCapture("pretransposed_rhs", default=True, required=False),), +) + + +LINEAR_BIAS_GRAPH = _graph( + "linear_bias", + "linear_bias_add", + ("linear_mm", "linear_bias_add"), + edge_names=("linear_mm_to_bias_add",), + inputs=(_input("x", "linear_mm", 0), _input("weight", "linear_mm", 1), _input("bias", "linear_bias_add", 1)), + attr_captures=(M.AttrCapture("pretransposed_rhs", default=False, required=False),), + constraints={ + "linear_weight_layout": {"input_role": "x", "weight_role": "weight", "output_node": "linear_mm"}, + "input_value_kind": {"role": "bias", "allowed_value_kinds": (M.ValueKind.PARAMETER, M.ValueKind.BUFFER)}, + }, +) + + +RMS_NORM_GRAPH = _graph( + "rms_norm", + "rms_weight_mul", + ("rms_square", "rms_mean", "rms_eps_add", "rms_inv", "rms_scale", "rms_weight_mul"), + edge_names=E.EDGE_GROUPS["rms_norm"], + inputs=(_input("x", "rms_square", 0), _input("weight", "rms_weight_mul", 1)), + shared_inputs=(_shared_input("rms_square", 0, "rms_scale", 0),), + attr_captures=(M.AttrCapture("epsilon", "rms_eps_add", "other", default=1e-6, required=False),), + constraints={ + "node_attr_equals": {"node": "rms_mean", "attr": "dim", "value": -1}, + }, +) + + +RMS_NORM_POW_GRAPH = _graph( + "rms_norm_pow", + "rms_weight_mul", + ("rms_square", "rms_mean", "rms_eps_add", "rms_inv_pow", "rms_scale", "rms_weight_mul"), + edge_names=E.EDGE_GROUPS["rms_norm_pow"], + inputs=(_input("x", "rms_square", 0), _input("weight", "rms_weight_mul", 1)), + shared_inputs=(_shared_input("rms_square", 0, "rms_scale", 0),), + attr_captures=(M.AttrCapture("epsilon", "rms_eps_add", "other", default=1e-6, required=False),), + constraints={ + "node_attr_equals": {"node": "rms_mean", "attr": "dim", "value": -1}, + }, +) + + +RMS_NORM_NO_WEIGHT_GRAPH = _graph( + "rms_norm_no_weight", + "rms_scale", + ("rms_square", "rms_mean", "rms_eps_add", "rms_inv", "rms_scale"), + edge_names=E.EDGE_GROUPS["rms_norm_no_weight"], + inputs=(_input("x", "rms_square", 0),), + shared_inputs=(_shared_input("rms_square", 0, "rms_scale", 0),), + attr_captures=(M.AttrCapture("epsilon", "rms_eps_add", "other", default=1e-6, required=False),), +) + + +RMS_NORM_POW_NO_WEIGHT_GRAPH = _graph( + "rms_norm_pow_no_weight", + "rms_scale", + ("rms_square", "rms_mean", "rms_eps_add", "rms_inv_pow", "rms_scale"), + edge_names=E.EDGE_GROUPS["rms_norm_pow_no_weight"], + inputs=(_input("x", "rms_square", 0),), + shared_inputs=(_shared_input("rms_square", 0, "rms_scale", 0),), + attr_captures=(M.AttrCapture("epsilon", "rms_eps_add", "other", default=1e-6, required=False),), +) + + +LAYERNORM_NO_BIAS_GRAPH = _graph( + "layernorm_no_bias", + "ln_weight_mul", + ("ln_mean", "ln_center", "ln_square", "ln_var", "ln_eps_add", "ln_inv", "ln_norm", "ln_weight_mul"), + edge_names=( + "ln_mean_to_center", + "ln_center_to_square", + "ln_square_to_var", + "ln_var_to_eps_add", + "ln_eps_add_to_inv", + "ln_center_to_norm", + "ln_inv_to_norm", + "ln_norm_to_weight_mul", + ), + inputs=(_input("x", "ln_mean", 0), _input("weight", "ln_weight_mul", 1)), + shared_inputs=(_shared_input("ln_mean", 0, "ln_center", 0),), + attr_captures=(M.AttrCapture("epsilon", "ln_eps_add", "other", default=1e-5, required=False),), + constraints={ + "node_attrs_equal": {"left_node": "ln_mean", "left_attr": "dim", "right_node": "ln_var", "right_attr": "dim"}, + }, +) + + +LAYERNORM_GRAPH = _graph( + "layernorm", + "ln_bias_add", + ("ln_mean", "ln_center", "ln_square", "ln_var", "ln_eps_add", "ln_inv", "ln_norm", "ln_weight_mul", "ln_bias_add"), + edge_names=E.EDGE_GROUPS["layernorm"], + inputs=(_input("x", "ln_mean", 0), _input("weight", "ln_weight_mul", 1), _input("bias", "ln_bias_add", 1)), + shared_inputs=(_shared_input("ln_mean", 0, "ln_center", 0),), + attr_captures=(M.AttrCapture("epsilon", "ln_eps_add", "other", default=1e-5, required=False),), + constraints={ + "node_attrs_equal": {"left_node": "ln_mean", "left_attr": "dim", "right_node": "ln_var", "right_attr": "dim"}, + "input_value_kind": {"role": "bias", "allowed_value_kinds": (M.ValueKind.PARAMETER, M.ValueKind.BUFFER)}, + }, +) + + +SILU_DECOMPOSED_GRAPH = _graph( + "silu_decomposed", + "silu_product", + ("silu_sigmoid", "silu_product"), + edge_names=E.EDGE_GROUPS["silu_decomposed"], + inputs=(_input("x", "silu_product", 0),), + shared_inputs=(_shared_input("silu_product", 0, "silu_sigmoid", 0),), +) + + +GLU_DECOMPOSED_GRAPH = _graph( + "glu_decomposed", + "glu_product", + ("glu_left_slice", "glu_right_slice", "glu_sigmoid", "glu_product"), + edge_names=E.EDGE_GROUPS["glu_decomposed"], + inputs=(_input("x", "glu_left_slice", 0),), + shared_inputs=(_shared_input("glu_left_slice", 0, "glu_right_slice", 0),), + attr_captures=(M.AttrCapture("axis", "glu_left_slice", "axis", default=-1, required=False),), + constraints={ + "slice_halves": { + "left_node": "glu_left_slice", + "right_node": "glu_right_slice", + }, + }, +) + + +SWIGLU_MLP_GRAPH = _graph( + "swiglu_mlp", + "down_proj", + ("gate_proj", "up_proj", "mlp_activation", "mlp_product", "down_proj"), + edge_names=E.EDGE_GROUPS["swiglu_mlp"], + inputs=(_input("hidden", "gate_proj", 0), _input("gate_weight", "gate_proj", 1), _input("up_weight", "up_proj", 1), _input("down_weight", "down_proj", 1)), + shared_inputs=(_shared_input("gate_proj", 0, "up_proj", 0),), + attr_captures=(M.AttrCapture("product_scale", default=1.0, required=False),), +) + + +GELU_MLP_GRAPH = _graph( + "gelu_mlp", + "down_proj", + ("up_proj", "gelu", "down_proj"), + inputs=(_input("hidden", "up_proj", 0), _input("up_weight", "up_proj", 1), _input("down_weight", "down_proj", 1)), + edge_names=E.EDGE_GROUPS["gelu_mlp"], +) + + +GEMMA4_GEGLU_MLP_GRAPH = _graph( + "gemma4_geglu_mlp", + "down_proj_view", + ( + "gate_input_view", + "gate_weight_transpose", + "gate_proj", + "gate_proj_view", + "gelu", + "up_input_view", + "up_weight_transpose", + "up_proj", + "up_proj_view", + "mlp_product", + "mlp_product_view", + "down_weight_transpose", + "down_proj", + "down_proj_view", + ), + edge_names=E.EDGE_GROUPS["gemma4_geglu_mlp"], + inputs=( + _input("hidden", "gate_input_view", 0), + _input("gate_weight", "gate_weight_transpose", 0, allowed_value_kinds=(M.ValueKind.PARAMETER, M.ValueKind.BUFFER, M.ValueKind.LIFTED_CONSTANT)), + _input("up_weight", "up_weight_transpose", 0, allowed_value_kinds=(M.ValueKind.PARAMETER, M.ValueKind.BUFFER, M.ValueKind.LIFTED_CONSTANT)), + _input("down_weight", "down_weight_transpose", 0, allowed_value_kinds=(M.ValueKind.PARAMETER, M.ValueKind.BUFFER, M.ValueKind.LIFTED_CONSTANT)), + ), + shared_inputs=(_shared_input("gate_input_view", 0, "up_input_view", 0),), + attr_captures=(M.AttrCapture("product_scale", default=1.0 / 64.0, required=False),), + constraints={ + "node_attr_equals": {"node": "gelu", "attr": "approximate", "value": "tanh"}, + }, +) diff --git a/python/cactus/transpiler/Fusions/fusion_special.py b/python/cactus/transpiler/Fusions/fusion_special.py new file mode 100644 index 000000000..361933567 --- /dev/null +++ b/python/cactus/transpiler/Fusions/fusion_special.py @@ -0,0 +1,103 @@ +from __future__ import annotations + +from . import edges as E +from . import models as M +from .fusion_builders import _cache_input, _graph, _input, _variadic_input + + +LSTM_CELL_GRAPH = _graph( + "lstm_cell", + "lstm_tanh", + ("lstm_gate_mm", "lstm_recurrent_mm", "lstm_gate_add", "lstm_sigmoid", "lstm_tanh"), + edge_names=E.EDGE_GROUPS["lstm_cell"], + inputs=(_input("input", "lstm_gate_mm", 0), _input("h_prev", "lstm_recurrent_mm", 0), _input("weight_ih", "lstm_gate_mm", 1), _input("weight_hh", "lstm_recurrent_mm", 1)), + constraints={ + "lstm_gate_structure": {"gate_node": "lstm_gate_add", "gate_count": 4, "require_explicit_gate_split": True}, + }, +) + + +DELTANET_DECODE_GRAPH = _graph( + "gated_deltanet_decode", + "delta_gate", + ("delta_q", "delta_k", "delta_v", "delta_gate"), + inputs=(_input("query", "delta_q", 0), _input("key", "delta_k", 0), _input("value", "delta_v", 0), _input("gate_log", "delta_gate", 0), _input("beta", "delta_gate", 1), _input("initial_state", "delta_gate", 2)), + attr_captures=(M.AttrCapture("scale", default=None, required=False),), + cache_inputs=(_cache_input("initial_state", "delta_gate", 2, cache_kind=M.CacheKind.RECURRENT, tensor_role=M.CacheTensorRole.STATE),), +) + + +DELTANET_PREFILL_GRAPH = _graph( + "gated_deltanet_prefill", + "delta_prefill", + ("delta_q", "delta_k", "delta_v", "delta_gate", "delta_prefill"), + inputs=(_input("query", "delta_q", 0), _input("key", "delta_k", 0), _input("value", "delta_v", 0), _input("gate_log", "delta_gate", 0), _input("beta", "delta_gate", 1), _input("initial_state", "delta_prefill", 2)), + attr_captures=(M.AttrCapture("chunk_size", "delta_prefill", "chunk_size", required=False), M.AttrCapture("scale", default=None, required=False),), + cache_inputs=(_cache_input("initial_state", "delta_prefill", 2, cache_kind=M.CacheKind.RECURRENT, tensor_role=M.CacheTensorRole.STATE),), +) + + +REL_POS_BIAS_GRAPH = _graph( + "rel_pos_bias", + "rel_pos_bias", + ("rel_pos_bias",), + inputs=( + _input("query", "rel_pos_bias", 0, tensor_constraints=(M.TensorConstraint(rank=4),)), + _input("relative_key", "rel_pos_bias", 1, tensor_constraints=(M.TensorConstraint(rank=4),)), + ), + attr_captures=(M.AttrCapture("scale", "rel_pos_bias", "scale", required=False),), +) + + +SAMPLE_GRAPH = _graph( + "sample", + "sample_softmax", + ("sample_topk", "sample_softmax"), + edge_names=("sample_topk_to_softmax",), + inputs=(_input("logits", "sample_topk", 0),), + attr_captures=(M.AttrCapture("temperature", default=0.6, required=False), M.AttrCapture("top_p", default=0.95, required=False), M.AttrCapture("top_k", "sample_topk", "k", default=20, required=False)), +) + + +SCATTER_TOPK_GRAPH = _graph( + "scatter_topk", + "scatter", + ("scatter",), + inputs=(_input("indices", "scatter", 0), _input("values", "scatter", 1)), + attr_captures=(M.AttrCapture("num_classes", "scatter", "num_classes", required=False),), +) + + +GAUSSIAN_TOPK_GRAPH = _graph( + "gaussian_topk", + "gaussian_topk", + ("gaussian_topk",), + inputs=(_input("x", "gaussian_topk", 0),), + attr_captures=(M.AttrCapture("ppf", "gaussian_topk", "ppf", required=False),), +) + + +ALTUP_PREDICT_GRAPH = _graph( + "altup_predict", + "altup_predict", + ("altup_predict",), + inputs=(_input("coefs", "altup_predict", 0), _variadic_input("streams", "altup_predict", 1)), +) + + +ALTUP_CORRECT_GRAPH = _graph( + "altup_correct", + "altup_correct", + ("altup_correct",), + inputs=(_input("coefs", "altup_correct", 0), _input("innovation", "altup_correct", 1), _variadic_input("predictions", "altup_correct", 2)), +) + + +DSP_GRAPHS: dict[str, M.FusionGraph] = { + "stft": _graph("stft", "stft", ("stft",), inputs=(_input("x", "stft", 0), _input("weight", "stft", 1)), attr_captures=(M.AttrCapture("stride", "stft", "stride", required=False), M.AttrCapture("num_fft_bins", "stft", "num_fft_bins", required=False))), + "rfft": _graph("rfft", "rfft", ("rfft",), inputs=(_input("x", "rfft", 0),)), + "irfft": _graph("irfft", "irfft", ("irfft",), inputs=(_input("x", "irfft", 0),), attr_captures=(M.AttrCapture("output_length", "irfft", "output_length", required=False),)), + "spectrogram": _graph("spectrogram", "spectrogram", ("spectrogram",), inputs=(_input("waveform", "spectrogram", 0), _input("mel_filters", "spectrogram", 1)), metadata={"requires_preprocessing_attrs": True}), + "image_preprocess": _graph("image_preprocess", "image_preprocess", ("image_preprocess",), inputs=(_input("pixel_input", "image_preprocess", 0),), metadata={"requires_preprocessing_attrs": True}), + "bilinear_interpolation": _graph("bilinear_interpolation", "bilinear_interpolation", ("bilinear_interpolation",), inputs=(_input("pos_embeds", "bilinear_interpolation", 0),), attr_captures=(M.AttrCapture("dst_height", "bilinear_interpolation", "dst_height", required=False), M.AttrCapture("dst_width", "bilinear_interpolation", "dst_width", required=False))), +} diff --git a/python/cactus/transpiler/Fusions/fusions.py b/python/cactus/transpiler/Fusions/fusions.py index bbdd4497c..b629b4230 100644 --- a/python/cactus/transpiler/Fusions/fusions.py +++ b/python/cactus/transpiler/Fusions/fusions.py @@ -1,1227 +1,20 @@ -from . import edges as E -from . import models as M -from . import nodes as N - - -def _nodes(*names: str) -> dict[str, M.FusionNode]: - return {name: N.NODES[name] for name in names} - - -def _edges(*names: str) -> tuple[M.FusionEdge, ...]: - return tuple(E.EDGES[name] for name in names) - - -def _input(role: str, node: str, parent_index: int | None = None, **kwargs) -> M.FusionInput: - return M.FusionInput(role=role, source=M.NodeRef(node=node, parent_index=parent_index), **kwargs) - - -def _variadic_input( - role: str, - node: str, - start_parent_index: int, - *, - min_count: int = 1, - max_count: int | None = None, - end_parent_index: int | None = None, -) -> M.FusionInput: - return M.FusionInput( - role=role, - source=M.NodeRef(node=node, parent_index=start_parent_index), - variadic=True, - min_count=min_count, - max_count=max_count, - end_parent_index=end_parent_index, - ) - - -def _shared_input( - left_node: str, - left_parent_index: int, - right_node: str, - right_parent_index: int, -) -> tuple[M.NodeRef, M.NodeRef]: - return ( - M.NodeRef(left_node, left_parent_index), - M.NodeRef(right_node, right_parent_index), - ) - - -def _output(node: str, role: str = "out", output_index: int | None = None) -> M.FusionOutput: - return M.FusionOutput(role=role, node=node, output_index=output_index) - - -def _cache_input( - role: str, - node: str, - parent_index: int, - *, - cache_kind: str = M.CacheKind.KV, - tensor_role: str | None = None, - optional: bool = False, -) -> M.CacheInput: - return M.CacheInput( - role, - M.NodeRef(node, parent_index), - cache_kind=cache_kind, - tensor_role=tensor_role, - optional=optional, - ) - - -def _cache_output( - role: str, - node: str, - *, - cache_kind: str = M.CacheKind.KV, - tensor_role: str | None = None, -) -> M.CacheOutput: - return M.CacheOutput(role, node, cache_kind=cache_kind, tensor_role=tensor_role) - - -def _cache_mutation( - name: str, - *, - cache_kind: str = M.CacheKind.KV, - read_roles: tuple[str, ...] = (), - write_roles: tuple[str, ...] = (), -) -> M.CacheMutation: - return M.CacheMutation( - name, - cache_kind=cache_kind, - read_roles=read_roles, - write_roles=write_roles, - ) - - -def _required_attrs(**attrs) -> dict: - return {"required_attrs": attrs} - - -def _graph( - name: str, - root: str, - node_names: tuple[str, ...], - *, - edge_names: tuple[str, ...] = (), - inputs: tuple[M.FusionInput, ...] = (), - shared_inputs: tuple[tuple[M.NodeRef, M.NodeRef], ...] = (), - outputs: tuple[M.FusionOutput, ...] | None = None, - attr_captures: tuple[M.AttrCapture, ...] = (), - repeated_subgraphs: tuple[M.RepeatedSubgraph, ...] = (), - cache_inputs: tuple[M.CacheInput, ...] = (), - cache_outputs: tuple[M.CacheOutput, ...] = (), - cache_mutations: tuple[M.CacheMutation, ...] = (), - constraints: dict[str, M.ConstraintValue] | None = None, - variants: tuple[str, ...] = (), - metadata: dict | None = None, - allow_root_external_children: bool = True, -) -> M.FusionGraph: - return M.FusionGraph( - name=name, - root=root, - nodes=_nodes(*node_names), - edges=_edges(*edge_names), - inputs=inputs, - shared_inputs=shared_inputs, - outputs=outputs or (_output(root),), - attr_captures=attr_captures, - repeated_subgraphs=repeated_subgraphs, - cache_inputs=cache_inputs, - cache_outputs=cache_outputs, - cache_mutations=cache_mutations, - constraints=constraints or {}, - variants=variants, - metadata=metadata or {}, - allow_root_external_children=allow_root_external_children, - ) - - -def _single_node_graph( - name: str, - root: str, - input_roles: tuple[str, ...] = (), - *, - attr_captures: tuple[M.AttrCapture, ...] = (), - metadata: dict | None = None, -) -> M.FusionGraph: - return _graph( - name, - root, - (root,), - inputs=tuple(_input(role, root, index) for index, role in enumerate(input_roles)), - attr_captures=attr_captures, - metadata=metadata, - ) - - -def _definition( - name: str, - cactus_op: str, - graph: M.FusionGraph, - *, - fusion_fields: tuple[str, ...] = ("generic",), - supported_inference_modes: tuple[str, ...] = (), - supported_modalities: tuple[str, ...] = (), - metadata: dict | None = None, -) -> M.FusionDefinition: - return M.FusionDefinition( - name=name, - target=f"cactus.{cactus_op}", - graph=graph, - fusion_fields=fusion_fields, - supported_inference_modes=supported_inference_modes, - supported_modalities=supported_modalities, - cactus_op=cactus_op, - metadata=metadata or {}, - ) +from __future__ import annotations - -def _index_by_target(fusions: dict[str, M.FusionDefinition]) -> dict[str, M.FusionDefinition]: - return {fusion.target: fusion for fusion in fusions.values()} - - -def _index_by_cactus_op( - fusions: dict[str, M.FusionDefinition], -) -> dict[str, tuple[M.FusionDefinition, ...]]: - index: dict[str, tuple[M.FusionDefinition, ...]] = {} - - for fusion in fusions.values(): - if fusion.cactus_op is None: - continue - - index[fusion.cactus_op] = (*index.get(fusion.cactus_op, ()), fusion) - - return index - - -def _index_by_field( - fusions: dict[str, M.FusionDefinition], -) -> dict[str, tuple[M.FusionDefinition, ...]]: - index: dict[str, tuple[M.FusionDefinition, ...]] = {} - - for fusion in fusions.values(): - for field in fusion.fusion_fields: - index[field] = (*index.get(field, ()), fusion) - - return index - - -def _index_by_root_op( - fusions: dict[str, M.FusionDefinition], -) -> dict[str, tuple[M.FusionDefinition, ...]]: - index: dict[str, tuple[M.FusionDefinition, ...]] = {} - - for fusion in fusions.values(): - root_node = fusion.graph.nodes.get(fusion.graph.root) - if root_node is None: - continue - - for op in root_node.ops: - index[op] = (*index.get(op, ()), fusion) - - return index - - -DIRECT_GRAPHS: dict[str, M.FusionGraph] = { - "add": _single_node_graph("add", "add", ("a", "b")), - "subtract": _single_node_graph("subtract", "subtract", ("a", "b")), - "multiply": _single_node_graph("multiply", "multiply", ("a", "b")), - "divide": _single_node_graph("divide", "divide", ("a", "b")), - "abs": _single_node_graph("abs", "abs", ("x",)), - "pow": _single_node_graph( - "pow", - "pow", - ("x",), - attr_captures=(M.AttrCapture("exponent", "pow", "exponent", required=False),), - ), - "sqrt": _single_node_graph("sqrt", "sqrt", ("x",)), - "scalar_add": _single_node_graph( - "scalar_add", - "scalar_add", - ("x",), - attr_captures=(M.AttrCapture("value", "scalar_add", "other", required=False),), - ), - "scalar_subtract": _single_node_graph( - "scalar_subtract", - "scalar_subtract", - ("x",), - attr_captures=(M.AttrCapture("value", "scalar_subtract", "other", required=False),), - ), - "scalar_multiply": _single_node_graph( - "scalar_multiply", - "scalar_multiply", - ("x",), - attr_captures=(M.AttrCapture("value", "scalar_multiply", "other", required=False),), - ), - "scalar_divide": _single_node_graph( - "scalar_divide", - "scalar_divide", - ("x",), - attr_captures=(M.AttrCapture("value", "scalar_divide", "other", required=False),), - ), - "scalar_floor_divide": _single_node_graph( - "scalar_floor_divide", - "scalar_floor_divide", - ("x",), - attr_captures=(M.AttrCapture("value", "scalar_floor_divide", "other", required=False),), - ), - "scalar_exp": _single_node_graph("scalar_exp", "scalar_exp", ("x",)), - "scalar_sqrt": _single_node_graph("scalar_sqrt", "scalar_sqrt", ("x",)), - "scalar_cos": _single_node_graph("scalar_cos", "scalar_cos", ("x",)), - "scalar_sin": _single_node_graph("scalar_sin", "scalar_sin", ("x",)), - "scalar_log": _single_node_graph("scalar_log", "scalar_log", ("x",)), - "scalar_not_equal": _single_node_graph( - "scalar_not_equal", - "scalar_not_equal", - ("x",), - attr_captures=(M.AttrCapture("value", "scalar_not_equal", "other", required=False),), - ), - "scalar_equal": _single_node_graph( - "scalar_equal", - "scalar_equal", - ("x",), - attr_captures=(M.AttrCapture("value", "scalar_equal", "other", required=False),), - ), - "scalar_less": _single_node_graph( - "scalar_less", - "scalar_less", - ("x",), - attr_captures=(M.AttrCapture("value", "scalar_less", "other", required=False),), - ), - "scalar_less_equal": _single_node_graph( - "scalar_less_equal", - "scalar_less_equal", - ("x",), - attr_captures=(M.AttrCapture("value", "scalar_less_equal", "other", required=False),), - ), - "scalar_greater": _single_node_graph( - "scalar_greater", - "scalar_greater", - ("x",), - attr_captures=(M.AttrCapture("value", "scalar_greater", "other", required=False),), - ), - "scalar_greater_equal": _single_node_graph( - "scalar_greater_equal", - "scalar_greater_equal", - ("x",), - attr_captures=(M.AttrCapture("value", "scalar_greater_equal", "other", required=False),), - ), - "mean": _single_node_graph( - "mean", - "mean", - ("x",), - attr_captures=(M.AttrCapture("axis", "mean", "dim", default=-1, required=False),), - ), - "sum": _single_node_graph( - "sum", - "sum", - ("x",), - attr_captures=(M.AttrCapture("axis", "sum", "dim", default=-1, required=False),), - ), - "variance": _single_node_graph( - "variance", - "variance", - ("x",), - attr_captures=(M.AttrCapture("axis", "variance", "dim", default=-1, required=False),), - ), - "min": _single_node_graph( - "min", - "min", - ("x",), - attr_captures=(M.AttrCapture("axis", "min", "dim", default=-1, required=False),), - ), - "max": _single_node_graph( - "max", - "max", - ("x",), - attr_captures=(M.AttrCapture("axis", "max", "dim", default=-1, required=False),), - ), - "cumsum": _single_node_graph( - "cumsum", - "cumsum", - ("x",), - attr_captures=(M.AttrCapture("axis", "cumsum", "dim", default=-1, required=False),), - ), - "softmax": _single_node_graph( - "softmax", - "softmax", - ("x",), - attr_captures=(M.AttrCapture("axis", "softmax", "dim", default=-1, required=False),), - ), - "topk": _single_node_graph( - "topk", - "topk_direct", - ("x",), - attr_captures=(M.AttrCapture("k", "topk_direct", "k", required=True),), - ), - "where": _single_node_graph("where", "where", ("condition", "x", "y")), - "not_equal": _single_node_graph("not_equal", "not_equal", ("a", "b")), - "equal": _single_node_graph("equal", "equal", ("a", "b")), - "less": _single_node_graph("less", "less", ("a", "b")), - "less_equal": _single_node_graph("less_equal", "less_equal", ("a", "b")), - "greater": _single_node_graph("greater", "greater", ("a", "b")), - "greater_equal": _single_node_graph("greater_equal", "greater_equal", ("a", "b")), - "bitwise_and": _single_node_graph("bitwise_and", "bitwise_and", ("a", "b")), - "bitwise_or": _single_node_graph("bitwise_or", "bitwise_or", ("a", "b")), - "logical_not": _single_node_graph("logical_not", "logical_not", ("x",)), - "bitwise_not": _single_node_graph("bitwise_not", "bitwise_not", ("x",)), - "view": _single_node_graph( - "view", - "view", - ("x",), - attr_captures=(M.AttrCapture("shape", "view", "shape", required=False),), - ), - "reshape": _single_node_graph( - "reshape", - "reshape", - ("x",), - attr_captures=(M.AttrCapture("shape", "reshape", "shape", required=False),), - ), - "flatten": _single_node_graph( - "flatten", - "flatten", - ("x",), - attr_captures=( - M.AttrCapture("start_dim", "flatten", "start_dim", default=0, required=False), - M.AttrCapture("end_dim", "flatten", "end_dim", default=-1, required=False), - ), - ), - "expand": _single_node_graph( - "expand", - "expand", - ("x",), - attr_captures=(M.AttrCapture("shape", "expand", "shape", required=False),), - ), - "transpose": _single_node_graph( - "transpose", - "transpose", - ("x",), - attr_captures=(M.AttrCapture("permutation", "transpose", "permutation", required=False),), - ), - "slice": _single_node_graph( - "slice", - "slice", - ("x",), - attr_captures=( - M.AttrCapture("axis", "slice", "dim", required=False), - M.AttrCapture("start", "slice", "start", required=False), - M.AttrCapture("end", "slice", "end", required=False), - M.AttrCapture("step", "slice", "step", required=False), - ), - ), - "unfold": _single_node_graph( - "unfold", - "unfold", - ("x",), - attr_captures=( - M.AttrCapture("dimension", "unfold", "dim", required=False), - M.AttrCapture("size", "unfold", "size", required=False), - M.AttrCapture("step", "unfold", "step", default=1, required=False), - ), - ), - "index": _single_node_graph( - "index", - "index", - ("x",), - attr_captures=( - M.AttrCapture("index_value", "index", "index", required=False), - M.AttrCapture("axis", "index", "dim", default=0, required=False), - ), - ), - "cat": _graph("cat", "cat", ("cat",), inputs=(_variadic_input("values", "cat", 0, min_count=2),), attr_captures=(M.AttrCapture("axis", "cat", "dim", default=0, required=False),)), - "gather": _single_node_graph( - "gather", - "gather", - ("tensor", "indices"), - attr_captures=(M.AttrCapture("axis", "gather", "dim", default=0, required=False),), - ), - "precision_cast": _single_node_graph( - "precision_cast", - "precision_cast", - ("x",), - attr_captures=(M.AttrCapture("dtype", "precision_cast", "dtype", required=False),), - ), - "pad": _single_node_graph( - "pad", - "pad", - ("x",), - attr_captures=( - M.AttrCapture("pad", "pad", "pad", required=False), - M.AttrCapture("value", "pad", "value", default=0.0, required=False), - ), - ), - "relu": _single_node_graph("relu", "relu", ("x",)), - "silu": _single_node_graph("silu", "silu", ("x",)), - "gelu": _single_node_graph("gelu", "gelu", ("x",)), - "gelu_erf": _single_node_graph("gelu_erf", "gelu_erf", ("x",)), - "sigmoid": _single_node_graph("sigmoid", "sigmoid", ("x",)), - "tanh": _single_node_graph("tanh", "tanh", ("x",)), - "clamp": _single_node_graph( - "clamp", - "clamp", - ("x",), - attr_captures=( - M.AttrCapture("lo", "clamp", "min", required=False), - M.AttrCapture("hi", "clamp", "max", required=False), - ), - ), - "groupnorm": _single_node_graph( - "groupnorm", - "groupnorm", - ("x", "weight", "bias"), - attr_captures=( - M.AttrCapture("num_groups", "groupnorm", "num_groups", required=False), - M.AttrCapture("epsilon", "groupnorm", "eps", default=1e-5, required=False), - ), - ), - "batchnorm": _single_node_graph( - "batchnorm", - "batchnorm", - ("x", "weight", "bias", "running_mean", "running_var"), - attr_captures=( - M.AttrCapture("axis", "batchnorm", "axis", default=1, required=False), - M.AttrCapture("epsilon", "batchnorm", "eps", default=1e-5, required=False), - ), - ), - "embedding_from_tensor": _single_node_graph( - "embedding_from_tensor", - "embedding", - ("embedding_tensor", "indices"), - ), -} - - -LINEAR_GRAPH = _graph( - "linear", - "linear_mm", - ("linear_mm",), - inputs=(_input("x", "linear_mm", 0), _input("weight", "linear_mm", 1)), - attr_captures=(M.AttrCapture("pretransposed_rhs", default=False, required=False),), - constraints={ - "linear_weight_layout": {"input_role": "x", "weight_role": "weight", "output_node": "linear_mm"}, - }, -) - -LINEAR_TRANSPOSED_GRAPH = _graph( - "linear_transposed", - "linear_mm", - ("linear_weight_transpose", "linear_mm"), - edge_names=E.EDGE_GROUPS["linear_transposed"], - inputs=( - _input("x", "linear_mm", 0), - _input("weight", "linear_weight_transpose", 0, allowed_value_kinds=(M.ValueKind.PARAMETER, M.ValueKind.BUFFER, M.ValueKind.LIFTED_CONSTANT)), - ), - attr_captures=(M.AttrCapture("pretransposed_rhs", default=True, required=False),), -) - -LINEAR_BIAS_GRAPH = _graph( - "linear_bias", - "linear_bias_add", - ("linear_mm", "linear_bias_add"), - edge_names=("linear_mm_to_bias_add",), - inputs=(_input("x", "linear_mm", 0), _input("weight", "linear_mm", 1), _input("bias", "linear_bias_add", 1)), - attr_captures=(M.AttrCapture("pretransposed_rhs", default=False, required=False),), - constraints={ - "linear_weight_layout": {"input_role": "x", "weight_role": "weight", "output_node": "linear_mm"}, - "input_value_kind": {"role": "bias", "allowed_value_kinds": (M.ValueKind.PARAMETER, M.ValueKind.BUFFER)}, - }, -) - -RMS_NORM_GRAPH = _graph( - "rms_norm", - "rms_weight_mul", - ("rms_square", "rms_mean", "rms_eps_add", "rms_inv", "rms_scale", "rms_weight_mul"), - edge_names=E.EDGE_GROUPS["rms_norm"], - inputs=(_input("x", "rms_square", 0), _input("weight", "rms_weight_mul", 1)), - shared_inputs=(_shared_input("rms_square", 0, "rms_scale", 0),), - attr_captures=(M.AttrCapture("epsilon", "rms_eps_add", "other", default=1e-6, required=False),), - constraints={ - "node_attr_equals": {"node": "rms_mean", "attr": "dim", "value": -1}, - }, -) - -RMS_NORM_POW_GRAPH = _graph( - "rms_norm_pow", - "rms_weight_mul", - ("rms_square", "rms_mean", "rms_eps_add", "rms_inv_pow", "rms_scale", "rms_weight_mul"), - edge_names=E.EDGE_GROUPS["rms_norm_pow"], - inputs=(_input("x", "rms_square", 0), _input("weight", "rms_weight_mul", 1)), - shared_inputs=(_shared_input("rms_square", 0, "rms_scale", 0),), - attr_captures=(M.AttrCapture("epsilon", "rms_eps_add", "other", default=1e-6, required=False),), - constraints={ - "node_attr_equals": {"node": "rms_mean", "attr": "dim", "value": -1}, - }, -) - -RMS_NORM_NO_WEIGHT_GRAPH = _graph( - "rms_norm_no_weight", - "rms_scale", - ("rms_square", "rms_mean", "rms_eps_add", "rms_inv", "rms_scale"), - edge_names=E.EDGE_GROUPS["rms_norm_no_weight"], - inputs=(_input("x", "rms_square", 0),), - shared_inputs=(_shared_input("rms_square", 0, "rms_scale", 0),), - attr_captures=(M.AttrCapture("epsilon", "rms_eps_add", "other", default=1e-6, required=False),), -) - -RMS_NORM_POW_NO_WEIGHT_GRAPH = _graph( - "rms_norm_pow_no_weight", - "rms_scale", - ("rms_square", "rms_mean", "rms_eps_add", "rms_inv_pow", "rms_scale"), - edge_names=E.EDGE_GROUPS["rms_norm_pow_no_weight"], - inputs=(_input("x", "rms_square", 0),), - shared_inputs=(_shared_input("rms_square", 0, "rms_scale", 0),), - attr_captures=(M.AttrCapture("epsilon", "rms_eps_add", "other", default=1e-6, required=False),), -) - -LAYERNORM_NO_BIAS_GRAPH = _graph( - "layernorm_no_bias", - "ln_weight_mul", - ("ln_mean", "ln_center", "ln_square", "ln_var", "ln_eps_add", "ln_inv", "ln_norm", "ln_weight_mul"), - edge_names=( - "ln_mean_to_center", - "ln_center_to_square", - "ln_square_to_var", - "ln_var_to_eps_add", - "ln_eps_add_to_inv", - "ln_center_to_norm", - "ln_inv_to_norm", - "ln_norm_to_weight_mul", - ), - inputs=(_input("x", "ln_mean", 0), _input("weight", "ln_weight_mul", 1)), - shared_inputs=(_shared_input("ln_mean", 0, "ln_center", 0),), - attr_captures=(M.AttrCapture("epsilon", "ln_eps_add", "other", default=1e-5, required=False),), - constraints={ - "node_attrs_equal": {"left_node": "ln_mean", "left_attr": "dim", "right_node": "ln_var", "right_attr": "dim"}, - }, -) - -LAYERNORM_GRAPH = _graph( - "layernorm", - "ln_bias_add", - ("ln_mean", "ln_center", "ln_square", "ln_var", "ln_eps_add", "ln_inv", "ln_norm", "ln_weight_mul", "ln_bias_add"), - edge_names=E.EDGE_GROUPS["layernorm"], - inputs=(_input("x", "ln_mean", 0), _input("weight", "ln_weight_mul", 1), _input("bias", "ln_bias_add", 1)), - shared_inputs=(_shared_input("ln_mean", 0, "ln_center", 0),), - attr_captures=(M.AttrCapture("epsilon", "ln_eps_add", "other", default=1e-5, required=False),), - constraints={ - "node_attrs_equal": {"left_node": "ln_mean", "left_attr": "dim", "right_node": "ln_var", "right_attr": "dim"}, - "input_value_kind": {"role": "bias", "allowed_value_kinds": (M.ValueKind.PARAMETER, M.ValueKind.BUFFER)}, - }, -) - -SILU_DECOMPOSED_GRAPH = _graph( - "silu_decomposed", - "silu_product", - ("silu_sigmoid", "silu_product"), - edge_names=E.EDGE_GROUPS["silu_decomposed"], - inputs=(_input("x", "silu_product", 0),), - shared_inputs=(_shared_input("silu_product", 0, "silu_sigmoid", 0),), -) - -GLU_DECOMPOSED_GRAPH = _graph( - "glu_decomposed", - "glu_product", - ("glu_left_slice", "glu_right_slice", "glu_sigmoid", "glu_product"), - edge_names=E.EDGE_GROUPS["glu_decomposed"], - inputs=(_input("x", "glu_left_slice", 0),), - shared_inputs=(_shared_input("glu_left_slice", 0, "glu_right_slice", 0),), - attr_captures=(M.AttrCapture("axis", "glu_left_slice", "axis", default=-1, required=False),), - constraints={ - "slice_halves": { - "left_node": "glu_left_slice", - "right_node": "glu_right_slice", - }, - }, -) - -SWIGLU_MLP_GRAPH = _graph( - "swiglu_mlp", - "down_proj", - ("gate_proj", "up_proj", "mlp_activation", "mlp_product", "down_proj"), - edge_names=E.EDGE_GROUPS["swiglu_mlp"], - inputs=(_input("hidden", "gate_proj", 0), _input("gate_weight", "gate_proj", 1), _input("up_weight", "up_proj", 1), _input("down_weight", "down_proj", 1)), - shared_inputs=(_shared_input("gate_proj", 0, "up_proj", 0),), - attr_captures=(M.AttrCapture("product_scale", default=1.0, required=False),), -) - -GELU_MLP_GRAPH = _graph( - "gelu_mlp", - "down_proj", - ("up_proj", "gelu", "down_proj"), - inputs=(_input("hidden", "up_proj", 0), _input("up_weight", "up_proj", 1), _input("down_weight", "down_proj", 1)), - edge_names=E.EDGE_GROUPS["gelu_mlp"], -) - -GEMMA4_GEGLU_MLP_GRAPH = _graph( - "gemma4_geglu_mlp", - "down_proj_view", - ( - "gate_input_view", - "gate_weight_transpose", - "gate_proj", - "gate_proj_view", - "gelu", - "up_input_view", - "up_weight_transpose", - "up_proj", - "up_proj_view", - "mlp_product", - "mlp_product_view", - "down_weight_transpose", - "down_proj", - "down_proj_view", - ), - edge_names=E.EDGE_GROUPS["gemma4_geglu_mlp"], - inputs=( - _input("hidden", "gate_input_view", 0), - _input("gate_weight", "gate_weight_transpose", 0, allowed_value_kinds=(M.ValueKind.PARAMETER, M.ValueKind.BUFFER, M.ValueKind.LIFTED_CONSTANT)), - _input("up_weight", "up_weight_transpose", 0, allowed_value_kinds=(M.ValueKind.PARAMETER, M.ValueKind.BUFFER, M.ValueKind.LIFTED_CONSTANT)), - _input("down_weight", "down_weight_transpose", 0, allowed_value_kinds=(M.ValueKind.PARAMETER, M.ValueKind.BUFFER, M.ValueKind.LIFTED_CONSTANT)), - ), - shared_inputs=(_shared_input("gate_input_view", 0, "up_input_view", 0),), - attr_captures=(M.AttrCapture("product_scale", default=1.0, required=False),), - constraints={ - "node_attr_equals": {"node": "gelu", "attr": "approximate", "value": "tanh"}, - }, -) - -ATTENTION_DIRECT_GRAPH = _graph( - "scaled_dot_product_attention", - "attn_sdp", - ("attn_sdp",), - inputs=(_input("query", "attn_sdp", 0), _input("key", "attn_sdp", 1), _input("value", "attn_sdp", 2), _input("mask", "attn_sdp", 3, optional=True)), - attr_captures=(M.AttrCapture("scale", "attn_sdp", "scale", required=False), M.AttrCapture("is_causal", "attn_sdp", "is_causal", default=True, required=False)), -) - -ATTENTION_CORE_GRAPH = _graph( - "attention_core", - "attn_value", - ("attn_qk", "attn_scale", "attn_softmax", "attn_value"), - edge_names=("attn_qk_to_scale", "attn_scale_to_softmax", "attn_softmax_to_value"), - inputs=(_input("query", "attn_qk", 0), _input("key", "attn_qk", 1), _input("value", "attn_value", 1)), - attr_captures=(M.AttrCapture("scale", "attn_scale", "other", required=False), M.AttrCapture("is_causal", default=True, required=False), M.AttrCapture("window_size", default=0, required=False)), - constraints={ - "node_attr_equals": {"node": "attn_softmax", "attr": "dim", "value": -1}, - "parent_tensor_dim_equal": {"left_node": "attn_qk", "left_parent_index": 0, "left_dim": -1, "right_node": "attn_qk", "right_parent_index": 1, "right_dim": -1, "allow_missing": True}, - }, -) - -ATTENTION_MASKED_GRAPH = _graph( - "attention_masked", - "attn_value", - ("attn_qk", "attn_scale", "attn_mask_add", "attn_softmax", "attn_value"), - edge_names=("attn_qk_to_scale", "attn_scale_to_mask_add", "attn_mask_add_to_softmax", "attn_softmax_to_value"), - inputs=(_input("query", "attn_qk", 0), _input("key", "attn_qk", 1), _input("mask", "attn_mask_add", 1), _input("value", "attn_value", 1)), - attr_captures=(M.AttrCapture("scale", "attn_scale", "other", required=False), M.AttrCapture("use_mask", default=True), M.AttrCapture("additive_mask", default=True, required=False)), - constraints={ - "input_broadcastable_to_node": {"role": "mask", "node": "attn_qk", "allow_missing": True}, - }, -) - -GEMMA4_ATTENTION_LAYOUT_GRAPH = _single_node_graph( - "gemma4_attention_layout", - "view", - ("x",), - attr_captures=( - M.AttrCapture("scale", default=1.0, required=False), - M.AttrCapture("input_layout", default="bhqd_bhds_bhsd", required=False), - M.AttrCapture("output_layout", default="bhqd", required=False), - M.AttrCapture("window_size", default=0, required=False), - ), - metadata={"special_matcher": "gemma4_attention_layout"}, -) - -LFM_BMM_MASKED_ATTENTION_GRAPH = _graph( - "lfm_bmm_masked_attention", - "lfm_attn_output_view", - ( - "lfm_attn_query_expand", - "lfm_attn_key_expand", - "lfm_attn_value_expand", - "lfm_attn_query_flat", - "lfm_attn_key_flat", - "lfm_attn_value_flat", - "lfm_attn_qk", - "lfm_attn_logits_view", - "lfm_attn_mask_where", - "lfm_attn_mask_zero", - "lfm_attn_mask_fill", - "lfm_attn_mask_add", - "lfm_attn_softmax", - "lfm_attn_null_eq", - "lfm_attn_valid_not", - "lfm_attn_any_valid", - "lfm_attn_invalid_not", - "lfm_attn_zero_like", - "lfm_attn_probs_where", - "lfm_attn_probs_expand", - "lfm_attn_probs_flat", - "lfm_attn_output_bmm", - "lfm_attn_output_view", - ), - edge_names=E.EDGE_GROUPS["lfm_bmm_masked_attention"], - inputs=( - _input("query", "lfm_attn_query_expand", 0), - _input("key", "lfm_attn_key_expand", 0), - _input("value", "lfm_attn_value_expand", 0), - _input("mask_predicate", "lfm_attn_mask_where", 0, metadata={"drop_after_fusion": True}), - ), - attr_captures=( - M.AttrCapture("scale", default=1.0, required=False), - M.AttrCapture("is_causal", default=True, required=False), - M.AttrCapture("position_offset", default=0, required=False), - M.AttrCapture("window_size", default=0, required=False), - M.AttrCapture("input_layout", default="bhqd_bhds_bhsd", required=False), - M.AttrCapture("output_layout", default="bhqd", required=False), - M.AttrCapture("dropped_mask_builder", default=True, required=False), - ), - constraints={ - "node_attr_equals": ( - {"node": "lfm_attn_softmax", "attr": "axis", "value": -1, "allow_missing": True}, - {"node": "lfm_attn_null_eq", "attr": "other", "value": None}, - {"node": "lfm_attn_any_valid", "attr": "arg_1", "value": -1}, - ), - }, -) - -ROPE_GRAPH = _graph( - "rope", - "rope_add", - ("rope_slice_even", "rope_slice_odd", "rope_neg", "rope_rotate_cat", "rope_cos_mul", "rope_sin_mul", "rope_add"), - edge_names=E.EDGE_GROUPS["rope"], - inputs=(_input("x", "rope_slice_even", 0), _input("cos", "rope_cos_mul", 1), _input("sin", "rope_sin_mul", 1)), - shared_inputs=( - _shared_input("rope_slice_even", 0, "rope_slice_odd", 0), - _shared_input("rope_slice_even", 0, "rope_cos_mul", 0), - ), - attr_captures=(M.AttrCapture("theta", default=None, required=False), M.AttrCapture("position_offset", default=0, required=False)), - constraints={ - "rope_tables_compatible": {"x_role": "x", "cos_role": "cos", "sin_role": "sin"}, - }, -) - -CONV_GRAPH = _graph( - "conv", - "conv", - ("conv",), - inputs=(_input("x", "conv", 0), _input("weight", "conv", 1), _input("bias", "conv", 2, optional=True)), - attr_captures=(M.AttrCapture("stride", "conv", "stride", default=1, required=False), M.AttrCapture("padding", "conv", "padding", default=0, required=False), M.AttrCapture("dilation", "conv", "dilation", default=1, required=False), M.AttrCapture("groups", "conv", "groups", default=1, required=False)), -) - -LFM_SHORT_CONV_PREFILL_GRAPH = _graph( - "lfm_short_conv_prefill", - "slice", - ("conv", "slice"), - edge_names=E.EDGE_GROUPS["short_conv_prefill"], - inputs=(_input("x", "conv", 0), _input("weight", "conv", 1), _input("bias", "conv", 2, optional=True)), - attr_captures=( - M.AttrCapture("kernel_size", default=3, required=False), - M.AttrCapture("dilation", "conv", "dilation", default=1, required=False), - M.AttrCapture("causal", default=True, required=False), - M.AttrCapture("layout", default="batch_hidden_sequence", required=False), - ), - constraints={ - "short_conv_prefill_structure": {"conv_node": "conv", "slice_node": "slice", "x_role": "x", "weight_role": "weight"}, - }, -) - -LFM_SHORT_CONV_DECODE_GRAPH = _graph( - "lfm_short_conv_decode", - "sum", - ("multiply", "sum"), - edge_names=E.EDGE_GROUPS["short_conv_decode"], - inputs=( - _input("cache_window", "multiply", 0, allowed_value_kinds=(M.ValueKind.CACHE_OUTPUT, M.ValueKind.CACHE_STATE, M.ValueKind.ACTIVATION)), - _input("weight", "multiply", 1, allowed_value_kinds=(M.ValueKind.PARAMETER, M.ValueKind.BUFFER, M.ValueKind.ACTIVATION)), - ), - attr_captures=( - M.AttrCapture("window_size", default=3, required=False), - M.AttrCapture("layout", default="batch_hidden_sequence", required=False), - ), - constraints={ - "short_conv_decode_structure": {"cache_window_role": "cache_window", "multiply_node": "multiply", "sum_node": "sum", "weight_role": "weight"}, - }, -) - -CONV_BIAS_GRAPH = _graph( - "conv_bias", - "conv_bias_add", - ("conv", "conv_bias_add"), - edge_names=("conv_to_bias_add",), - inputs=(_input("x", "conv", 0), _input("weight", "conv", 1), _input("bias", "conv_bias_add", 1)), - attr_captures=(M.AttrCapture("stride", "conv", "stride", default=1, required=False), M.AttrCapture("padding", "conv", "padding", default=0, required=False), M.AttrCapture("dilation", "conv", "dilation", default=1, required=False), M.AttrCapture("groups", "conv", "groups", default=1, required=False)), -) - -KV_CACHE_APPEND_GRAPH = _graph( - "kv_cache_append", - "kv_cache_append", - ("kv_cache_append",), - inputs=( - _input("new_kv", "kv_cache_append", 1), - _input("cache_state", "kv_cache_append", 0, allowed_value_kinds=(M.ValueKind.CACHE_INPUT, M.ValueKind.CACHE_STATE, M.ValueKind.CACHE_OUTPUT)), - ), - cache_inputs=(_cache_input("cache_state", "kv_cache_append", 0),), - cache_outputs=(_cache_output("updated_cache_state", "kv_cache_append"),), - cache_mutations=(_cache_mutation("append_kv", read_roles=("cache_state",), write_roles=("updated_cache_state",)),), - constraints={ - "cache_output_consumers": {"node": "kv_cache_append", "allow_output_child": True}, - }, -) - -KV_CACHE_INITIAL_APPEND_GRAPH = _graph( - "kv_cache_initial_append", - "kv_cache_append", - ("kv_cache_append",), - inputs=( - _input("new_kv", "kv_cache_append", 1), - _input("empty_cache", "kv_cache_append", 0, metadata={"drop_after_fusion": True}), - ), - cache_outputs=(_cache_output("initialized_cache_state", "kv_cache_append"),), - cache_mutations=(_cache_mutation("initialize_append_kv", write_roles=("initialized_cache_state",)),), - constraints={ - "empty_cache_initializer": {"empty_role": "empty_cache", "new_value_role": "new_kv", "output_node": "kv_cache_append"}, - "cache_output_consumers": {"node": "kv_cache_append", "allow_output_child": True}, - }, -) - -ATTENTION_CACHED_GRAPH = _graph( - "attention_cached", - "attn_value", - ("attn_qk", "attn_scale", "attn_softmax", "attn_value"), - edge_names=("attn_qk_to_scale", "attn_scale_to_softmax", "attn_softmax_to_value"), - inputs=(_input("query", "attn_qk", 0), _input("key_new", "attn_qk", 1), _input("value_new", "attn_value", 1)), - cache_inputs=( - _cache_input("key_cache_state", "attn_qk", 1, tensor_role=M.CacheTensorRole.KEY, optional=True), - _cache_input("value_cache_state", "attn_value", 1, tensor_role=M.CacheTensorRole.VALUE, optional=True), - ), - attr_captures=(M.AttrCapture("scale", "attn_scale", "other", required=False), M.AttrCapture("position_offset", default=0, required=False), M.AttrCapture("window_size", default=0, required=False)), - constraints={ - "same_layer": {"nodes": ("attn_qk", "attn_value"), "allow_missing": True}, - }, -) - -EXPERT_BRANCH_GRAPH = _graph( - "moe_expert_branch", - "expert_weighted", - ("expert_gate", "expert_up", "expert_activation", "expert_product", "expert_down", "expert_weighted"), - edge_names=("expert_gate_to_activation", "expert_activation_to_product", "expert_up_to_product", "expert_product_to_down", "expert_down_to_weighted"), - inputs=(_input("hidden", "expert_gate", 0), _input("gate_weight", "expert_gate", 1), _input("up_weight", "expert_up", 1), _input("down_weight", "expert_down", 1), _input("routing_weight", "expert_weighted", 1)), - shared_inputs=(_shared_input("expert_gate", 0, "expert_up", 0),), -) - -MOE_GATED_GRAPH = _graph( - "moe_layer_gated", - "moe_combine", - ("moe_combine",), - inputs=(_variadic_input("expert_outputs", "moe_combine", 0),), - repeated_subgraphs=(M.RepeatedSubgraph("experts", EXPERT_BRANCH_GRAPH, min_count=1),), - attr_captures=(M.AttrCapture("num_experts", default=None, required=False), M.AttrCapture("num_experts_per_tok", "topk", "k", required=False), M.AttrCapture("normalize_routing", default=True, required=False), M.AttrCapture("epsilon", default=1e-6, required=False), M.AttrCapture("routed_scaling_factor", default=1.0, required=False)), - constraints={ - "moe_expert_branch_routing": {"repeated_subgraph": "experts", "combine_node": "moe_combine", "routing_weight_role": "routing_weight", "max_depth": 64}, - "moe_routing_weights_combine": {"repeated_subgraph": "experts", "combine_node": "moe_combine", "weighted_node": "expert_weighted", "routing_weight_role": "routing_weight", "max_depth": 64}, - }, -) - -LFM_GROUPED_MOE_GRAPH = _graph( - "lfm_grouped_moe", - "moe_grouped_combine", - ( - "moe_hidden_view", - "moe_router_weight_transpose", - "moe_router_logits", - "moe_router_sigmoid", - "moe_router_bias", - "moe_topk", - "moe_topk_indices", - "moe_routing_gather", - "moe_routing_sum", - "moe_routing_eps", - "moe_routing_norm", - "moe_routing_scale", - "moe_token_arange", - "moe_token_unsqueeze", - "moe_token_expand", - "moe_token_clone", - "moe_token_flat", - "moe_routing_flat", - "moe_topk_flat", - "moe_sort", - "moe_sort_indices", - "moe_unsort_empty", - "moe_unsort_arange", - "moe_unsort_index_put", - "moe_sorted_expert_index", - "moe_sorted_expert_copy", - "moe_histc", - "moe_offsets", - "moe_sorted_routing_index", - "moe_token_positions_index", - "moe_sorted_hidden_index", - "moe_gate_up_weight_transpose", - "moe_grouped_gate_up", - "moe_gate_up_split", - "moe_gate_tensor", - "moe_up_tensor", - "moe_gate_copy", - "moe_gate_sigmoid", - "moe_gate_activation", - "moe_gate_activation_copy", - "moe_expert_product", - "moe_down_weight_transpose", - "moe_grouped_down", - "moe_routing_unsqueeze", - "moe_grouped_weighted", - "moe_unsort_index", - "moe_output_view", - "moe_grouped_combine", - ), - edge_names=E.EDGE_GROUPS["lfm_grouped_moe"], - inputs=( - _input("hidden", "moe_hidden_view", 0), - _input("router_weight", "moe_router_weight_transpose", 0, allowed_value_kinds=(M.ValueKind.PARAMETER, M.ValueKind.BUFFER)), - _input("router_bias", "moe_router_bias", 1, allowed_value_kinds=(M.ValueKind.PARAMETER, M.ValueKind.BUFFER)), - _input("gate_up_weight", "moe_gate_up_weight_transpose", 0, allowed_value_kinds=(M.ValueKind.PARAMETER, M.ValueKind.BUFFER)), - _input("down_weight", "moe_down_weight_transpose", 0, allowed_value_kinds=(M.ValueKind.PARAMETER, M.ValueKind.BUFFER)), - ), - attr_captures=( - M.AttrCapture("num_experts", "moe_histc", "bins", required=False), - M.AttrCapture("num_experts_per_tok", "moe_topk", "k", required=True), - M.AttrCapture("epsilon", "moe_routing_eps", "other", default=1e-6, required=False), - M.AttrCapture("routed_scaling_factor", "moe_routing_scale", "other", default=1.0, required=False), - ), - constraints={ - "grouped_moe_structure": { - "topk_node": "moe_topk", - "topk_indices_node": "moe_topk_indices", - "router_bias_role": "router_bias", - "gate_up_weight_role": "gate_up_weight", - "down_weight_role": "down_weight", - "gate_up_grouped_node": "moe_grouped_gate_up", - "down_grouped_node": "moe_grouped_down", - "split_node": "moe_gate_up_split", - "offsets_node": "moe_offsets", - "combine_node": "moe_grouped_combine", - }, - }, -) - -LFM_GROUPED_MOE_NO_TOKEN_CLONE_GRAPH = _graph( - "lfm_grouped_moe_no_token_clone", - "moe_grouped_combine", - tuple(node_name for node_name in LFM_GROUPED_MOE_GRAPH.nodes if node_name != "moe_token_clone"), - edge_names=E.EDGE_GROUPS["lfm_grouped_moe_no_token_clone"], - inputs=LFM_GROUPED_MOE_GRAPH.inputs, - attr_captures=LFM_GROUPED_MOE_GRAPH.attr_captures, - constraints=LFM_GROUPED_MOE_GRAPH.constraints, -) - -LSTM_CELL_GRAPH = _graph( - "lstm_cell", - "lstm_tanh", - ("lstm_gate_mm", "lstm_recurrent_mm", "lstm_gate_add", "lstm_sigmoid", "lstm_tanh"), - edge_names=E.EDGE_GROUPS["lstm_cell"], - inputs=(_input("input", "lstm_gate_mm", 0), _input("h_prev", "lstm_recurrent_mm", 0), _input("weight_ih", "lstm_gate_mm", 1), _input("weight_hh", "lstm_recurrent_mm", 1)), - constraints={ - "lstm_gate_structure": {"gate_node": "lstm_gate_add", "gate_count": 4, "require_explicit_gate_split": True}, - }, -) - -DELTANET_DECODE_GRAPH = _graph( - "gated_deltanet_decode", - "delta_gate", - ("delta_q", "delta_k", "delta_v", "delta_gate"), - inputs=(_input("query", "delta_q", 0), _input("key", "delta_k", 0), _input("value", "delta_v", 0), _input("gate_log", "delta_gate", 0), _input("beta", "delta_gate", 1), _input("initial_state", "delta_gate", 2)), - attr_captures=(M.AttrCapture("scale", default=None, required=False),), - cache_inputs=(_cache_input("initial_state", "delta_gate", 2, cache_kind=M.CacheKind.RECURRENT, tensor_role=M.CacheTensorRole.STATE),), -) - -DELTANET_PREFILL_GRAPH = _graph( - "gated_deltanet_prefill", - "delta_prefill", - ("delta_q", "delta_k", "delta_v", "delta_gate", "delta_prefill"), - inputs=(_input("query", "delta_q", 0), _input("key", "delta_k", 0), _input("value", "delta_v", 0), _input("gate_log", "delta_gate", 0), _input("beta", "delta_gate", 1), _input("initial_state", "delta_prefill", 2)), - attr_captures=(M.AttrCapture("chunk_size", "delta_prefill", "chunk_size", required=False), M.AttrCapture("scale", default=None, required=False),), - cache_inputs=(_cache_input("initial_state", "delta_prefill", 2, cache_kind=M.CacheKind.RECURRENT, tensor_role=M.CacheTensorRole.STATE),), -) - -REL_POS_BIAS_GRAPH = _graph( - "rel_pos_bias", - "rel_pos_bias", - ("rel_pos_bias",), - inputs=( - _input("query", "rel_pos_bias", 0, tensor_constraints=(M.TensorConstraint(rank=4),)), - _input("relative_key", "rel_pos_bias", 1, tensor_constraints=(M.TensorConstraint(rank=4),)), - ), - attr_captures=(M.AttrCapture("scale", "rel_pos_bias", "scale", required=False),), -) - -KV_CACHE_STATE_GRAPH = _graph( - "kv_cache_state", - "kv_cache_state", - ("kv_cache_state",), - outputs=(M.FusionOutput("cache_state", "kv_cache_state"),), - attr_captures=(M.AttrCapture("max_seq_len", default=None, required=False), M.AttrCapture("num_kv_heads", default=None, required=False), M.AttrCapture("head_dim", default=None, required=False), M.AttrCapture("window_size", default=0, required=False), M.AttrCapture("sink_size", default=0, required=False), M.AttrCapture("num_slots", default=1, required=False)), -) - -CONV_CACHE_STATE_GRAPH = _graph( - "conv_cache_state", - "conv_cache_state", - ("conv_cache_state",), - outputs=(M.FusionOutput("cache_state", "conv_cache_state"),), - attr_captures=(M.AttrCapture("window_size", default=None, required=False), M.AttrCapture("hidden_dim", default=None, required=False)), -) - -CONV_CACHE_APPEND_GRAPH = _graph( - "conv_cache_append", - "conv_cache_append", - ("conv_cache_append",), - inputs=( - _input("new_data", "conv_cache_append", 0), - _input("cache_state", "conv_cache_append", 1, allowed_value_kinds=(M.ValueKind.CACHE_INPUT, M.ValueKind.CACHE_STATE, M.ValueKind.CACHE_OUTPUT)), - ), - cache_inputs=(_cache_input("cache_state", "conv_cache_append", 1, cache_kind=M.CacheKind.CONV, tensor_role=M.CacheTensorRole.STATE),), - cache_outputs=(_cache_output("updated_cache_state", "conv_cache_append", cache_kind=M.CacheKind.CONV, tensor_role=M.CacheTensorRole.STATE),), - cache_mutations=(_cache_mutation("append_conv_state", cache_kind=M.CacheKind.CONV, read_roles=("cache_state",), write_roles=("updated_cache_state",)),), -) - -LFM_CONV_CACHE_ROLL_APPEND_GRAPH = _graph( - "lfm_conv_cache_roll_append", - "lfm_conv_cache_output_copy", - ( - "lfm_conv_cache_clone", - "lfm_conv_cache_roll_arange", - "lfm_conv_cache_roll_add", - "lfm_conv_cache_roll_mod", - "lfm_conv_cache_gather", - "lfm_conv_cache_scatter_slice_dim0", - "lfm_conv_cache_scatter_slice_dim1", - "lfm_conv_cache_value_slice_dim0", - "lfm_conv_cache_value_slice_dim1", - "lfm_conv_cache_value_slice_last", - "lfm_conv_cache_value_copy", - "lfm_conv_cache_scatter_dim2", - "lfm_conv_cache_scatter_dim1", - "lfm_conv_cache_scatter_dim0", - "lfm_conv_cache_output_copy", - ), - edge_names=E.EDGE_GROUPS["lfm_conv_cache_roll_append"], - inputs=( - _input("new_data", "lfm_conv_cache_value_copy", 1), - _input("cache_state", "lfm_conv_cache_clone", 0, allowed_value_kinds=(M.ValueKind.CACHE_INPUT, M.ValueKind.CACHE_STATE, M.ValueKind.CACHE_OUTPUT)), - ), - cache_inputs=( - _cache_input("cache_state", "lfm_conv_cache_clone", 0, cache_kind=M.CacheKind.CONV, tensor_role=M.CacheTensorRole.STATE), - ), - cache_outputs=( - _cache_output("updated_cache_state", "lfm_conv_cache_output_copy", cache_kind=M.CacheKind.CONV, tensor_role=M.CacheTensorRole.STATE), - ), - cache_mutations=( - _cache_mutation("roll_append_conv_state", cache_kind=M.CacheKind.CONV, read_roles=("cache_state",), write_roles=("updated_cache_state",)), - ), - attr_captures=( - M.AttrCapture("window_size", default=3, required=False), - ), - constraints={ - "cache_roll_append_structure": { - "cache_role": "cache_state", - "new_data_role": "new_data", - "clone_node": "lfm_conv_cache_clone", - "gather_node": "lfm_conv_cache_gather", - "arange_node": "lfm_conv_cache_roll_arange", - "add_node": "lfm_conv_cache_roll_add", - "mod_node": "lfm_conv_cache_roll_mod", - "scatter_slice_dim0_node": "lfm_conv_cache_scatter_slice_dim0", - "scatter_slice_dim1_node": "lfm_conv_cache_scatter_slice_dim1", - "value_slice_dim0_node": "lfm_conv_cache_value_slice_dim0", - "value_slice_dim1_node": "lfm_conv_cache_value_slice_dim1", - "value_slice_last_node": "lfm_conv_cache_value_slice_last", - "scatter_dim2_node": "lfm_conv_cache_scatter_dim2", - "scatter_dim1_node": "lfm_conv_cache_scatter_dim1", - "scatter_dim0_node": "lfm_conv_cache_scatter_dim0", - }, - }, -) - -CONV_CACHE_INITIALIZE_GRAPH = _graph( - "conv_cache_initialize", - "conv_cache_initialize", - ("conv_cache_initialize",), - inputs=( - _input("rows", "conv_cache_initialize", 1), - _input("empty_cache", "conv_cache_initialize", 0, metadata={"drop_after_fusion": True}), - ), - cache_outputs=(_cache_output("initialized_cache_state", "conv_cache_initialize", cache_kind=M.CacheKind.CONV, tensor_role=M.CacheTensorRole.STATE),), - cache_mutations=(_cache_mutation("initialize_conv_state", cache_kind=M.CacheKind.CONV, write_roles=("initialized_cache_state",)),), - constraints={ - "conv_cache_initialize_structure": {"rows_role": "rows", "empty_role": "empty_cache", "output_node": "conv_cache_initialize"}, - }, -) - -RECURRENT_CACHE_STATE_GRAPH = _graph( - "recurrent_cache_state", - "recurrent_cache_state", - ("recurrent_cache_state",), - outputs=(M.FusionOutput("cache_state", "recurrent_cache_state"),), - attr_captures=(M.AttrCapture("shape", default=None, required=False), M.AttrCapture("precision", default=None, required=False)), -) - -RECURRENT_CACHE_WRITE_GRAPH = _graph( - "recurrent_cache_write", - "recurrent_cache_write", - ("recurrent_cache_write",), - inputs=(_input("new_value", "recurrent_cache_write", 0), _input("cache_input", "recurrent_cache_write", 1)), - cache_inputs=(_cache_input("cache_input", "recurrent_cache_write", 1, cache_kind=M.CacheKind.RECURRENT, tensor_role=M.CacheTensorRole.STATE),), - cache_outputs=(_cache_output("updated_cache", "recurrent_cache_write", cache_kind=M.CacheKind.RECURRENT, tensor_role=M.CacheTensorRole.STATE),), - cache_mutations=(_cache_mutation("write_recurrent_state", cache_kind=M.CacheKind.RECURRENT, read_roles=("cache_input",), write_roles=("updated_cache",)),), -) - -SAMPLE_GRAPH = _graph( - "sample", - "sample_softmax", - ("sample_topk", "sample_softmax"), - edge_names=("sample_topk_to_softmax",), - inputs=(_input("logits", "sample_topk", 0),), - attr_captures=(M.AttrCapture("temperature", default=0.6, required=False), M.AttrCapture("top_p", default=0.95, required=False), M.AttrCapture("top_k", "sample_topk", "k", default=20, required=False)), -) - -SCATTER_TOPK_GRAPH = _graph( - "scatter_topk", - "scatter", - ("scatter",), - inputs=(_input("indices", "scatter", 0), _input("values", "scatter", 1)), - attr_captures=(M.AttrCapture("num_classes", "scatter", "num_classes", required=False),), -) - -GAUSSIAN_TOPK_GRAPH = _graph( - "gaussian_topk", - "gaussian_topk", - ("gaussian_topk",), - inputs=(_input("x", "gaussian_topk", 0),), - attr_captures=(M.AttrCapture("ppf", "gaussian_topk", "ppf", required=False),), -) - -ALTUP_PREDICT_GRAPH = _graph( - "altup_predict", - "altup_predict", - ("altup_predict",), - inputs=(_input("coefs", "altup_predict", 0), _variadic_input("streams", "altup_predict", 1)), -) - -ALTUP_CORRECT_GRAPH = _graph( - "altup_correct", - "altup_correct", - ("altup_correct",), - inputs=(_input("coefs", "altup_correct", 0), _input("innovation", "altup_correct", 1), _variadic_input("predictions", "altup_correct", 2)), -) - -DSP_GRAPHS: dict[str, M.FusionGraph] = { - "stft": _graph("stft", "stft", ("stft",), inputs=(_input("x", "stft", 0), _input("weight", "stft", 1)), attr_captures=(M.AttrCapture("stride", "stft", "stride", required=False), M.AttrCapture("num_fft_bins", "stft", "num_fft_bins", required=False))), - "rfft": _graph("rfft", "rfft", ("rfft",), inputs=(_input("x", "rfft", 0),)), - "irfft": _graph("irfft", "irfft", ("irfft",), inputs=(_input("x", "irfft", 0),), attr_captures=(M.AttrCapture("output_length", "irfft", "output_length", required=False),)), - "spectrogram": _graph("spectrogram", "spectrogram", ("spectrogram",), inputs=(_input("waveform", "spectrogram", 0), _input("mel_filters", "spectrogram", 1)), metadata={"requires_preprocessing_attrs": True}), - "image_preprocess": _graph("image_preprocess", "image_preprocess", ("image_preprocess",), inputs=(_input("pixel_input", "image_preprocess", 0),), metadata={"requires_preprocessing_attrs": True}), - "bilinear_interpolation": _graph("bilinear_interpolation", "bilinear_interpolation", ("bilinear_interpolation",), inputs=(_input("pos_embeds", "bilinear_interpolation", 0),), attr_captures=(M.AttrCapture("dst_height", "bilinear_interpolation", "dst_height", required=False), M.AttrCapture("dst_width", "bilinear_interpolation", "dst_width", required=False))), -} +from . import models as M +from .fusion_attention import * +from .fusion_builders import ( + _definition, + _index_by_cactus_op, + _index_by_field, + _index_by_root_op, + _index_by_target, + _required_attrs, +) +from .fusion_cache_conv import * +from .fusion_direct import * +from .fusion_moe import * +from .fusion_neural import * +from .fusion_special import * GRAPH_BY_NAME: dict[str, M.FusionGraph] = { @@ -1280,6 +73,7 @@ def _index_by_root_op( "sqrt": "scalar_sqrt", } + DIRECT_DEFINITIONS: dict[str, M.FusionDefinition] = { name: _definition(name, cactus_op=DIRECT_CACTUS_OPS.get(name, name), graph=graph, fusion_fields=("generic", "direct")) for name, graph in DIRECT_GRAPHS.items() @@ -1359,7 +153,14 @@ def _index_by_root_op( FUSION_CATALOG = M.FusionCatalog(fusions=tuple(FUSIONS.values())) + FUSIONS_BY_TARGET: dict[str, M.FusionDefinition] = _index_by_target(FUSIONS) + + FUSIONS_BY_CACTUS_OP: dict[str, tuple[M.FusionDefinition, ...]] = _index_by_cactus_op(FUSIONS) + + FUSIONS_BY_FIELD: dict[str, tuple[M.FusionDefinition, ...]] = _index_by_field(FUSIONS) + + FUSIONS_BY_ROOT_OP: dict[str, tuple[M.FusionDefinition, ...]] = _index_by_root_op(FUSIONS) diff --git a/python/cactus/transpiler/Generator/component_split.py b/python/cactus/transpiler/Generator/component_split.py index a631a5e64..eb0cf108c 100644 --- a/python/cactus/transpiler/Generator/component_split.py +++ b/python/cactus/transpiler/Generator/component_split.py @@ -1,666 +1,7 @@ from __future__ import annotations -from collections.abc import Mapping -from dataclasses import dataclass, field -from typing import Any - -from ..Fusions import models as FModels -from ..IR import models as IRModels - - -PREFILL_WITH_CACHE_TASK = "prefill_with_cache" -DECODE_WITH_CACHE_TASK = "decode_with_cache" -GEMMA4_PREFILL_CHUNK_TOKENS = 128 - - -@dataclass(slots=True, frozen=True) -class PlaceholderSpec: - name: str - logical_name: str - source_node: str | None = None - tensor_node: str | None = None - target: str | None = None - value_kind: str = FModels.ValueKind.USER_INPUT - - -@dataclass(slots=True, frozen=True) -class OutputSpec: - node: str - logical_name: str - row_limit: int | None = None - - -@dataclass(slots=True, frozen=True) -class ComponentSplitSpec: - name: str - graph: IRModels.Graph - outputs: tuple[OutputSpec, ...] - placeholders: tuple[PlaceholderSpec, ...] = () - ref_aliases: Mapping[str, str] = field(default_factory=dict) - input_aliases: Mapping[str, str] = field(default_factory=dict) - chunk_tokens: int | None = None - - -@dataclass(slots=True, frozen=True) -class Gemma4Boundaries: - vision_features: str - audio_features: str - merged_inputs_embeds: str - prefill_text_inputs_embeds: str - prefill_token_ids: str - prefill_per_layer_inputs: str - decode_text_inputs_embeds: str - decode_token_ids: str - decode_per_layer_inputs: str - decode_position_ids: str - prefill_logits: str - decode_logits: str - - -def split_component_graphs( - graphs: Mapping[str, IRModels.Graph], - model_profile: Any | None = None, -) -> dict[str, IRModels.Graph] | None: - if not is_gemma4_profile(model_profile): - return None - - prefill = find_graph_by_task(graphs, PREFILL_WITH_CACHE_TASK) - decode = find_graph_by_task(graphs, DECODE_WITH_CACHE_TASK) - - if prefill is None or decode is None: - return None - - return split_gemma4_components(prefill, decode) - - -def split_gemma4_components(prefill: IRModels.Graph, decode: IRModels.Graph) -> dict[str, IRModels.Graph]: - boundaries = find_gemma4_boundaries(prefill, decode) - position_step = PlaceholderSpec("position_ids", "position_ids", source_node=boundaries.decode_position_ids) - position_chunk = PlaceholderSpec("position_ids", "position_ids", tensor_node="input_ids") - inputs_embeds_step = PlaceholderSpec("inputs_embeds", "inputs_embeds", source_node=boundaries.decode_text_inputs_embeds) - inputs_embeds_chunk = PlaceholderSpec("inputs_embeds", "inputs_embeds", source_node=boundaries.merged_inputs_embeds) - per_layer_step = PlaceholderSpec("per_layer_inputs", "per_layer_inputs", source_node=boundaries.decode_per_layer_inputs) - per_layer_chunk = PlaceholderSpec("per_layer_inputs", "per_layer_inputs", source_node=boundaries.prefill_per_layer_inputs) - - specs = ( - ComponentSplitSpec( - name="vision_encoder", - graph=prefill, - outputs=(OutputSpec(boundaries.vision_features, "image_features", row_limit=256),), - input_aliases={"image_position_ids": "pixel_position_ids"}, - ), - ComponentSplitSpec( - name="audio_encoder", - graph=prefill, - outputs=(OutputSpec(boundaries.audio_features, "audio_features"),), - ), - ComponentSplitSpec( - name="lm_encoder_step", - graph=decode, - outputs=( - OutputSpec(boundaries.decode_text_inputs_embeds, "inputs_embeds"), - OutputSpec(boundaries.decode_per_layer_inputs, "per_layer_inputs"), - OutputSpec(position_step.name, "position_ids"), - ), - placeholders=(position_step,), - ref_aliases={boundaries.decode_token_ids: "input_ids"}, - ), - ComponentSplitSpec( - name="lm_encoder_text_chunk", - graph=prefill, - outputs=( - OutputSpec(boundaries.prefill_text_inputs_embeds, "inputs_embeds"), - OutputSpec(boundaries.prefill_per_layer_inputs, "per_layer_inputs"), - OutputSpec(position_chunk.name, "position_ids"), - ), - placeholders=(position_chunk,), - ref_aliases={ - boundaries.merged_inputs_embeds: boundaries.prefill_text_inputs_embeds, - boundaries.prefill_token_ids: "input_ids", - }, - chunk_tokens=GEMMA4_PREFILL_CHUNK_TOKENS, - ), - ComponentSplitSpec( - name="lm_encoder_media_chunk", - graph=prefill, - outputs=( - OutputSpec(inputs_embeds_chunk.name, "inputs_embeds"), - OutputSpec(boundaries.prefill_per_layer_inputs, "per_layer_inputs"), - OutputSpec(position_chunk.name, "position_ids"), - ), - placeholders=(inputs_embeds_chunk, position_chunk), - ref_aliases={ - boundaries.merged_inputs_embeds: inputs_embeds_chunk.name, - boundaries.prefill_token_ids: "input_ids", - }, - chunk_tokens=GEMMA4_PREFILL_CHUNK_TOKENS, - ), - ComponentSplitSpec( - name="lm_encoder_media_step", - graph=decode, - outputs=( - OutputSpec(inputs_embeds_step.name, "inputs_embeds"), - OutputSpec(boundaries.decode_per_layer_inputs, "per_layer_inputs"), - OutputSpec(position_step.name, "position_ids"), - ), - placeholders=(inputs_embeds_step, position_step), - ref_aliases={boundaries.decode_token_ids: "input_ids"}, - ), - ComponentSplitSpec( - name="decoder_prefill_chunk", - graph=prefill, - outputs=(OutputSpec(boundaries.prefill_logits, "logits"),), - placeholders=(inputs_embeds_chunk, per_layer_chunk, position_chunk), - chunk_tokens=GEMMA4_PREFILL_CHUNK_TOKENS, - ), - ComponentSplitSpec( - name="decoder_step", - graph=decode, - outputs=(OutputSpec(boundaries.decode_logits, "logits"),), - placeholders=(inputs_embeds_step, per_layer_step, position_step), - ), - ) - - return {spec.name: extract_component_graph(spec) for spec in specs} - - -def is_gemma4_profile(model_profile: Any | None) -> bool: - profile_name = str(getattr(model_profile, "model_profiles", "") or "").lower() - return "gemma4" in profile_name or "gemma_4" in profile_name - - -def find_graph_by_task(graphs: Mapping[str, IRModels.Graph], task: str) -> IRModels.Graph | None: - for graph in graphs.values(): - if graph.task == task: - return graph - - for name, graph in graphs.items(): - if task in name: - return graph - - return None - - -def find_gemma4_boundaries(prefill: IRModels.Graph, decode: IRModels.Graph) -> Gemma4Boundaries: - vision_merge, audio_merge = find_gemma4_masked_scatter_merges(prefill) - - return Gemma4Boundaries( - vision_features=feature_parent(vision_merge, "vision"), - audio_features=feature_parent(audio_merge, "audio"), - merged_inputs_embeds=audio_merge.name, - prefill_text_inputs_embeds=find_text_inputs_embeds(prefill), - prefill_token_ids=find_embedding_token_ids(prefill), - prefill_per_layer_inputs=find_per_layer_inputs(prefill), - decode_text_inputs_embeds=find_text_inputs_embeds(decode), - decode_token_ids=find_embedding_token_ids(decode), - decode_per_layer_inputs=find_per_layer_inputs(decode), - decode_position_ids=find_decode_position_ids(decode), - prefill_logits=find_logits_output(prefill), - decode_logits=find_logits_output(decode), - ) - - -def find_gemma4_masked_scatter_merges(graph: IRModels.Graph) -> tuple[IRModels.Node, IRModels.Node]: - merges = [ - node - for node in graph.nodes - if node.target == "aten.masked_scatter.default" - and tensor_rank(node) == 3 - and tensor_last_dim(node) == 1536 - and len(node.parents) >= 3 - ] - - if len(merges) < 2: - raise ValueError("Gemma4 component split requires the vision/audio masked_scatter merge nodes") - - return tuple(sorted(merges, key=lambda node: node.index)[:2]) - - -def feature_parent(merge: IRModels.Node, role: str) -> str: - if len(merge.parents) < 3: - raise ValueError(f"{merge.name}: Gemma4 {role} merge does not have a feature parent") - - return merge.parents[2].name - - -def find_text_inputs_embeds(graph: IRModels.Graph) -> str: - for node in graph.nodes: - if node.target not in {"aten.mul.Tensor", "cactus.multiply"} or tensor_rank(node) != 3 or tensor_last_dim(node) != 1536: - continue - - if any(parent.target in {"aten.embedding.default", "cactus.embedding_from_tensor"} for parent in node.parents): - return node.name - - raise ValueError("Gemma4 component split could not find text inputs_embeds") - - -def find_embedding_token_ids(graph: IRModels.Graph) -> str: - text_inputs_embeds = graph.nodes_map[find_text_inputs_embeds(graph)] - - for parent in text_inputs_embeds.parents: - if parent.target not in {"aten.embedding.default", "cactus.embedding_from_tensor"}: - continue - - for embedding_parent in parent.parents: - if embedding_parent.value_kind != FModels.ValueKind.PARAMETER: - return embedding_parent.name - - raise ValueError("Gemma4 component split could not find embedding token-id input") - - -def find_per_layer_inputs(graph: IRModels.Graph) -> str: - candidates = [ - node - for node in graph.nodes - if node.target in {"aten.mul.Tensor", "aten.mul.Scalar", "cactus.multiply", "cactus.scalar_multiply"} - and tensor_shape(node)[-2:] == [35, 256] - and len(node.children) >= 10 - ] - - if not candidates: - raise ValueError("Gemma4 component split could not find per_layer_inputs") - - return max(candidates, key=lambda node: len(node.children)).name - - -def find_decode_position_ids(graph: IRModels.Graph) -> str: - for node in graph.nodes: - if node.target == "aten.unsqueeze.default" and tensor_shape(node) == [1, 1]: - return node.name - - for node in graph.nodes: - if node.name == "cache_position": - return node.name - - raise ValueError("Gemma4 component split could not find decode position_ids") - - -def find_logits_output(graph: IRModels.Graph) -> str: - output = graph.outputs[0] if graph.outputs else None - refs = IRModels.extract_node_refs((output.args, output.kwargs)) if output is not None else () - - if refs: - return refs[0] - - for node in reversed(graph.nodes): - if tensor_rank(node) == 3 and tensor_last_dim(node) == 262144: - return node.name - - raise ValueError("Gemma4 component split could not find logits output") - - -def extract_component_graph(spec: ComponentSplitSpec) -> IRModels.Graph: - placeholder_specs = {placeholder.name: placeholder for placeholder in spec.placeholders} - ref_replacements = { - **dict(spec.ref_aliases), - **{ - placeholder.source_node: placeholder.name - for placeholder in spec.placeholders - if placeholder.source_node is not None - }, - } - required_names = collect_required_node_names(spec, placeholder_specs, ref_replacements) - output_names = tuple(ref_replacements.get(output.node, output.node) for output in spec.outputs) - output_logical_names = { - ref_replacements.get(output.node, output.node): output.logical_name - for output in spec.outputs - if output.row_limit is None - } - used_names: set[str] = set() - nodes: list[IRModels.Node] = [] - - for placeholder in spec.placeholders: - if placeholder.name not in required_names and placeholder.name not in output_names: - continue - - nodes.append(create_placeholder_node(spec, placeholder, output_logical_names, used_names)) - - for node in spec.graph.nodes: - if node.name not in required_names: - continue - - if node.name in ref_replacements: - continue - - nodes.append(clone_component_node(spec, node, ref_replacements, output_logical_names, used_names)) - - final_output_names = [] - - for output, output_name in zip(spec.outputs, output_names): - final_output_name = output_name - - if output.row_limit is not None: - final_output_name = add_row_limited_output_node(nodes, output_name, output.logical_name, output.row_limit, used_names) - - output_logical_names[final_output_name] = output.logical_name - final_output_names.append(final_output_name) - - nodes.append(create_output_node(spec.name, tuple(final_output_names), nodes)) - graph = IRModels.rebuild_graph(tuple(nodes), spec.graph) - - if spec.chunk_tokens is not None: - graph = retarget_chunk_graph_sequence_length(graph, spec.chunk_tokens) - - return graph - - -def add_row_limited_output_node( - nodes: list[IRModels.Node], - source_name: str, - logical_name: str, - row_limit: int, - used_names: set[str], -) -> str: - source = next((node for node in nodes if node.name == source_name), None) - - if source is None: - raise ValueError(f"Component split cannot add row-limited output for missing node {source_name}") - - source_shape = tensor_shape(source) - - if not source_shape: - raise ValueError(f"{source_name}: row-limited output requires tensor shape metadata") - - output_shape = list(source_shape) - output_shape[0] = min(row_limit, int(output_shape[0])) if isinstance(output_shape[0], int) else row_limit - name = unique_component_node_name(f"{source_name}_first_{row_limit}", used_names) - - source_meta = source.tensor_output_meta if isinstance(source.tensor_output_meta, dict) else {} - nodes.append( - IRModels.Node( - index=max((node.index for node in nodes), default=-1) + 1, - name=name, - node_type="call_function", - target="cactus.slice", - args=[{"node": source_name}], - kwargs={}, - users=(), - tensor_output_meta={**source_meta, "shape": output_shape}, - module_stack=source.module_stack, - value_kind=FModels.ValueKind.ACTIVATION, - attrs={"axis": 0, "start": 0, "length": row_limit, "step": 1}, - ir_metadata={"logical_output": logical_name}, - cache=None, - ) - ) - - return name - - -def retarget_chunk_graph_sequence_length(graph: IRModels.Graph, chunk_tokens: int) -> IRModels.Graph: - source_lengths = chunk_source_lengths(graph) - - if not source_lengths: - return graph - - cache_capacity = max(source_lengths) - nodes = tuple(retarget_chunk_node(node, source_lengths, chunk_tokens, cache_capacity) for node in graph.nodes) - return IRModels.rebuild_graph(nodes, graph) - - -def chunk_source_lengths(graph: IRModels.Graph) -> frozenset[int]: - lengths: set[int] = set() - - for node in graph.nodes: - if not node.is_placeholder: - continue - - logical_name = str(node.ir_metadata.get("logical_input") or node.target) - - if logical_name not in {"input_ids", "position_ids", "inputs_embeds", "per_layer_inputs"}: - continue - - shape = tensor_shape(node) - - if len(shape) >= 2 and shape[0] == 1 and isinstance(shape[1], int) and shape[1] > 1: - lengths.add(shape[1]) - - return frozenset(lengths) - - -def retarget_chunk_node( - node: IRModels.Node, - source_lengths: frozenset[int], - chunk_tokens: int, - cache_capacity: int, -) -> IRModels.Node: - metadata = dict(node.ir_metadata) - - if is_empty_cache_cat_node(node): - metadata.setdefault("prefill_cache_capacity", cache_capacity) - - if node.is_placeholder and metadata.get("logical_input") in {"input_ids", "position_ids", "inputs_embeds", "per_layer_inputs"}: - metadata.setdefault("prefill_chunk_tokens", chunk_tokens) - - return IRModels.Node( - index=node.index, - name=node.name, - node_type=node.node_type, - target=node.target, - args=retarget_sequence_value(node.args, source_lengths, chunk_tokens), - kwargs=retarget_sequence_value(node.kwargs, source_lengths, chunk_tokens), - users=node.users, - tensor_output_meta=retarget_sequence_value(node.tensor_output_meta, source_lengths, chunk_tokens), - module_stack=node.module_stack, - value_kind=node.value_kind, - attrs=retarget_sequence_value(dict(node.attrs), source_lengths, chunk_tokens), - ir_metadata=metadata, - cache=node.cache, - ) - - -def retarget_sequence_value(value: Any, source_lengths: frozenset[int], chunk_tokens: int) -> Any: - if isinstance(value, bool): - return value - - if isinstance(value, int): - return chunk_tokens if value in source_lengths else value - - if isinstance(value, list): - return [retarget_sequence_value(item, source_lengths, chunk_tokens) for item in value] - - if isinstance(value, tuple): - return tuple(retarget_sequence_value(item, source_lengths, chunk_tokens) for item in value) - - if isinstance(value, dict): - return {key: retarget_sequence_value(item, source_lengths, chunk_tokens) for key, item in value.items()} - - return value - - -def is_empty_cache_cat_node(node: IRModels.Node) -> bool: - if node.target not in {"cactus.cat", "aten.cat.default"} or len(node.parents) < 2: - return False - - first_parent = node.parents[0] - return first_parent.value_kind == FModels.ValueKind.LIFTED_CONSTANT and element_count(tensor_shape(first_parent)) == 0 - - -def element_count(shape: list[Any]) -> int | None: - product = 1 - - for dim in shape: - if not isinstance(dim, int): - return None - - product *= dim - - return product - - -def collect_required_node_names( - spec: ComponentSplitSpec, - placeholder_specs: Mapping[str, PlaceholderSpec], - ref_replacements: Mapping[str, str], -) -> frozenset[str]: - required: set[str] = set() - - for output in spec.outputs: - collect_node_ancestors(output.node, spec.graph, required, placeholder_specs, ref_replacements) - - for placeholder in spec.placeholders: - if placeholder.source_node is None: - required.add(placeholder.name) - - return frozenset(required) - - -def collect_node_ancestors( - node_name: str, - graph: IRModels.Graph, - required: set[str], - placeholder_specs: Mapping[str, PlaceholderSpec], - ref_replacements: Mapping[str, str], -) -> None: - stack = [node_name] - - while stack: - current_name = stack.pop() - effective_name = ref_replacements.get(current_name, current_name) - - if effective_name in required: - continue - - required.add(effective_name) - - if effective_name in placeholder_specs: - continue - - node = graph.nodes_map.get(effective_name) - - if node is None: - raise ValueError(f"Component split references missing node {effective_name}") - - stack.extend(parent.name for parent in node.parents) - - -def create_placeholder_node( - spec: ComponentSplitSpec, - placeholder: PlaceholderSpec, - output_logical_names: Mapping[str, str], - used_names: set[str], -) -> IRModels.Node: - template = source_template_node(spec.graph, placeholder) - name = unique_component_node_name(placeholder.name, used_names) - metadata = dict(template.ir_metadata) if template is not None else {} - metadata["logical_input"] = placeholder.logical_name - - if placeholder.name in output_logical_names: - metadata["logical_output"] = output_logical_names[placeholder.name] - - return IRModels.Node( - index=template.index if template is not None else -1, - name=name, - node_type="placeholder", - target=placeholder.target or placeholder.logical_name, - args=[], - kwargs={}, - users=(), - tensor_output_meta=template.tensor_output_meta if template is not None else None, - module_stack=template.module_stack if template is not None else None, - value_kind=placeholder.value_kind, - attrs={}, - ir_metadata=metadata, - cache=template.cache if template is not None else None, - ) - - -def source_template_node(graph: IRModels.Graph, placeholder: PlaceholderSpec) -> IRModels.Node | None: - for name in (placeholder.tensor_node, placeholder.source_node, placeholder.name): - if name is not None and name in graph.nodes_map: - return graph.nodes_map[name] - - return None - - -def clone_component_node( - spec: ComponentSplitSpec, - node: IRModels.Node, - ref_replacements: Mapping[str, str], - output_logical_names: Mapping[str, str], - used_names: set[str], -) -> IRModels.Node: - metadata = dict(node.ir_metadata) - - if node.is_placeholder: - logical_input = spec.input_aliases.get(node.name) or spec.input_aliases.get(node.target) - - if logical_input is not None: - metadata["logical_input"] = logical_input - - if node.name in output_logical_names: - metadata["logical_output"] = output_logical_names[node.name] - - name = unique_component_node_name(node.name, used_names) - - return IRModels.Node( - index=node.index, - name=name, - node_type=node.node_type, - target=node.target, - args=IRModels.rewrite_node_refs(node.args, dict(ref_replacements)), - kwargs=IRModels.rewrite_node_refs(node.kwargs, dict(ref_replacements)), - users=(), - tensor_output_meta=node.tensor_output_meta, - module_stack=node.module_stack, - value_kind=node.value_kind, - attrs=dict(node.attrs), - ir_metadata=metadata, - cache=node.cache, - ) - - -def create_output_node( - component_name: str, - output_names: tuple[str, ...], - nodes: list[IRModels.Node], -) -> IRModels.Node: - max_index = max((node.index for node in nodes), default=-1) - - return IRModels.Node( - index=max_index + 1, - name=f"{component_name}_output", - node_type="output", - target="output", - args=[[{"node": name} for name in output_names]], - kwargs={}, - users=(), - tensor_output_meta=None, - module_stack=None, - value_kind=FModels.ValueKind.OUTPUT, - attrs={}, - ir_metadata={}, - ) - - -def unique_component_node_name(name: str, used_names: set[str]) -> str: - if name not in used_names: - used_names.add(name) - return name - - suffix = 1 - candidate = f"{name}_{suffix}" - - while candidate in used_names: - suffix += 1 - candidate = f"{name}_{suffix}" - - used_names.add(candidate) - return candidate - - -def tensor_shape(node: IRModels.Node) -> list[Any]: - if not isinstance(node.tensor_output_meta, dict): - return [] - - shape = node.tensor_output_meta.get("shape") - return list(shape) if isinstance(shape, list) else [] - - -def tensor_rank(node: IRModels.Node) -> int: - return len(tensor_shape(node)) - - -def tensor_last_dim(node: IRModels.Node) -> Any | None: - shape = tensor_shape(node) - return shape[-1] if shape else None +from .component_split_boundaries import * +from .component_split_builder import * +from .component_split_plans import * +from .component_split_types import * +from .component_split_utils import * diff --git a/python/cactus/transpiler/Generator/component_split_boundaries.py b/python/cactus/transpiler/Generator/component_split_boundaries.py new file mode 100644 index 000000000..9c6c35b17 --- /dev/null +++ b/python/cactus/transpiler/Generator/component_split_boundaries.py @@ -0,0 +1,234 @@ +from __future__ import annotations + +from .component_split_types import Gemma4Boundaries, OutputSpec, WhisperBoundaries +from .component_split_utils import tensor_last_dim, tensor_rank, tensor_shape +from ..IR import models as IRModels + + +def find_graph_by_task(graphs: Mapping[str, IRModels.Graph], task: str) -> IRModels.Graph | None: + for graph in graphs.values(): + if graph.task == task: + return graph + + for name, graph in graphs.items(): + if task in name: + return graph + + return None + + +def find_whisper_boundaries(prefill: IRModels.Graph, decode: IRModels.Graph) -> WhisperBoundaries: + return WhisperBoundaries( + encoder_hidden_states=find_whisper_encoder_hidden_states(prefill), + cross_key_values=find_whisper_cross_key_values(prefill), + decode_logits=find_logits_output(decode), + ) + + +def find_whisper_encoder_hidden_states(graph: IRModels.Graph) -> str: + candidates = [ + node + for node in graph.nodes + if node.target == "" + and tensor_shape(node) == [1, 1500, 384] + and len(node.children) >= 4 + ] + + if not candidates: + raise ValueError("Whisper component split could not find encoder_hidden_states") + + return max(candidates, key=lambda node: len(node.children)).name + + +def find_whisper_cross_key_values(graph: IRModels.Graph) -> tuple[tuple[str, str], ...]: + pairs: list[tuple[str, str]] = [] + + for node in graph.nodes: + if node.target != "cactus.attention" or len(node.parents) < 3: + continue + + key_node, value_node = node.parents[1], node.parents[2] + + if ( + key_node.target == "cactus.cat" + and value_node.target == "cactus.cat" + and tensor_shape(key_node) == [1, 6, 1500, 64] + and tensor_shape(value_node) == [1, 6, 1500, 64] + ): + pairs.append((key_node.name, value_node.name)) + + if not pairs: + raise ValueError("Whisper component split could not find decoder cross-attention KV outputs") + + return tuple(pairs) + + +def find_decoder_token_input(graph: IRModels.Graph) -> str: + for name in ("decoder_input_ids", "input_ids"): + if name in graph.nodes_map: + return name + + raise ValueError("Whisper component split could not find decoder token input") + + +def find_whisper_decode_position_index(graph: IRModels.Graph) -> str: + for node in graph.nodes: + if node.target == "aten.repeat.default" and tensor_shape(node) == [1, 1]: + if any(child.target == "aten.index.Tensor" for child in node.children): + return node.name + + raise ValueError("Whisper component split could not find decoder position index") + + +def whisper_decoder_cross_input_aliases(graph: IRModels.Graph, layer_count: int) -> dict[str, str]: + aliases: dict[str, str] = {} + + for node in graph.nodes: + if node.cache is None or node.cache.tensor_index is None: + continue + + tensor_index = int(node.cache.tensor_index) + layer_index = tensor_index // 4 + + if layer_index >= layer_count: + continue + + role_index = tensor_index % 4 + + if role_index == 2: + aliases[node.name] = f"cross_k_{layer_index}" + elif role_index == 3: + aliases[node.name] = f"cross_v_{layer_index}" + + return aliases + + +def find_gemma4_boundaries(prefill: IRModels.Graph, decode: IRModels.Graph) -> Gemma4Boundaries: + vision_merge, audio_merge = find_gemma4_masked_scatter_merges(prefill) + + return Gemma4Boundaries( + vision_features=feature_parent(vision_merge, "vision"), + audio_features=feature_parent(audio_merge, "audio"), + merged_inputs_embeds=audio_merge.name, + prefill_text_inputs_embeds=find_text_inputs_embeds(prefill), + prefill_token_ids=find_embedding_token_ids(prefill), + prefill_per_layer_inputs=find_per_layer_inputs(prefill), + decode_text_inputs_embeds=find_text_inputs_embeds(decode), + decode_token_ids=find_embedding_token_ids(decode), + decode_per_layer_inputs=find_per_layer_inputs(decode), + decode_position_ids=find_decode_position_ids(decode), + prefill_logits=find_logits_output(prefill), + decode_logits=find_logits_output(decode), + ) + + +def find_gemma4_masked_scatter_merges(graph: IRModels.Graph) -> tuple[IRModels.Node, IRModels.Node]: + merges = [ + node + for node in graph.nodes + if node.target == "aten.masked_scatter.default" + and tensor_rank(node) == 3 + and tensor_last_dim(node) == 1536 + and len(node.parents) >= 3 + ] + + if len(merges) < 2: + raise ValueError("Gemma4 component split requires the vision/audio masked_scatter merge nodes") + + return tuple(sorted(merges, key=lambda node: node.index)[:2]) + + +def feature_parent(merge: IRModels.Node, role: str) -> str: + if len(merge.parents) < 3: + raise ValueError(f"{merge.name}: Gemma4 {role} merge does not have a feature parent") + + return merge.parents[2].name + + +def find_text_inputs_embeds(graph: IRModels.Graph) -> str: + for node in graph.nodes: + if node.target not in {"aten.mul.Tensor", "cactus.multiply"} or tensor_rank(node) != 3 or tensor_last_dim(node) != 1536: + continue + + if any(parent.target in {"aten.embedding.default", "cactus.embedding_from_tensor"} for parent in node.parents): + return node.name + + raise ValueError("Gemma4 component split could not find text inputs_embeds") + + +def find_embedding_token_ids(graph: IRModels.Graph) -> str: + text_inputs_embeds = graph.nodes_map[find_text_inputs_embeds(graph)] + + for parent in text_inputs_embeds.parents: + if parent.target not in {"aten.embedding.default", "cactus.embedding_from_tensor"}: + continue + + for embedding_parent in parent.parents: + if embedding_parent.value_kind != FModels.ValueKind.PARAMETER: + return embedding_parent.name + + raise ValueError("Gemma4 component split could not find embedding token-id input") + + +def find_per_layer_inputs(graph: IRModels.Graph) -> str: + candidates = [ + node + for node in graph.nodes + if node.target in {"aten.mul.Tensor", "aten.mul.Scalar", "cactus.multiply", "cactus.scalar_multiply"} + and tensor_shape(node)[-2:] == [35, 256] + and len(node.children) >= 10 + ] + + if not candidates: + raise ValueError("Gemma4 component split could not find per_layer_inputs") + + return max(candidates, key=lambda node: len(node.children)).name + + +def find_decode_position_ids(graph: IRModels.Graph) -> str: + for node in graph.nodes: + if node.target == "aten.unsqueeze.default" and tensor_shape(node) == [1, 1]: + return node.name + + for node in graph.nodes: + if node.name == "cache_position": + return node.name + + raise ValueError("Gemma4 component split could not find decode position_ids") + + +def find_logits_output(graph: IRModels.Graph) -> str: + output = graph.outputs[0] if graph.outputs else None + refs = IRModels.extract_node_refs((output.args, output.kwargs)) if output is not None else () + + if refs: + return refs[0] + + for node in reversed(graph.nodes): + if tensor_rank(node) == 3 and tensor_last_dim(node) == 262144: + return node.name + + raise ValueError("Gemma4 component split could not find logits output") + + +def graph_output_specs(graph: IRModels.Graph) -> tuple[OutputSpec, ...]: + refs = graph_output_refs(graph) + + if not refs: + return (OutputSpec(find_logits_output(graph), "logits"),) + + logits = refs[0] + + return tuple( + OutputSpec(ref, "logits" if ref == logits else ref) + for ref in refs + ) + + +def graph_output_refs(graph: IRModels.Graph) -> tuple[str, ...]: + refs: list[str] = [] + + for output in graph.outputs: + refs.extend(IRModels.extract_node_refs((output.args, output.kwargs))) + + return tuple(dict.fromkeys(refs)) diff --git a/python/cactus/transpiler/Generator/component_split_builder.py b/python/cactus/transpiler/Generator/component_split_builder.py new file mode 100644 index 000000000..23dda1cba --- /dev/null +++ b/python/cactus/transpiler/Generator/component_split_builder.py @@ -0,0 +1,374 @@ +from __future__ import annotations + +from collections.abc import Mapping +from typing import Any + +from .component_split_types import ComponentSplitSpec, PlaceholderSpec +from .component_split_utils import element_count, tensor_shape +from ..Fusions import models as FModels +from ..IR import models as IRModels + + +def extract_component_graph(spec: ComponentSplitSpec) -> IRModels.Graph: + placeholder_specs = {placeholder.name: placeholder for placeholder in spec.placeholders} + ref_replacements = { + **dict(spec.ref_aliases), + **{ + placeholder.source_node: placeholder.name + for placeholder in spec.placeholders + if placeholder.source_node is not None + }, + } + required_names = collect_required_node_names(spec, placeholder_specs, ref_replacements) + output_names = tuple(ref_replacements.get(output.node, output.node) for output in spec.outputs) + output_logical_names = { + ref_replacements.get(output.node, output.node): output.logical_name + for output in spec.outputs + if output.row_limit is None + } + used_names: set[str] = set() + nodes: list[IRModels.Node] = [] + + for placeholder in spec.placeholders: + if not placeholder.force and placeholder.name not in required_names and placeholder.name not in output_names: + continue + + nodes.append(create_placeholder_node(spec, placeholder, output_logical_names, used_names)) + + for node in spec.graph.nodes: + if node.name not in required_names: + continue + + if node.name in ref_replacements: + continue + + nodes.append(clone_component_node(spec, node, ref_replacements, output_logical_names, used_names)) + + final_output_names = [] + + for output, output_name in zip(spec.outputs, output_names): + final_output_name = output_name + + if output.row_limit is not None: + final_output_name = add_row_limited_output_node(nodes, output_name, output.logical_name, output.row_limit, used_names) + + output_logical_names[final_output_name] = output.logical_name + final_output_names.append(final_output_name) + + nodes.append(create_output_node(spec.name, tuple(final_output_names), nodes)) + graph = IRModels.rebuild_graph(tuple(nodes), spec.graph) + graph.metadata.update({str(key): str(value) for key, value in spec.metadata.items()}) + + if spec.chunk_tokens is not None: + graph = retarget_chunk_graph_sequence_length(graph, spec.chunk_tokens) + + return graph + + +def add_row_limited_output_node( + nodes: list[IRModels.Node], + source_name: str, + logical_name: str, + row_limit: int, + used_names: set[str], +) -> str: + source = next((node for node in nodes if node.name == source_name), None) + + if source is None: + raise ValueError(f"Component split cannot add row-limited output for missing node {source_name}") + + source_shape = tensor_shape(source) + + if not source_shape: + raise ValueError(f"{source_name}: row-limited output requires tensor shape metadata") + + output_shape = list(source_shape) + output_shape[0] = min(row_limit, int(output_shape[0])) if isinstance(output_shape[0], int) else row_limit + name = unique_component_node_name(f"{source_name}_first_{row_limit}", used_names) + + source_meta = source.tensor_output_meta if isinstance(source.tensor_output_meta, dict) else {} + nodes.append( + IRModels.Node( + index=max((node.index for node in nodes), default=-1) + 1, + name=name, + node_type="call_function", + target="cactus.slice", + args=[{"node": source_name}], + kwargs={}, + users=(), + tensor_output_meta={**source_meta, "shape": output_shape}, + module_stack=source.module_stack, + value_kind=FModels.ValueKind.ACTIVATION, + attrs={"axis": 0, "start": 0, "length": row_limit, "step": 1}, + ir_metadata={"logical_output": logical_name}, + cache=None, + ) + ) + + return name + + +def retarget_chunk_graph_sequence_length(graph: IRModels.Graph, chunk_tokens: int) -> IRModels.Graph: + source_lengths = chunk_source_lengths(graph) + + if not source_lengths: + return graph + + cache_capacity = max(source_lengths) + nodes = tuple(retarget_chunk_node(node, source_lengths, chunk_tokens, cache_capacity) for node in graph.nodes) + return IRModels.rebuild_graph(nodes, graph) + + +def chunk_source_lengths(graph: IRModels.Graph) -> frozenset[int]: + lengths: set[int] = set() + + for node in graph.nodes: + if not node.is_placeholder: + continue + + logical_name = str(node.ir_metadata.get("logical_input") or node.target) + + if logical_name not in {"input_ids", "position_ids", "inputs_embeds", "per_layer_inputs"}: + continue + + shape = tensor_shape(node) + + if len(shape) >= 2 and shape[0] == 1 and isinstance(shape[1], int) and shape[1] > 1: + lengths.add(shape[1]) + + return frozenset(lengths) + + +def retarget_chunk_node( + node: IRModels.Node, + source_lengths: frozenset[int], + chunk_tokens: int, + cache_capacity: int, +) -> IRModels.Node: + metadata = dict(node.ir_metadata) + + if is_empty_cache_cat_node(node): + metadata.setdefault("prefill_cache_capacity", cache_capacity) + + if node.is_placeholder and metadata.get("logical_input") in {"input_ids", "position_ids", "inputs_embeds", "per_layer_inputs"}: + metadata.setdefault("prefill_chunk_tokens", chunk_tokens) + + return IRModels.Node( + index=node.index, + name=node.name, + node_type=node.node_type, + target=node.target, + args=retarget_sequence_value(node.args, source_lengths, chunk_tokens), + kwargs=retarget_sequence_value(node.kwargs, source_lengths, chunk_tokens), + users=node.users, + tensor_output_meta=retarget_sequence_value(node.tensor_output_meta, source_lengths, chunk_tokens), + module_stack=node.module_stack, + value_kind=node.value_kind, + attrs=retarget_sequence_value(dict(node.attrs), source_lengths, chunk_tokens), + ir_metadata=metadata, + cache=node.cache, + ) + + +def retarget_sequence_value(value: Any, source_lengths: frozenset[int], chunk_tokens: int) -> Any: + if isinstance(value, bool): + return value + + if isinstance(value, int): + return chunk_tokens if value in source_lengths else value + + if isinstance(value, list): + return [retarget_sequence_value(item, source_lengths, chunk_tokens) for item in value] + + if isinstance(value, tuple): + return tuple(retarget_sequence_value(item, source_lengths, chunk_tokens) for item in value) + + if isinstance(value, dict): + return { + key: item if is_dimension_index_attr(key) else retarget_sequence_value(item, source_lengths, chunk_tokens) + for key, item in value.items() + } + + return value + + +def is_dimension_index_attr(key: Any) -> bool: + return str(key) in { + "axis", + "dim", + "dim0", + "dim1", + "dims", + "end_dim", + "permutation", + "start_dim", + } + + +def is_empty_cache_cat_node(node: IRModels.Node) -> bool: + if node.target not in {"cactus.cat", "aten.cat.default"} or len(node.parents) < 2: + return False + + first_parent = node.parents[0] + return first_parent.value_kind == FModels.ValueKind.LIFTED_CONSTANT and element_count(tensor_shape(first_parent)) == 0 + + +def collect_required_node_names( + spec: ComponentSplitSpec, + placeholder_specs: Mapping[str, PlaceholderSpec], + ref_replacements: Mapping[str, str], +) -> frozenset[str]: + required: set[str] = set() + + for output in spec.outputs: + collect_node_ancestors(output.node, spec.graph, required, placeholder_specs, ref_replacements) + + for placeholder in spec.placeholders: + if placeholder.source_node is None: + required.add(placeholder.name) + + return frozenset(required) + + +def collect_node_ancestors( + node_name: str, + graph: IRModels.Graph, + required: set[str], + placeholder_specs: Mapping[str, PlaceholderSpec], + ref_replacements: Mapping[str, str], +) -> None: + stack = [node_name] + + while stack: + current_name = stack.pop() + effective_name = ref_replacements.get(current_name, current_name) + + if effective_name in required: + continue + + required.add(effective_name) + + if effective_name in placeholder_specs: + continue + + node = graph.nodes_map.get(effective_name) + + if node is None: + raise ValueError(f"Component split references missing node {effective_name}") + + stack.extend(parent.name for parent in node.parents) + + +def create_placeholder_node( + spec: ComponentSplitSpec, + placeholder: PlaceholderSpec, + output_logical_names: Mapping[str, str], + used_names: set[str], +) -> IRModels.Node: + template = source_template_node(spec.graph, placeholder) + name = unique_component_node_name(placeholder.name, used_names) + metadata = dict(template.ir_metadata) if template is not None else {} + metadata["logical_input"] = placeholder.logical_name + + if placeholder.name in output_logical_names: + metadata["logical_output"] = output_logical_names[placeholder.name] + + return IRModels.Node( + index=template.index if template is not None else -1, + name=name, + node_type="placeholder", + target=placeholder.target or placeholder.logical_name, + args=[], + kwargs={}, + users=(), + tensor_output_meta=template.tensor_output_meta if template is not None else None, + module_stack=template.module_stack if template is not None else None, + value_kind=placeholder.value_kind, + attrs={}, + ir_metadata=metadata, + cache=template.cache if template is not None else None, + ) + + +def source_template_node(graph: IRModels.Graph, placeholder: PlaceholderSpec) -> IRModels.Node | None: + for name in (placeholder.tensor_node, placeholder.source_node, placeholder.name): + if name is not None and name in graph.nodes_map: + return graph.nodes_map[name] + + return None + + +def clone_component_node( + spec: ComponentSplitSpec, + node: IRModels.Node, + ref_replacements: Mapping[str, str], + output_logical_names: Mapping[str, str], + used_names: set[str], +) -> IRModels.Node: + metadata = dict(node.ir_metadata) + + if node.is_placeholder: + logical_input = spec.input_aliases.get(node.name) or spec.input_aliases.get(node.target) + + if logical_input is not None: + metadata["logical_input"] = logical_input + + if node.name in output_logical_names: + metadata["logical_output"] = output_logical_names[node.name] + + name = unique_component_node_name(node.name, used_names) + + return IRModels.Node( + index=node.index, + name=name, + node_type=node.node_type, + target=node.target, + args=IRModels.rewrite_node_refs(node.args, dict(ref_replacements)), + kwargs=IRModels.rewrite_node_refs(node.kwargs, dict(ref_replacements)), + users=(), + tensor_output_meta=node.tensor_output_meta, + module_stack=node.module_stack, + value_kind=node.value_kind, + attrs=dict(node.attrs), + ir_metadata=metadata, + cache=node.cache, + ) + + +def create_output_node( + component_name: str, + output_names: tuple[str, ...], + nodes: list[IRModels.Node], +) -> IRModels.Node: + max_index = max((node.index for node in nodes), default=-1) + + return IRModels.Node( + index=max_index + 1, + name=f"{component_name}_output", + node_type="output", + target="output", + args=[[{"node": name} for name in output_names]], + kwargs={}, + users=(), + tensor_output_meta=None, + module_stack=None, + value_kind=FModels.ValueKind.OUTPUT, + attrs={}, + ir_metadata={}, + ) + + +def unique_component_node_name(name: str, used_names: set[str]) -> str: + if name not in used_names: + used_names.add(name) + return name + + suffix = 1 + candidate = f"{name}_{suffix}" + + while candidate in used_names: + suffix += 1 + candidate = f"{name}_{suffix}" + + used_names.add(candidate) + return candidate diff --git a/python/cactus/transpiler/Generator/component_split_plans.py b/python/cactus/transpiler/Generator/component_split_plans.py new file mode 100644 index 000000000..d076b799c --- /dev/null +++ b/python/cactus/transpiler/Generator/component_split_plans.py @@ -0,0 +1,235 @@ +from __future__ import annotations + +from collections.abc import Mapping +from typing import Any + +from .component_split_boundaries import * +from .component_split_builder import extract_component_graph +from .component_split_types import * +from ..IR import models as IRModels + + +def split_component_graphs( + graphs: Mapping[str, IRModels.Graph], + model_profile: Any | None = None, +) -> dict[str, IRModels.Graph] | None: + prefill = find_graph_by_task(graphs, PREFILL_WITH_CACHE_TASK) + decode = find_graph_by_task(graphs, DECODE_WITH_CACHE_TASK) + + if prefill is None or decode is None: + return None + + if is_gemma4_profile(model_profile): + return split_gemma4_components(prefill, decode) + + if is_whisper_profile(model_profile): + return split_whisper_components(prefill, decode) + + if is_causal_lm_profile(model_profile): + return split_causal_lm_components(prefill, decode) + + return None + + +def split_gemma4_components(prefill: IRModels.Graph, decode: IRModels.Graph) -> dict[str, IRModels.Graph]: + boundaries = find_gemma4_boundaries(prefill, decode) + position_step = PlaceholderSpec("position_ids", "position_ids", source_node=boundaries.decode_position_ids) + position_chunk = PlaceholderSpec("position_ids", "position_ids", tensor_node="input_ids") + inputs_embeds_step = PlaceholderSpec("inputs_embeds", "inputs_embeds", source_node=boundaries.decode_text_inputs_embeds) + inputs_embeds_chunk = PlaceholderSpec("inputs_embeds", "inputs_embeds", source_node=boundaries.merged_inputs_embeds) + per_layer_step = PlaceholderSpec("per_layer_inputs", "per_layer_inputs", source_node=boundaries.decode_per_layer_inputs) + per_layer_chunk = PlaceholderSpec("per_layer_inputs", "per_layer_inputs", source_node=boundaries.prefill_per_layer_inputs) + + specs = ( + ComponentSplitSpec( + name="vision_encoder", + graph=prefill, + outputs=(OutputSpec(boundaries.vision_features, "image_features", row_limit=256),), + input_aliases={"image_position_ids": "pixel_position_ids"}, + ), + ComponentSplitSpec( + name="audio_encoder", + graph=prefill, + outputs=(OutputSpec(boundaries.audio_features, "audio_features"),), + ), + ComponentSplitSpec( + name="lm_encoder_step", + graph=decode, + outputs=( + OutputSpec(boundaries.decode_text_inputs_embeds, "inputs_embeds"), + OutputSpec(boundaries.decode_per_layer_inputs, "per_layer_inputs"), + OutputSpec(position_step.name, "position_ids"), + ), + placeholders=(position_step,), + ref_aliases={boundaries.decode_token_ids: "input_ids"}, + ), + ComponentSplitSpec( + name="lm_encoder_text_chunk", + graph=prefill, + outputs=( + OutputSpec(boundaries.prefill_text_inputs_embeds, "inputs_embeds"), + OutputSpec(boundaries.prefill_per_layer_inputs, "per_layer_inputs"), + OutputSpec(position_chunk.name, "position_ids"), + ), + placeholders=(position_chunk,), + ref_aliases={ + boundaries.merged_inputs_embeds: boundaries.prefill_text_inputs_embeds, + boundaries.prefill_token_ids: "input_ids", + }, + chunk_tokens=GEMMA4_PREFILL_CHUNK_TOKENS, + ), + ComponentSplitSpec( + name="lm_encoder_media_chunk", + graph=prefill, + outputs=( + OutputSpec(inputs_embeds_chunk.name, "inputs_embeds"), + OutputSpec(boundaries.prefill_per_layer_inputs, "per_layer_inputs"), + OutputSpec(position_chunk.name, "position_ids"), + ), + placeholders=(inputs_embeds_chunk, position_chunk), + ref_aliases={ + boundaries.merged_inputs_embeds: inputs_embeds_chunk.name, + boundaries.prefill_token_ids: "input_ids", + }, + chunk_tokens=GEMMA4_PREFILL_CHUNK_TOKENS, + ), + ComponentSplitSpec( + name="lm_encoder_media_step", + graph=decode, + outputs=( + OutputSpec(inputs_embeds_step.name, "inputs_embeds"), + OutputSpec(boundaries.decode_per_layer_inputs, "per_layer_inputs"), + OutputSpec(position_step.name, "position_ids"), + ), + placeholders=(inputs_embeds_step, position_step), + ref_aliases={boundaries.decode_token_ids: "input_ids"}, + ), + ComponentSplitSpec( + name="decoder_prefill_chunk", + graph=prefill, + outputs=(OutputSpec(boundaries.prefill_logits, "logits"),), + placeholders=(inputs_embeds_chunk, per_layer_chunk, position_chunk), + chunk_tokens=GEMMA4_PREFILL_CHUNK_TOKENS, + ), + ComponentSplitSpec( + name="decoder_step", + graph=decode, + outputs=(OutputSpec(boundaries.decode_logits, "logits"),), + placeholders=(inputs_embeds_step, per_layer_step, position_step), + ), + ) + + return {spec.name: extract_component_graph(spec) for spec in specs} + + +def is_gemma4_profile(model_profile: Any | None) -> bool: + profile_name = str(getattr(model_profile, "model_profiles", "") or "").lower() + return "gemma4" in profile_name or "gemma_4" in profile_name + + +def is_whisper_profile(model_profile: Any | None) -> bool: + profile_name = str(getattr(model_profile, "model_profiles", "") or "").lower() + return "whisper" in profile_name + + +def is_causal_lm_profile(model_profile: Any | None) -> bool: + return str(getattr(model_profile, "load_strategy", "") or "") == "causal_lm" + + +def split_causal_lm_components(prefill: IRModels.Graph, decode: IRModels.Graph) -> dict[str, IRModels.Graph]: + step_position = PlaceholderSpec("position_ids", "position_ids", tensor_node="input_ids", force=True) + chunk_position = PlaceholderSpec("position_ids", "position_ids", tensor_node="input_ids", force=True) + + specs = ( + ComponentSplitSpec( + name="lm_encoder_step", + graph=decode, + outputs=( + OutputSpec("input_ids", "inputs_embeds"), + OutputSpec(step_position.name, "position_ids"), + ), + placeholders=(step_position,), + ), + ComponentSplitSpec( + name="lm_encoder_text_chunk", + graph=prefill, + outputs=( + OutputSpec("input_ids", "inputs_embeds"), + OutputSpec(chunk_position.name, "position_ids"), + ), + placeholders=(chunk_position,), + chunk_tokens=GENERIC_CAUSAL_PREFILL_CHUNK_TOKENS, + ), + ComponentSplitSpec( + name="decoder_step", + graph=decode, + outputs=graph_output_specs(decode), + placeholders=(step_position,), + input_aliases={"input_ids": "inputs_embeds"}, + ), + ComponentSplitSpec( + name="decoder_prefill_chunk", + graph=prefill, + outputs=graph_output_specs(prefill), + placeholders=(chunk_position,), + input_aliases={"input_ids": "inputs_embeds"}, + chunk_tokens=GENERIC_CAUSAL_PREFILL_CHUNK_TOKENS, + ), + ) + + return {spec.name: extract_component_graph(spec) for spec in specs} + + +def split_whisper_components(prefill: IRModels.Graph, decode: IRModels.Graph) -> dict[str, IRModels.Graph]: + boundaries = find_whisper_boundaries(prefill, decode) + position_index = find_whisper_decode_position_index(decode) + hidden = PlaceholderSpec( + "encoder_hidden_states", + "encoder_hidden_states", + source_node=boundaries.encoder_hidden_states, + ) + position = PlaceholderSpec( + "position_ids", + "position_ids", + tensor_node=position_index, + ) + route_metadata = {"runtime_route": "encoder_cross_kv_decoder_step"} + decoder_input_aliases = whisper_decoder_cross_input_aliases(decode, len(boundaries.cross_key_values)) + + specs = ( + ComponentSplitSpec( + name="audio_encoder", + graph=prefill, + outputs=(OutputSpec(boundaries.encoder_hidden_states, "encoder_hidden_states"),), + metadata={ + **route_metadata, + "runtime_role": "source_encoder", + "source_kind": "audio_features", + }, + ), + ComponentSplitSpec( + name="decoder_cross_kv", + graph=prefill, + outputs=tuple( + output + for layer_index, (key_node, value_node) in enumerate(boundaries.cross_key_values) + for output in ( + OutputSpec(key_node, f"cross_k_{layer_index}"), + OutputSpec(value_node, f"cross_v_{layer_index}"), + ) + ), + placeholders=(hidden,), + metadata={**route_metadata, "runtime_role": "decoder_cross_kv"}, + ), + ComponentSplitSpec( + name="decoder_step", + graph=decode, + outputs=(OutputSpec(boundaries.decode_logits, "logits"),), + placeholders=(position,), + ref_aliases={position_index: "position_ids"}, + input_aliases=decoder_input_aliases, + metadata={**route_metadata, "runtime_role": "decoder_step"}, + ), + ) + + return {spec.name: extract_component_graph(spec) for spec in specs} diff --git a/python/cactus/transpiler/Generator/component_split_types.py b/python/cactus/transpiler/Generator/component_split_types.py new file mode 100644 index 000000000..a4d7f73cd --- /dev/null +++ b/python/cactus/transpiler/Generator/component_split_types.py @@ -0,0 +1,66 @@ +from __future__ import annotations + +from collections.abc import Mapping +from dataclasses import dataclass, field + +from ..Fusions import models as FModels +from ..IR import models as IRModels + + +PREFILL_WITH_CACHE_TASK = "prefill_with_cache" +DECODE_WITH_CACHE_TASK = "decode_with_cache" +GEMMA4_PREFILL_CHUNK_TOKENS = 128 +GENERIC_CAUSAL_PREFILL_CHUNK_TOKENS = 4 + + +@dataclass(slots=True, frozen=True) +class PlaceholderSpec: + name: str + logical_name: str + source_node: str | None = None + tensor_node: str | None = None + target: str | None = None + value_kind: str = FModels.ValueKind.USER_INPUT + force: bool = False + + +@dataclass(slots=True, frozen=True) +class OutputSpec: + node: str + logical_name: str + row_limit: int | None = None + + +@dataclass(slots=True, frozen=True) +class ComponentSplitSpec: + name: str + graph: IRModels.Graph + outputs: tuple[OutputSpec, ...] + placeholders: tuple[PlaceholderSpec, ...] = () + ref_aliases: Mapping[str, str] = field(default_factory=dict) + input_aliases: Mapping[str, str] = field(default_factory=dict) + metadata: Mapping[str, str] = field(default_factory=dict) + chunk_tokens: int | None = None + + +@dataclass(slots=True, frozen=True) +class Gemma4Boundaries: + vision_features: str + audio_features: str + merged_inputs_embeds: str + prefill_text_inputs_embeds: str + prefill_token_ids: str + prefill_per_layer_inputs: str + decode_text_inputs_embeds: str + decode_token_ids: str + decode_per_layer_inputs: str + decode_position_ids: str + prefill_logits: str + decode_logits: str + + +@dataclass(slots=True, frozen=True) +class WhisperBoundaries: + encoder_hidden_states: str + cross_key_values: tuple[tuple[str, str], ...] + decode_logits: str diff --git a/python/cactus/transpiler/Generator/component_split_utils.py b/python/cactus/transpiler/Generator/component_split_utils.py new file mode 100644 index 000000000..2e31a6a57 --- /dev/null +++ b/python/cactus/transpiler/Generator/component_split_utils.py @@ -0,0 +1,30 @@ +from __future__ import annotations + +from typing import Any + +from ..IR import models as IRModels + + +def element_count(shape: list[Any]) -> int | None: + product = 1 + + for dim in shape: + if not isinstance(dim, int): + return None + + product *= dim + + return product + + +def tensor_shape(node: IRModels.Node) -> list[Any]: + return list(IRModels.tensor_shape(node)) + + +def tensor_rank(node: IRModels.Node) -> int: + return len(tensor_shape(node)) + + +def tensor_last_dim(node: IRModels.Node) -> Any | None: + shape = tensor_shape(node) + return shape[-1] if shape else None diff --git a/python/cactus/transpiler/Generator/constants.py b/python/cactus/transpiler/Generator/constants.py index d20670e0a..8b404f5fe 100644 --- a/python/cactus/transpiler/Generator/constants.py +++ b/python/cactus/transpiler/Generator/constants.py @@ -34,6 +34,7 @@ FP16_RUNTIME_INPUTS = { "input_features", + "encoder_hidden_states", "inputs_embeds", "pixel_values", } @@ -210,6 +211,10 @@ "aten.flatten.using_ints", } +REPEAT_TARGETS = { + "aten.repeat.default", +} + TRANSPOSE_TARGETS = { "cactus.transpose", "aten.t.default", diff --git a/python/cactus/transpiler/Generator/errors.py b/python/cactus/transpiler/Generator/errors.py new file mode 100644 index 000000000..1f4b44f94 --- /dev/null +++ b/python/cactus/transpiler/Generator/errors.py @@ -0,0 +1,2 @@ +class UnsupportedLoweringError(NotImplementedError): + pass diff --git a/python/cactus/transpiler/Generator/generate.py b/python/cactus/transpiler/Generator/generate.py index 67a7ea35f..208d259e4 100644 --- a/python/cactus/transpiler/Generator/generate.py +++ b/python/cactus/transpiler/Generator/generate.py @@ -1,32 +1,27 @@ from __future__ import annotations import json -import math import shutil -import struct +import tempfile from collections.abc import Mapping from dataclasses import replace from pathlib import Path from typing import Any -import numpy as np from . import component_split from . import constants from . import models +from . import lowerings as lowering_engine from ..Converter import models as CModels -from ..Fusions import models as FModels from ..IR import models as IRModels from ..RuntimePlan import models as RPModels -class UnsupportedLoweringError(NotImplementedError): - pass GraphInput = CModels.LayerMap | IRModels.Graph ComponentInput = GraphInput | Mapping[str, GraphInput] -SIZE_T_MAX = (1 << 64) - 1 def generate( @@ -55,10 +50,10 @@ def generate( ) prepare_generation_output_dir(generator_config.output_dir) components = component_graphs_from_input(ir_output, generator_config, component_name, model_profile) - lowering_rules = build_lowering_rules(lowerings) + lowering_rules = lowering_engine.build_lowering_rules(lowerings) for component in components: - lower_component(component, generator_config, lowering_rules) + lowering_engine.lower_component(component, generator_config, lowering_rules) return models.GenerationResult.from_components(components) @@ -137,6 +132,8 @@ def materialize_runtime_bundle_files(bundle_dir: Path, weights_dir: Path | None, for source in model_profile_metadata_files(model_profile): materialize_bundle_file(source, bundle_dir / source.name, overwrite=False) + ensure_tokenizer_sidecars(bundle_dir) + def model_profile_metadata_files(model_profile: Any | None) -> tuple[Path, ...]: if model_profile is None: @@ -193,6 +190,93 @@ def materialize_bundle_file(source: Path, target: Path, *, overwrite: bool = Tru shutil.copy2(source, target) +def ensure_tokenizer_sidecars(bundle_dir: Path) -> None: + if (bundle_dir / "vocab.txt").is_file(): + return + + if not (bundle_dir / "tokenizer.json").is_file(): + return + + try: + from transformers import AutoTokenizer + from cactus.convert.cactus_adapters import tokenizer as tokenizer_utils + except Exception: + return + + try: + with tempfile.TemporaryDirectory(prefix="cactus_tokenizer_") as temp_dir_name: + temp_dir = Path(temp_dir_name) + + for filename in tokenizer_source_filenames(): + source = bundle_dir / filename + + if source.is_file(): + shutil.copy2(source, temp_dir / filename) + + tokenizer = AutoTokenizer.from_pretrained(str(temp_dir), local_files_only=True, trust_remote_code=True) + original_hf_download = tokenizer_utils.hf_hub_download + tokenizer_utils.hf_hub_download = None + + try: + tokenizer_utils.convert_hf_tokenizer( + tokenizer, + temp_dir, + model_id=str(bundle_dir), + model_type=bundle_model_type(bundle_dir), + ) + finally: + tokenizer_utils.hf_hub_download = original_hf_download + + for filename in tokenizer_sidecar_filenames(): + source = temp_dir / filename + + if source.is_file(): + shutil.copy2(source, bundle_dir / filename) + except Exception: + return + + +def tokenizer_source_filenames() -> tuple[str, ...]: + return ( + "config.json", + "tokenizer.json", + "tokenizer_config.json", + "special_tokens_map.json", + "added_tokens.json", + "chat_template.jinja", + ) + + +def tokenizer_sidecar_filenames() -> tuple[str, ...]: + return ( + "vocab.txt", + "merges.txt", + "special_tokens.json", + "tokenizer_config.txt", + "chat_template.jinja2", + ) + + +def bundle_model_type(bundle_dir: Path) -> str | None: + for filename in ("config.json", "hf_config.json"): + config_path = bundle_dir / filename + + if not config_path.is_file(): + continue + + try: + config = json.loads(config_path.read_text(encoding="utf-8")) + except json.JSONDecodeError: + continue + + model_type = config.get("model_type") + + if model_type is not None: + return str(model_type) + + return None + + def prepare_generation_output_dir(output_dir: Path) -> None: output_dir.mkdir(parents=True, exist_ok=True) constants_dir = output_dir / "constants" @@ -251,2966 +335,3 @@ def default_component_name(graph: IRModels.Graph, model_profile: Any | None) -> return f"{model_name}_{task}" return model_name - - -def create_cactus_graph() -> Any: - try: - from cactus import Graph - except ModuleNotFoundError: - from bindings.cactus import Graph - - return Graph() - - -def component_weight_resolver( - component: models.ComponentGraph, - config: models.GeneratorConfig, -) -> models.WeightResolver | None: - if config.weights_dir is None: - return None - - resolver = models.WeightResolver.from_graph( - component.ir_graph, - config.weights_dir, - config.weights_manifest_path, - ) - - if not resolver.records: - component.warn(f"{component.name}: no converted weight records found in {config.weights_dir}") - - return resolver - - -def build_lowering_rules( - extra_lowerings: Mapping[str, models.LoweringRule | models.LoweringFn] | None = None, -) -> dict[str, models.LoweringRule]: - rules: dict[str, models.LoweringRule] = {} - - add_rules(rules, constants.BINARY_TARGETS, lower_binary) - add_rules(rules, constants.SCALAR_TARGETS, lower_scalar) - add_rules(rules, constants.UNARY_TARGETS, lower_unary) - add_rules(rules, constants.LOG1P_TARGETS, lower_log1p) - add_rules(rules, constants.POW_TARGETS, lower_pow) - add_rules(rules, constants.REDUCE_TARGETS, lower_reduce) - add_rules(rules, constants.SHAPE_TARGETS, lower_shape) - add_rules(rules, constants.EXPAND_TARGETS, lower_expand) - add_rules(rules, constants.UNSQUEEZE_TARGETS, lower_unsqueeze) - add_rules(rules, constants.FLATTEN_TARGETS, lower_flatten) - add_rules(rules, constants.TRANSPOSE_TARGETS, lower_transpose) - add_rules(rules, constants.SLICE_TARGETS, lower_slice) - add_rules(rules, constants.INDEX_TARGETS, lower_index) - add_rules(rules, constants.WHERE_TARGETS, lower_where) - add_rules(rules, constants.UNFOLD_TARGETS, lower_unfold) - add_rules(rules, constants.CAT_TARGETS, lower_cat) - add_rules(rules, constants.MATMUL_TARGETS, lower_matmul) - add_rules(rules, constants.ADDM_CONST_TARGETS, lower_addmm) - add_rules(rules, constants.SPLIT_TARGETS, lower_split) - add_rules(rules, constants.GETITEM_TARGETS, lower_getitem) - add_rules(rules, constants.TO_COPY_TARGETS, lower_to_copy) - add_rules(rules, constants.PASS_THROUGH_TARGETS, lower_pass_through) - add_rules(rules, constants.COPY_TARGETS, lower_copy) - add_rules(rules, constants.PAD_TARGETS, lower_constant_pad_nd) - add_rules(rules, constants.CONSTANT_INPUT_TARGETS, lower_constant_input) - add_rules(rules, constants.NEG_TARGETS, lower_neg) - add_rules(rules, constants.CUMSUM_TARGETS, lower_cumsum) - add_rules(rules, constants.SOFTMAX_TARGETS, lower_softmax) - add_rules(rules, constants.TOPK_TARGETS, lower_topk) - add_rules(rules, constants.GATHER_TARGETS, lower_gather) - add_rules(rules, constants.EMBEDDING_TARGETS, lower_embedding) - add_rules(rules, constants.CLAMP_TARGETS, lower_clamp) - add_rules(rules, constants.NORM_TARGETS, lower_norm) - add_rules(rules, constants.CONV_TARGETS, lower_conv) - add_rules(rules, constants.ATTENTION_TARGETS, lower_attention) - add_rules(rules, constants.CACHE_TARGETS, lower_cache) - add_rules(rules, constants.MOE_TARGETS, lower_moe) - add_rules(rules, constants.SPECIAL_CACTUS_TARGETS, lower_special_cactus) - add_rules(rules, constants.UNSUPPORTED_SEMANTIC_TARGETS, lower_unsupported_semantic) - - if extra_lowerings: - for target, rule in extra_lowerings.items(): - if isinstance(rule, models.LoweringRule): - rules[target] = rule - else: - rules[target] = models.LoweringRule(target, rule) - - return rules - - -def add_rules(rules: dict[str, models.LoweringRule], targets: Any, lower: models.LoweringFn) -> None: - target_names = targets.keys() if isinstance(targets, dict) else targets - - for target in target_names: - rules[target] = models.LoweringRule(target, lower) - - -def cache_attention_generation_plan(graph: IRModels.Graph) -> tuple[frozenset[str], frozenset[str], dict[str, IRModels.CacheAnnotation]]: - skip_names: set[str] = set() - cache_state_names: set[str] = set() - prefill_annotations = prefill_cache_cat_annotations(graph) - - for node in graph.nodes: - if node.target != "cactus.attention" or len(node.parents) < 3: - continue - - key_cat = find_cache_cat_ancestor(node.parents[1], FModels.CacheTensorRole.KEY) - value_cat = find_cache_cat_ancestor(node.parents[2], FModels.CacheTensorRole.VALUE) - - if key_cat is None or value_cat is None: - continue - - key_cache = key_cat.parents[0] - value_cache = value_cat.parents[0] - - if key_cache.cache is None or value_cache.cache is None: - continue - - if key_cache.cache.layer_index != value_cache.cache.layer_index: - continue - - skip_names.update(cache_wrapper_path_names(node.parents[1], key_cat)) - skip_names.update(cache_wrapper_path_names(node.parents[2], value_cat)) - cache_state_names.update((key_cache.name, value_cache.name)) - - for cat_name in prefill_annotations: - cat_node = graph.nodes_map.get(cat_name) - - if cat_node is not None and cat_node.parents: - skip_names.add(cat_node.parents[0].name) - - return frozenset(skip_names), frozenset(cache_state_names), prefill_annotations - - -def prefill_cache_cat_annotations(graph: IRModels.Graph) -> dict[str, IRModels.CacheAnnotation]: - annotations: dict[str, IRModels.CacheAnnotation] = {} - candidates = [node for node in graph.nodes if is_empty_cache_cat(node)] - - for index, node in enumerate(candidates): - shape = concrete_shape(meta_shape(node)) - - if shape is None or len(shape) != 4: - continue - - _, num_kv_heads, sequence_length, head_dim = shape - cache_capacity = int(node.ir_metadata.get("prefill_cache_capacity", sequence_length)) - layer_index = index // 2 - annotations[node.name] = IRModels.CacheAnnotation( - kind=FModels.CacheKind.KV, - role=FModels.CacheTensorRole.KEY if index % 2 == 0 else FModels.CacheTensorRole.VALUE, - tensor_index=index, - layer_index=layer_index, - shape=tuple(shape), - layout="batch_heads_sequence_head_dim", - sequence_length=cache_capacity, - window_size=prefill_cache_window_size(graph, layer_index), - num_kv_heads=int(num_kv_heads), - head_dim=int(head_dim), - source="prefill_empty_cache_cat", - ) - - return annotations - - -def prefill_cache_window_size(graph: IRModels.Graph, layer_index: int) -> int | None: - model_name = str(getattr(graph, "model_name", "")).lower() - - if "gemma" not in model_name: - return None - - return 0 if layer_index % 5 == 4 else 512 - - -def is_empty_cache_cat(node: IRModels.Node) -> bool: - if node.target not in {"cactus.cat", "aten.cat.default"} or len(node.parents) < 2: - return False - - first_parent = node.parents[0] - return first_parent.value_kind == FModels.ValueKind.LIFTED_CONSTANT and element_count(meta_shape(first_parent)) == 0 - - -def find_cache_cat_ancestor(node: IRModels.Node, role: str, max_depth: int = 16) -> IRModels.Node | None: - current = node - - for _ in range(max_depth): - if current.target in {"cactus.cat", "aten.cat.default"} and len(current.parents) >= 2: - cache_parent = current.parents[0] - - if cache_parent.cache is not None and cache_parent.cache.kind == FModels.CacheKind.KV and cache_parent.cache.role == role: - return current - - if len(current.parents) != 1: - return None - - current = current.parents[0] - - return None - - -def cache_wrapper_path_names(start: IRModels.Node, stop: IRModels.Node) -> set[str]: - names: set[str] = set() - current = start - - while True: - names.add(current.name) - - if current is stop: - return names - - if len(current.parents) != 1: - return set() - - current = current.parents[0] - - -def lower_component( - component: models.ComponentGraph, - config: models.GeneratorConfig, - lowering_rules: dict[str, models.LoweringRule], -) -> models.ComponentGraph: - component.graph = create_cactus_graph() - component.weight_resolver = component_weight_resolver(component, config) - context = models.GenerationContext(component=component, config=config, lowerings=lowering_rules) - ( - context.skip_node_names, - context.cache_state_placeholder_names, - context.prefill_cache_cat_annotations, - ) = cache_attention_generation_plan(component.ir_graph) - - for node in IRModels.topological_sort(component.ir_graph): - try: - lower_node(context, node) - except UnsupportedLoweringError as e: - component.mark_unsupported(node) - component.warn(str(e)) - - if config.strict and not config.allow_unsupported_ops: - raise - except KeyError as e: - component.mark_unsupported(node) - component.warn(f"{node.name}: missing lowered input {e}") - - if config.strict and not config.allow_unsupported_ops: - raise UnsupportedLoweringError(f"{node.name}: missing lowered input {e}") from e - - component.save() - return component - - -def lower_node(context: models.GenerationContext, node: IRModels.Node) -> None: - if node.name in context.skip_node_names: - return - - if node.is_placeholder: - context.bind(node, lower_placeholder(context, node)) - return - - if node.is_output: - context.bind(node, lower_output(context, node)) - return - - rule = context.lowering_for(node) - - if rule is None: - raise UnsupportedLoweringError(f"{node.name}: no lowering rule for {node.target}") - - result = rule.lower(context, node) - - if result is not None: - context.bind(node, result) - - -def lower_placeholder(context: models.GenerationContext, node: IRModels.Node) -> Any: - if should_lower_cache_placeholder_as_state(context, node): - return lower_cache_placeholder(context, node) - - shape, dynamic_dims = graph_input_shape(node) - dtype = placeholder_precision(context, node) - tensor = context.graph.input(shape, dtype=dtype, dynamic_dims=dynamic_dims if any(dynamic_dims) else None) - logical_name = logical_input_name(context.component.ir_graph, node) - - if node.value_kind in constants.WEIGHT_VALUE_KINDS: - bind_weight_placeholder(context, node, tensor, logical_name) - return tensor - - context.component.add_runtime_input(tensor, logical_name) - return tensor - - -def placeholder_precision(context: models.GenerationContext, node: IRModels.Node) -> int: - if node.value_kind in constants.WEIGHT_VALUE_KINDS and context.component.weight_resolver is not None: - record = context.component.weight_resolver.resolve(node.name) - if record is not None and record.precision and hasattr(context.graph, record.precision): - return int(getattr(context.graph, record.precision)) - - logical_name = logical_input_name(context.component.ir_graph, node) - if logical_name in constants.FP16_RUNTIME_INPUTS and str(tensor_dtype(node)) in {"torch.float32", "torch.float"}: - return int(context.graph.FP16) - - return cactus_precision(context.graph, tensor_dtype(node)) - - -def should_lower_cache_placeholder_as_state(context: models.GenerationContext, node: IRModels.Node) -> bool: - if node.cache is None: - return False - - if node.name in context.cache_state_placeholder_names: - return True - - if node.cache.kind == FModels.CacheKind.CONV: - return any(child.target in {"cactus.conv_cache_append", "cactus.conv_cache_initialize"} for child in node.children) - - if node.cache.kind == FModels.CacheKind.KV: - if any(child.target in {"cactus.cat", "aten.cat.default"} for child in node.children): - return True - - return bool(node.children) and all(child.target in {"cactus.kv_cache_append", "cactus.attention_cached"} for child in node.children) - - return False - - -def lower_cache_placeholder(context: models.GenerationContext, node: IRModels.Node) -> Any: - annotation = require_cache_annotation(node) - - if annotation.kind == FModels.CacheKind.KV: - cache_state = context.graph.kv_cache_state( - kv_cache_capacity(context, annotation), - required_cache_int(annotation.num_kv_heads, node, "num_kv_heads"), - required_cache_int(annotation.head_dim, node, "head_dim"), - window_size=int(annotation.window_size or 0), - sink_size=0, - num_slots=1, - ) - record_cache_state_binding(context, node, cache_state, annotation) - return cache_state - - if annotation.kind == FModels.CacheKind.CONV: - cache_state = context.graph.conv_cache_state( - required_cache_int(annotation.window_size, node, "window_size"), - required_cache_int(annotation.hidden_dim, node, "hidden_dim"), - ) - record_cache_state_binding(context, node, cache_state, annotation) - return cache_state - - raise UnsupportedLoweringError(f"{node.name}: unsupported cache placeholder kind {annotation.kind}") - - -def bind_weight_placeholder(context: models.GenerationContext, node: IRModels.Node, tensor: Any, logical_name: str | None = None) -> None: - resolver = context.component.weight_resolver - - if resolver is None: - context.component.warn(f"{node.name}: weight lowered as graph input because no weights_dir was provided") - context.component.add_runtime_input(tensor, logical_name) - return - - source_target = resolver.source_target_for(node.name) or node.target - record = resolver.resolve(node.name) - - if record is None or record.output_name is None: - message = f"{node.name}: could not resolve converted weight for source target {source_target}" - - if context.config.strict and not context.config.allow_unsupported_ops: - raise UnsupportedLoweringError(message) - - context.component.warn(message) - context.component.add_runtime_input(tensor, logical_name) - return - - node_id = models.tensor_node_id(tensor) - - if node_id is None: - message = f"{node.name}: lowered weight tensor does not expose a Cactus node id" - - if context.config.strict and not context.config.allow_unsupported_ops: - raise UnsupportedLoweringError(message) - - context.component.warn(message) - return - - context.component.add_weight_binding( - models.WeightBinding( - placeholder=node.name, - source_target=source_target, - node_id=node_id, - path=record.output_name, - output_name=record.output_name, - source_name=record.source_name or record.hf_name or source_target, - value_id=node.name, - precision=record.precision, - component=record.component, - scale_factor=record.scale_factor, - adapter_family=record.adapter_family, - transform=record.transform, - qdq_restore=record.qdq_restore, - ) - ) - - -def lower_output(context: models.GenerationContext, node: IRModels.Node) -> Any: - refs = IRModels.extract_node_refs((node.args, node.kwargs)) - values = tuple(context.require(ref) for ref in refs) - output_index = 0 - - for ref, value in zip(refs, values): - if isinstance(value, (tuple, list)): - for item in value: - context.component.add_output(item, logical_output_name(context.component.ir_graph, output_index, ref)) - output_index += 1 - else: - context.component.add_output(value, logical_output_name(context.component.ir_graph, output_index, ref)) - output_index += 1 - - if len(values) == 1: - return values[0] - - return values - - -def logical_input_name(graph: IRModels.Graph, node: IRModels.Node) -> str: - metadata_name = logical_name_from_metadata(node.ir_metadata, ("logical_input", "logical_name", "runtime_input")) - - if metadata_name: - return metadata_name - - spec_name = logical_input_name_from_specs(graph, node.name) - - if spec_name: - return spec_name - - if node.cache is not None: - return cache_logical_name(node) - - target_name = clean_logical_name(node.target) - - if target_name is not None: - return target_name - - return node.name - - -def logical_output_name(graph: IRModels.Graph, output_index: int, source_ref: str | None = None) -> str: - source = graph.nodes_map.get(source_ref) if source_ref is not None else None - - if source is not None: - metadata_name = logical_name_from_metadata(source.ir_metadata, ("logical_output", "logical_name", "runtime_output")) - - if metadata_name: - return metadata_name - - if source.cache is not None: - return cache_logical_name(source) - - if "logits" in source.name or "lm_head" in source.target: - return "logits" - - spec_name = logical_output_name_from_specs(graph, output_index) - - if spec_name: - return spec_name - - return f"output_{output_index}" - - -def logical_input_name_from_specs(graph: IRModels.Graph, node_name: str) -> str | None: - for spec in graph.input_specs: - if spec.arg_name != node_name: - continue - - return clean_logical_name(spec.target) or clean_logical_name(spec.arg_name) - - return None - - -def logical_output_name_from_specs(graph: IRModels.Graph, output_index: int) -> str | None: - if output_index < 0 or output_index >= len(graph.output_specs): - return None - - spec = graph.output_specs[output_index] - return clean_logical_name(spec.target) or clean_logical_name(spec.arg_name) - - -def logical_name_from_metadata(metadata: dict[str, Any], keys: tuple[str, ...]) -> str | None: - for key in keys: - value = clean_logical_name(metadata.get(key)) - - if value is not None: - return value - - return None - - -def clean_logical_name(value: Any) -> str | None: - if value is None: - return None - - name = str(value) - - if not name or name == "None": - return None - - return name - - -def cache_logical_name(node: IRModels.Node) -> str: - if node.cache is not None and node.cache.tensor_index is not None: - return f"past_key_values_{node.cache.tensor_index}" - - return node.name - - -def lower_binary(context: models.GenerationContext, node: IRModels.Node) -> Any: - inputs = context.inputs_for(node) - method = constants.BINARY_TARGETS[node.target] - - if len(inputs) == 2: - left, right = align_binary_inputs(context, node, inputs) - - if method == "add" and context.component.name in constants.GEMMA_ADD_CLIPPED_COMPONENTS and looks_like_gemma_residual_add(node): - return context.graph.add_clipped(left, right) - - if method == "multiply" and looks_like_gemma4_mlp_product(node): - if is_gemma4_mlp_gate_activation(node.parents[0]): - left = context.graph.scalar_multiply(left, constants.GEMMA4_MLP_PRODUCT_SCALE) - else: - right = context.graph.scalar_multiply(right, constants.GEMMA4_MLP_PRODUCT_SCALE) - - return getattr(context.graph, method)(left, right) - - scalar = scalar_attr(node, "other") - - if len(inputs) == 1 and scalar is not None: - scalar_method = binary_method_to_scalar_method(method) - return getattr(context.graph, scalar_method)(inputs[0], scalar) - - raise unsupported_arity(node, len(inputs), "two tensor inputs or one tensor plus scalar attr") - - -def looks_like_gemma_residual_add(node: IRModels.Node) -> bool: - if len(node.parents) != 2: - return False - - shape = meta_shape(node) - - if not shape or len(shape) > 4: - return False - - left, right = node.parents - return ( - is_gemma_residual_branch(left, right) - or is_gemma_residual_branch(right, left) - ) - - -def is_gemma_residual_branch(residual: IRModels.Node, branch: IRModels.Node) -> bool: - branch_source = strip_precision_passthrough(branch) - - if branch_source.target != "cactus.rms_norm" or not branch_source.parents: - return False - - return strip_precision_passthrough(residual) is not strip_precision_passthrough(branch_source.parents[0]) - - -def looks_like_gemma4_mlp_product(node: IRModels.Node) -> bool: - if len(node.parents) != 2: - return False - - text = node_context_text(node) - - if "gemma" not in text.lower() or "language_model" not in text or ".mlp" not in text: - return False - - left, right = node.parents - return ( - is_gemma4_mlp_gate_activation(left) and has_ancestor_text(right, "mlp_up_proj") - ) or ( - is_gemma4_mlp_gate_activation(right) and has_ancestor_text(left, "mlp_up_proj") - ) - - -def is_gemma4_mlp_gate_activation(node: IRModels.Node) -> bool: - source = strip_precision_passthrough(node) - - if source.target not in {"cactus.gelu", "aten.gelu.default"}: - return False - - return has_ancestor_text(source, "mlp_gate_proj") - - -def has_ancestor_text(node: IRModels.Node, pattern: str, max_depth: int = 8) -> bool: - pattern = pattern.lower() - worklist: list[tuple[IRModels.Node, int]] = [(node, 0)] - seen: set[str] = set() - - while worklist: - current, depth = worklist.pop(0) - - if current.name in seen: - continue - - seen.add(current.name) - - if pattern in node_context_text(current).lower(): - return True - - if depth >= max_depth: - continue - - for parent in current.parents: - worklist.append((parent, depth + 1)) - - return False - - -def node_context_text(node: IRModels.Node) -> str: - return f"{node.name} {node.target} {node.module_stack!r}" - - -def strip_precision_passthrough(node: IRModels.Node) -> IRModels.Node: - current = node - passthrough_targets = constants.PASS_THROUGH_TARGETS | constants.TO_COPY_TARGETS - - while current.target in passthrough_targets and len(current.parents) == 1: - current = current.parents[0] - - return current - - -def lower_scalar(context: models.GenerationContext, node: IRModels.Node) -> Any: - inputs = require_input_count(context, node, 1) - method, attr_name = constants.SCALAR_TARGETS[node.target] - value = scalar_attr(node, attr_name) - - if value is None: - value = scalar_attr(node, "arg_1") - - if value is None: - if method == "scalar_equal": - return context.graph.scalar_multiply(inputs[0], 0.0) - - if method == "scalar_not_equal": - return context.graph.scalar_add(context.graph.scalar_multiply(inputs[0], 0.0), 1.0) - - raise UnsupportedLoweringError(f"{node.name}: scalar lowering {node.target} missing attr {attr_name}") - - if method == "scalar_multiply" and float(value) == 1.0: - return inputs[0] - - return getattr(context.graph, method)(inputs[0], value) - - -def align_binary_inputs(context: models.GenerationContext, node: IRModels.Node, inputs: tuple[Any, ...]) -> tuple[Any, Any]: - target_shape = output_shape(node) - left = align_input_to_shape(context, inputs[0], node.parents[0], target_shape) - right = align_input_to_shape(context, inputs[1], node.parents[1], target_shape) - return align_binary_precision(context, left, right) - - -def align_binary_precision(context: models.GenerationContext, left: Any, right: Any) -> tuple[Any, Any]: - left_precision = getattr(left, "dtype", None) - right_precision = getattr(right, "dtype", None) - - if left_precision is None or right_precision is None or left_precision == right_precision: - return left, right - - return left, cast_to_precision(context, right, left_precision) - - -def align_input_to_shape(context: models.GenerationContext, value: Any, source_node: IRModels.Node, target_shape: tuple[Any, ...]) -> Any: - value = align_value_to_declared_shape(context, value, source_node) - source_shape = meta_shape(source_node) - - if shape_matches_tensor(source_node, target_shape): - return value - - if element_count(source_shape) == element_count(target_shape): - return context.graph.reshape(value, target_shape) - - return value - - -def align_value_to_declared_shape(context: models.GenerationContext, value: Any, source_node: IRModels.Node) -> Any: - actual_shape = tuple(getattr(value, "shape", ())) - declared_shape = meta_shape(source_node) - - if not actual_shape or not declared_shape or len(declared_shape) <= len(actual_shape): - return value - - resolved_shape = resolve_declared_shape_from_actual(declared_shape, actual_shape) - - if resolved_shape is None or tuple(resolved_shape) == actual_shape: - return value - - return context.graph.reshape(value, resolved_shape) - - -def resolve_declared_shape_from_actual(declared_shape: tuple[Any, ...], actual_shape: tuple[int, ...]) -> tuple[int, ...] | None: - actual_index = len(actual_shape) - 1 - resolved = [1 for _ in declared_shape] - - for declared_index in range(len(declared_shape) - 1, -1, -1): - declared_dim = concrete_dim(declared_shape[declared_index]) - - if actual_index >= 0 and (declared_dim is None or declared_dim == actual_shape[actual_index]): - resolved[declared_index] = actual_shape[actual_index] - actual_index -= 1 - continue - - if declared_dim == 1: - resolved[declared_index] = 1 - continue - - return None - - if actual_index >= 0: - return None - - return tuple(resolved) - - -def lower_unary(context: models.GenerationContext, node: IRModels.Node) -> Any: - inputs = require_input_count(context, node, 1) - method = constants.UNARY_TARGETS[node.target] - value = fp16_tensor(context, inputs[0]) if method in constants.FP16_UNARY_METHODS else inputs[0] - return getattr(context.graph, method)(value) - - -def lower_log1p(context: models.GenerationContext, node: IRModels.Node) -> Any: - inputs = require_input_count(context, node, 1) - - if node.parents and node.parents[0].target in {"aten.exp.default", "cactus.scalar_exp"} and node.parents[0].parents: - x = context.require(node.parents[0].parents[0].name) - return lower_stable_softplus(context, x) - - return context.graph.scalar_log(context.graph.scalar_add(inputs[0], 1.0)) - - -def lower_stable_softplus(context: models.GenerationContext, x: Any) -> Any: - x = cast_to_precision(context, x, context.graph.FP16) - neg_abs = context.graph.scalar_multiply(context.graph.abs(x), -1.0) - log_term = context.graph.scalar_log(context.graph.scalar_add(context.graph.scalar_exp(neg_abs), 1.0)) - positive_term = context.graph.clamp(x, 0.0, math.inf) - return context.graph.add(log_term, positive_term) - - -def lower_pow(context: models.GenerationContext, node: IRModels.Node) -> Any: - inputs = require_input_count(context, node, 1) - exponent = scalar_attr(node, "exponent") - - if exponent is None: - raise UnsupportedLoweringError(f"{node.name}: pow lowering missing exponent attr") - - return context.graph.pow(inputs[0], exponent) - - -def lower_reduce(context: models.GenerationContext, node: IRModels.Node) -> Any: - inputs = require_input_count(context, node, 1) - method = constants.REDUCE_TARGETS[node.target] - axis = axis_attr(node) - - if axis is None: - raise UnsupportedLoweringError(f"{node.name}: reduce lowering missing axis/dim attr") - - normalized_axis = normalize_dim(axis, len(meta_shape(node.parents[0]))) - reduced = getattr(context.graph, method)(inputs[0], normalized_axis) - target_shape = output_shape(node) - dropped_shape = reduction_dropped_shape(meta_shape(node.parents[0]), normalized_axis) - - if bool(node.attrs.get("keepdim", False)) or ( - tuple(dropped_shape) != tuple(target_shape) - and element_count(dropped_shape) == element_count(target_shape) - ): - return context.graph.reshape(reduced, target_shape) - - return reduced - - -def lower_shape(context: models.GenerationContext, node: IRModels.Node) -> Any: - inputs = require_input_count(context, node, 1) - method = constants.SHAPE_TARGETS[node.target] - shape = decoder_prefill_logits_shape(context, node, shape_attr(node)) - - return getattr(context.graph, method)(inputs[0], shape) - - -def lower_expand(context: models.GenerationContext, node: IRModels.Node) -> Any: - inputs = require_input_count(context, node, 1) - target_shape = shape_attr(node) - - if shape_matches_tensor(node.parents[0], target_shape): - return inputs[0] - - if element_count(meta_shape(node.parents[0])) == element_count(target_shape): - return context.graph.reshape(inputs[0], target_shape) - - try: - return context.graph.expand(inputs[0], target_shape) - except AttributeError as e: - raise UnsupportedLoweringError( - f"{node.name}: native Cactus expand is unavailable for broadcast shape " - f"{meta_shape(node.parents[0])} -> {target_shape}" - ) from e - - -def lower_unsqueeze(context: models.GenerationContext, node: IRModels.Node) -> Any: - inputs = require_input_count(context, node, 1) - return context.graph.reshape(inputs[0], output_shape(node)) - - -def lower_flatten(context: models.GenerationContext, node: IRModels.Node) -> Any: - inputs = require_input_count(context, node, 1) - return context.graph.flatten( - inputs[0], - start_dim=int(node.attrs.get("start_dim", 0)), - end_dim=int(node.attrs.get("end_dim", -1)), - ) - - -def lower_transpose(context: models.GenerationContext, node: IRModels.Node) -> Any: - inputs = require_input_count(context, node, 1) - - if can_alias_quantized_weight_transpose(context, node, inputs[0]): - return inputs[0] - - permutation = node.attrs.get("permutation") - - if permutation is not None: - return context.graph.permute(inputs[0], tuple_ints(permutation)) - - if node.target == "aten.t.default": - return context.graph.transpose(inputs[0]) - - dim0 = node.attrs.get("dim0") - dim1 = node.attrs.get("dim1") - - if dim0 is None or dim1 is None: - raise UnsupportedLoweringError(f"{node.name}: transpose lowering missing permutation or dim attrs") - - return context.graph.permute(inputs[0], swap_permutation(parent_rank(node), int(dim0), int(dim1))) - - -def can_alias_quantized_weight_transpose(context: models.GenerationContext, node: IRModels.Node, value: Any) -> bool: - return ( - is_weight_transpose_node(node) - and is_cq_tensor(context, value) - and bool(node.children) - and all(child.target in constants.MATMUL_TARGETS for child in node.children) - ) - - -def is_weight_transpose_node(node: IRModels.Node) -> bool: - if not node.parents: - return False - - parent = node.parents[0] - if not parent.is_placeholder or parent.value_kind not in constants.WEIGHT_VALUE_KINDS: - return False - - if node.target == "aten.t.default": - return True - - rank = len(meta_shape(parent)) - if rank < 2: - return False - - permutation = node.attrs.get("permutation") - if permutation is not None: - return tuple_ints(permutation) == swap_permutation(rank, rank - 2, rank - 1) - - dim0 = node.attrs.get("dim0") - dim1 = node.attrs.get("dim1") - if dim0 is None or dim1 is None: - return False - - return { - normalize_dim(int(dim0), rank), - normalize_dim(int(dim1), rank), - } == {rank - 2, rank - 1} - - -def is_cq_tensor(context: models.GenerationContext, value: Any) -> bool: - return getattr(value, "dtype", None) in { - context.graph.CQ1, - context.graph.CQ2, - context.graph.CQ3, - context.graph.CQ4, - } - - -def lower_slice(context: models.GenerationContext, node: IRModels.Node) -> Any: - inputs = require_input_count(context, node, 1) - axis_value = node.attrs.get("axis") - - if axis_value is None: - axis_value = node.attrs.get("dim") - - if axis_value is None: - if shape_matches_tensor(node.parents[0], output_shape(node)): - return inputs[0] - - raise UnsupportedLoweringError(f"{node.name}: slice lowering missing axis/dim attr") - - axis = normalize_dim(int(axis_value), len(meta_shape(node.parents[0]))) - - if node.target == "aten.select.int" or "index_value" in node.attrs: - return context.graph.index(inputs[0], int(node.attrs.get("index_value", node.attrs.get("index", 0))), axis=axis) - - start = int(node.attrs.get("start", 0) or 0) - step = node.attrs.get("step", 1) - - if is_full_axis_slice(node, axis, start, step): - return inputs[0] - - if step is not None and int(step) != 1: - target_shape = output_shape(node) - length = target_shape[axis] if axis < len(target_shape) else slice_length(node, start) - return context.graph.strided_slice(inputs[0], axis, start, length, int(step)) - - length = slice_length(node, start) - return context.graph.slice(inputs[0], axis, start, length=length) - - -def is_full_axis_slice(node: IRModels.Node, axis: int, start: int, step: Any) -> bool: - if start != 0: - return False - - if step is not None and int(step) != 1: - return False - - if not node.parents: - return False - - return shape_matches_tensor(node.parents[0], output_shape(node)) - - -def lower_index(context: models.GenerationContext, node: IRModels.Node) -> Any: - if node.target == "aten.index.Tensor": - return lower_index_tensor(context, node) - - inputs = require_input_count(context, node, 1) - index_value = node.attrs.get("index_value", node.attrs.get("index")) - - if index_value is None: - raise UnsupportedLoweringError(f"{node.name}: index lowering missing index_value/index attr") - - return context.graph.index(inputs[0], int(index_value), axis=int(node.attrs.get("axis", node.attrs.get("dim", 0)))) - - -def lower_index_tensor(context: models.GenerationContext, node: IRModels.Node) -> Any: - inputs = context.inputs_for(node) - - if len(inputs) < 2: - raise unsupported_arity(node, len(inputs), "source tensor plus at least one index tensor") - - source = inputs[0] - source_shape = meta_shape(node.parents[0]) - target_shape = meta_shape(node) - if len(source_shape) == 3 and len(target_shape) == 2 and concrete_dim(source_shape[0]) == 1: - feature_count = concrete_dim(source_shape[1]) - feature_dim = concrete_dim(source_shape[2]) - - if feature_count is not None and feature_dim is not None: - return context.graph.reshape(source, (feature_count, feature_dim)) - - if len(source_shape) == 2 and len(target_shape) == 4 and concrete_dim(source_shape[0]) == 1: - source_length = concrete_dim(source_shape[1]) - target_length = concrete_dim(target_shape[-1]) - - if source_length is not None and target_length is not None: - reshaped = context.graph.reshape(source, (1, 1, 1, source_length)) - - if target_length < source_length: - return context.graph.slice(reshaped, 3, 0, length=target_length) - - if target_length == source_length: - return reshaped - - concrete_target_shape = concrete_shape(target_shape) - - if concrete_target_shape is not None and element_count(source_shape) == element_count(concrete_target_shape): - return context.graph.reshape(source, concrete_target_shape) - - raise UnsupportedLoweringError(f"{node.name}: unsupported aten.index.Tensor pattern") - - -def lower_where(context: models.GenerationContext, node: IRModels.Node) -> Any: - inputs = require_input_count(context, node, 3) - output_precision = cactus_precision(context.graph, tensor_dtype(node)) - true_value = cast_to_precision(context, inputs[1], output_precision) - false_value = cast_to_precision(context, inputs[2], output_precision) - return context.graph.where(inputs[0], true_value, false_value) - - -def lower_unfold(context: models.GenerationContext, node: IRModels.Node) -> Any: - inputs = require_input_count(context, node, 1) - dimension = int(node.attrs.get("dimension", node.attrs.get("dim", node.attrs.get("arg_1", 0)))) - size = int(node.attrs.get("size", node.attrs.get("arg_2", 1))) - step = int(node.attrs.get("step", node.attrs.get("arg_3", 1))) - return context.graph.unfold(inputs[0], dimension, size, step) - - -def lower_cat(context: models.GenerationContext, node: IRModels.Node) -> Any: - if node.name in context.prefill_cache_cat_annotations: - prefill_cache = lower_prefill_cache_cat(context, node) - - if prefill_cache is not None: - return prefill_cache - - decode_cache = lower_decode_cache_cat(context, node) - - if decode_cache is not None: - return decode_cache - - inputs = context.inputs_for(node) - - if not inputs: - raise unsupported_arity(node, 0, "at least one input") - - passthrough = empty_cat_passthrough(node, inputs) - - if passthrough is not None: - return passthrough - - axis = normalize_dim(int(node.attrs.get("axis", node.attrs.get("dim", 0))), cat_rank(node)) - return context.graph.cat(fp16_cat_inputs(context, inputs), axis=axis) - - -def fp16_cat_inputs(context: models.GenerationContext, inputs: tuple[Any, ...]) -> tuple[Any, ...]: - fp16 = int(context.graph.FP16) - return tuple(value if getattr(value, "dtype", None) == fp16 else context.graph.precision_cast(value, fp16) for value in inputs) - - -def lower_decode_cache_cat(context: models.GenerationContext, node: IRModels.Node) -> Any | None: - if len(node.parents) < 2: - return None - - cache_parent = node.parents[0] - - if cache_parent.cache is None or cache_parent.cache.kind != FModels.CacheKind.KV: - return None - - return context.require(cache_parent.name) - - -def lower_prefill_cache_cat(context: models.GenerationContext, node: IRModels.Node) -> Any | None: - annotation = context.prefill_cache_cat_annotations.get(node.name) - - if annotation is None or len(node.parents) < 2: - return None - - new_kv_node = node.parents[1] - original_new_kv = context.require(new_kv_node.name) - new_kv = cache_new_kv_for_native(context, new_kv_node, original_new_kv, annotation.role or "") - new_kv = cast_to_precision(context, new_kv, context.graph.FP16) - cache_state = context.graph.kv_cache_state( - kv_cache_capacity(context, annotation), - required_cache_int(annotation.num_kv_heads, node, "num_kv_heads"), - required_cache_int(annotation.head_dim, node, "head_dim"), - window_size=int(annotation.window_size or 0), - sink_size=0, - num_slots=1, - ) - record_cache_state_binding(context, node, cache_state, annotation) - context.graph.kv_cache_append(new_kv, cache_state, window_size=int(annotation.window_size or 0), sink_size=0) - context.prefill_cache_cat_states[node.name] = cache_state - return original_new_kv - - -def lower_matmul(context: models.GenerationContext, node: IRModels.Node) -> Any: - inputs = require_input_count(context, node, 2) - - if node.target == "aten.bmm.default": - return lower_bmm(context, node, inputs) - - lhs = matmul_activation_operand(context, inputs[0]) - rhs = inputs[1] - lhs = slice_decoder_prefill_logits_lhs(context, node, lhs) - - if len(node.parents) >= 2 and is_weight_transpose_node(node.parents[1]): - original_weight_node = node.parents[1].parents[0] - original_weight = context.lookup(original_weight_node.name) - - if original_weight is not None and is_cq_tensor(context, original_weight): - return context.graph.matmul(lhs, original_weight, pretransposed_rhs=True) - - if not is_cq_tensor(context, rhs): - rhs = matmul_activation_operand(context, rhs) - - return context.graph.matmul( - lhs, - rhs, - pretransposed_rhs=bool(node.attrs.get("pretransposed_rhs", False)), - ) - - -def matmul_activation_operand(context: models.GenerationContext, value: Any) -> Any: - if getattr(value, "dtype", None) == context.graph.FP32: - return cast_to_precision(context, value, context.graph.FP16) - - return value - - -def slice_decoder_prefill_logits_lhs(context: models.GenerationContext, node: IRModels.Node, lhs: Any) -> Any: - if not is_decoder_prefill_logits_matmul(context, node): - return lhs - - lhs_shape = meta_shape(node.parents[0]) if node.parents else () - - if len(lhs_shape) != 2 or not isinstance(lhs_shape[0], int) or lhs_shape[0] <= 1: - return lhs - - return context.graph.slice(lhs, 0, int(lhs_shape[0]) - 1, length=1) - - -def is_decoder_prefill_logits_matmul(context: models.GenerationContext, node: IRModels.Node) -> bool: - shape = meta_shape(node) - - return ( - context.component.name == "decoder_prefill_chunk" - and node.target in constants.MATMUL_TARGETS - and len(shape) == 2 - and len(node.parents) >= 2 - and isinstance(shape[0], int) - and shape[0] > 1 - and isinstance(shape[-1], int) - and shape[-1] >= 100000 - ) - - -def decoder_prefill_logits_shape(context: models.GenerationContext, node: IRModels.Node, shape: tuple[int, ...]) -> tuple[int, ...]: - if ( - context.component.name == "decoder_prefill_chunk" - and len(shape) == 3 - and shape[0] == 1 - and isinstance(shape[1], int) - and shape[1] > 1 - and shape[2] >= 100000 - ): - return (1, 1, shape[2]) - - return shape - - -def empty_cat_passthrough(node: IRModels.Node, inputs: tuple[Any, ...]) -> Any | None: - target_shape = meta_shape(node) - non_empty_inputs = [] - - for input_value, parent in zip(inputs, node.parents): - parent_shape = meta_shape(parent) - parent_elements = element_count(parent_shape) - - if parent_elements == 0: - continue - - non_empty_inputs.append((input_value, parent_shape)) - - if len(non_empty_inputs) != 1: - return None - - input_value, input_shape = non_empty_inputs[0] - - if tuple(input_shape) == tuple(target_shape): - return input_value - - return None - - -def cat_rank(node: IRModels.Node) -> int: - target_shape = meta_shape(node) - - if target_shape: - return len(target_shape) - - return len(meta_shape(node.parents[0])) if node.parents else 0 - - -def lower_bmm(context: models.GenerationContext, node: IRModels.Node, inputs: tuple[Any, ...]) -> Any: - left_shape = meta_shape(node.parents[0]) - right_shape = meta_shape(node.parents[1]) - - if len(left_shape) != 3 or len(right_shape) != 3: - raise UnsupportedLoweringError(f"{node.name}: bmm lowering requires rank-3 inputs") - - if left_shape[0] != right_shape[0]: - raise UnsupportedLoweringError(f"{node.name}: bmm lowering requires matching batch dimensions") - - batch_size = concrete_dim(left_shape[0]) - - if batch_size is None: - raise UnsupportedLoweringError(f"{node.name}: bmm lowering requires a concrete batch size") - - if batch_size == 1: - left_2d = matmul_activation_operand(context, context.graph.reshape(inputs[0], (left_shape[1], left_shape[2]))) - right_2d = matmul_activation_operand(context, context.graph.reshape(inputs[1], (right_shape[1], right_shape[2]))) - product = context.graph.matmul(left_2d, right_2d) - return context.graph.reshape(product, output_shape(node)) - - batch_outputs = [] - - for batch_index in range(batch_size): - left_2d = matmul_activation_operand(context, context.graph.index(inputs[0], batch_index, axis=0)) - right_2d = matmul_activation_operand(context, context.graph.index(inputs[1], batch_index, axis=0)) - product = context.graph.matmul(left_2d, right_2d) - batch_outputs.append(context.graph.reshape(product, (1, left_shape[1], right_shape[2]))) - - return context.graph.cat(batch_outputs, axis=0) - - -def lower_addmm(context: models.GenerationContext, node: IRModels.Node) -> Any: - inputs = require_input_count(context, node, 3) - product = context.graph.matmul( - matmul_activation_operand(context, inputs[1]), - matmul_activation_operand(context, inputs[2]), - ) - return context.graph.add(inputs[0], product) - - -def lower_split(context: models.GenerationContext, node: IRModels.Node) -> Any: - inputs = require_input_count(context, node, 1) - split_sizes = node.attrs.get("split_sizes") - - if not isinstance(split_sizes, list): - raise UnsupportedLoweringError(f"{node.name}: split lowering requires split_sizes list") - - axis = normalize_dim(int(node.attrs.get("axis", node.attrs.get("dim", 0))), len(meta_shape(node.parents[0]))) - offset = 0 - outputs = [] - - for size in split_sizes: - size_int = int(size) - outputs.append(context.graph.slice(inputs[0], axis, offset, length=size_int)) - offset += size_int - - return tuple(outputs) - - -def lower_getitem(context: models.GenerationContext, node: IRModels.Node) -> Any: - inputs = require_input_count(context, node, 1) - index = int(node.attrs.get("index", 0)) - source = inputs[0] - - if isinstance(source, (tuple, list)): - return source[index] - - if index == 0: - return source - - raise UnsupportedLoweringError(f"{node.name}: getitem index {index} requires tuple-producing parent") - - -def lower_to_copy(context: models.GenerationContext, node: IRModels.Node) -> Any: - inputs = require_input_count(context, node, 1) - dtype = node.attrs.get("dtype") - - if dtype is None or dtype == "torch.bool": - return inputs[0] - - if str(dtype) == "torch.float32" and can_skip_float32_copy(node): - return inputs[0] - - precision_name = constants.DTYPE_TO_PRECISION.get(str(dtype)) - - if precision_name is None: - return inputs[0] - - return cast_to_precision(context, inputs[0], cactus_precision(context.graph, str(dtype))) - - -def can_skip_float32_copy(node: IRModels.Node) -> bool: - if not node.children or not node.parents: - return False - - source_dtype = tensor_dtype(node.parents[0]) - - if constants.DTYPE_TO_PRECISION.get(str(source_dtype)) != "FP16": - return False - - safe_consumers = { - "aten.bmm.default", - "cactus.attention", - "cactus.attention_cached", - "cactus.add", - "cactus.dense_mlp_tq_fused", - "cactus.multiply", - "cactus.pow", - "cactus.rope", - "cactus.rms_norm", - "cactus.scalar_multiply", - "cactus.silu", - "cactus.slice", - "cactus.transpose", - "cactus.view", - } - return all(child.target in safe_consumers for child in node.children) - - -def lower_pass_through(context: models.GenerationContext, node: IRModels.Node) -> Any: - inputs = require_at_least_one_input(context, node) - return inputs[0] - - -def lower_copy(context: models.GenerationContext, node: IRModels.Node) -> Any: - inputs = require_at_least_one_input(context, node) - return inputs[-1] - - -def lower_constant_pad_nd(context: models.GenerationContext, node: IRModels.Node) -> Any: - inputs = require_input_count(context, node, 1) - pads = node.attrs.get("pad", node.attrs.get("arg_1")) - value = scalar_attr(node, "value") - - if value is None: - value = scalar_attr(node, "arg_2") - - if value is None: - value = 0.0 - - if not isinstance(pads, (list, tuple)) or len(pads) % 2 != 0: - raise UnsupportedLoweringError(f"{node.name}: constant_pad_nd lowering requires an even pad list") - - result = inputs[0] - parent_shape = meta_shape(node.parents[0]) - current_shape = list(parent_shape) - rank = len(parent_shape) - positive_pads = [0 for _ in pads] - - for pair_index in range(0, len(pads), 2): - left_pad = int(pads[pair_index]) - right_pad = int(pads[pair_index + 1]) - positive_pads[pair_index] = max(left_pad, 0) - positive_pads[pair_index + 1] = max(right_pad, 0) - - if left_pad >= 0 and right_pad >= 0: - continue - - axis = rank - 1 - (pair_index // 2) - - if axis < 0: - raise UnsupportedLoweringError(f"{node.name}: pad rank exceeds input rank") - - axis_size = current_shape[axis] - - if not isinstance(axis_size, int): - raise UnsupportedLoweringError(f"{node.name}: cannot crop symbolic pad dimension {axis_size}") - - start = max(-left_pad, 0) - end_crop = max(-right_pad, 0) - length = axis_size - start - end_crop - - if length < 0: - raise UnsupportedLoweringError(f"{node.name}: pad crop removes more than the full axis") - - result = context.graph.slice(result, axis, start, length=length) - current_shape[axis] = length - - if any(positive_pads): - return context.graph.pad(result, positive_pads, float(value)) - - return result - - -def lower_constant_input(context: models.GenerationContext, node: IRModels.Node) -> Any: - if node.target == "aten.full_like.default": - return lower_full_like(context, node) - - constant = deterministic_constant_values(node) - - if constant is not None: - values, shape = constant - dtype = cactus_precision(context.graph, tensor_dtype(node)) - tensor = context.graph.input(shape, dtype=dtype) - path = write_constant_tensor(context, node, values, shape, dtype) - node_id = models.tensor_node_id(tensor) - - if node_id is None: - raise UnsupportedLoweringError(f"{node.name}: generated constant does not expose a Cactus node id") - - context.component.add_weight_binding( - models.WeightBinding( - placeholder=node.name, - source_target=node.target, - node_id=node_id, - path=path, - output_name=path, - source_name=node.target, - value_id=node.name, - precision=precision_name(context.graph, dtype), - component=context.component.name, - binding_kind="generated_constant", - ) - ) - return tensor - - context.component.warn(f"{node.name}: constant-producing op {node.target} lowered as graph input") - shape, dynamic_dims = graph_input_shape(node) - dtype = cactus_precision(context.graph, tensor_dtype(node)) - tensor = context.graph.input(shape, dtype=dtype, dynamic_dims=dynamic_dims if any(dynamic_dims) else None) - context.component.add_runtime_input(tensor, node.name) - return tensor - - -def lower_full_like(context: models.GenerationContext, node: IRModels.Node) -> Any: - inputs = require_input_count(context, node, 1) - fill_value = scalar_attr(node, "fill_value") - - if fill_value is None: - fill_value = scalar_attr(node, "arg_1") - - if fill_value is None: - fill_value = 0.0 - - output = context.graph.scalar_multiply(inputs[0], 0.0) - fill_float = float(fill_value) - - if fill_float != 0.0: - output = context.graph.scalar_add(output, fill_float) - - return cast_to_precision(context, output, cactus_precision(context.graph, tensor_dtype(node))) - - -def deterministic_constant_values(node: IRModels.Node) -> tuple[list[float], tuple[int, ...]] | None: - shape = concrete_shape(meta_shape(node)) - - if node.target.startswith("aten.arange"): - start = numeric_attr(node, "start", "arg_0", default=0) - end = numeric_attr(node, "end", "arg_1") - step = numeric_attr(node, "step", "arg_2", default=1) - - if end is None or step in {None, 0}: - return None - - values = arange_values(float(start), float(end), float(step)) - return values, (len(values),) - - if node.target == "aten.scalar_tensor.default": - value = numeric_attr(node, "value", "arg_0", default=0) - return [float(value or 0)], (1,) - - if node.target in {"aten.full.default", "aten.full_like.default"}: - if shape is None: - raw_size = node.attrs.get("size", node.attrs.get("arg_0")) - shape = tuple(int(dim) for dim in raw_size) if isinstance(raw_size, (list, tuple)) else None - - if shape is None: - return None - - fill_value = numeric_attr(node, "fill_value", "arg_1", default=0) - count = math.prod(shape) if shape else 1 - return [float(fill_value or 0)] * int(count), tuple(shape or (1,)) - - return None - - -def arange_values(start: float, end: float, step: float) -> list[float]: - values: list[float] = [] - current = start - - if step > 0: - while current < end: - values.append(current) - current += step - else: - while current > end: - values.append(current) - current += step - - return values - - -def numeric_attr(node: IRModels.Node, *names: str, default: Any | None = None) -> Any | None: - for name in names: - value = node.attrs.get(name) - - if value is not None: - return value - - return default - - -def write_constant_tensor( - context: models.GenerationContext, - node: IRModels.Node, - values: list[float], - shape: tuple[int, ...], - precision: int, -) -> str: - constants_dir = context.component.output_path.parent / "constants" - constants_dir.mkdir(parents=True, exist_ok=True) - filename = f"{models.sanitize_component_name(context.component.name)}__{models.sanitize_component_name(node.name)}.weights" - path = constants_dir / filename - data = constant_data_bytes(values, precision) - write_cactus_tensor_file(path, shape, precision, data) - return constant_binding_path(context, path) - - -def constant_data_bytes(values: list[float], precision: int) -> bytes: - if precision == 1: - return b"".join(struct.pack(" None: - alignment = 32 - header_size = 84 - ndim = len(shape) - original_n = shape[0] if shape else 0 - - with path.open("wb") as f: - f.write(struct.pack(" str: - try: - return str(path.relative_to(context.config.output_dir.parent)) - except ValueError: - return str(path) - - -def precision_name(graph: Any, precision: int) -> str: - for name in ("INT8", "FP16", "FP32", "CQ1", "CQ2", "CQ3", "CQ4"): - if int(getattr(graph, name)) == int(precision): - return name - - return str(precision) - - -def lower_neg(context: models.GenerationContext, node: IRModels.Node) -> Any: - inputs = require_input_count(context, node, 1) - return context.graph.scalar_multiply(inputs[0], -1.0) - - -def lower_cumsum(context: models.GenerationContext, node: IRModels.Node) -> Any: - inputs = require_input_count(context, node, 1) - axis = axis_attr(node) - - if axis is None: - raise UnsupportedLoweringError(f"{node.name}: cumsum lowering missing axis/dim attr") - - return context.graph.cumsum(inputs[0], normalize_dim(axis, len(meta_shape(node.parents[0])))) - - -def lower_softmax(context: models.GenerationContext, node: IRModels.Node) -> Any: - inputs = require_input_count(context, node, 1) - axis = axis_attr(node, default=-1) - - if axis is None: - raise UnsupportedLoweringError(f"{node.name}: softmax lowering missing axis/dim attr") - - return context.graph.softmax(fp16_tensor(context, inputs[0]), axis=normalize_dim(axis, len(meta_shape(node.parents[0])))) - - -def fp16_tensor(context: models.GenerationContext, value: Any) -> Any: - fp16 = int(context.graph.FP16) - return value if getattr(value, "dtype", None) == fp16 else context.graph.precision_cast(value, fp16) - - -def lower_topk(context: models.GenerationContext, node: IRModels.Node) -> Any: - inputs = require_input_count(context, node, 1) - k = node.attrs.get("k") - - if k is None: - raise UnsupportedLoweringError(f"{node.name}: topk lowering missing k attr") - - return context.graph.topk(inputs[0], int(k)) - - -def lower_gather(context: models.GenerationContext, node: IRModels.Node) -> Any: - inputs = require_input_count(context, node, 2) - axis = axis_attr(node, default=0) - return context.graph.gather(inputs[0], inputs[1], axis=axis if axis is not None else 0) - - -def lower_embedding(context: models.GenerationContext, node: IRModels.Node) -> Any: - inputs = require_input_count(context, node, 2) - return context.graph.embedding_from_tensor(inputs[0], inputs[1]) - - -def lower_clamp(context: models.GenerationContext, node: IRModels.Node) -> Any: - inputs = require_input_count(context, node, 1) - lo = scalar_weight_bound_value(context, node.parents[1]) if len(node.parents) > 1 else None - hi = scalar_weight_bound_value(context, node.parents[2]) if len(node.parents) > 2 else None - - if lo is None: - lo = inputs[1] if len(inputs) > 1 else node.attrs.get("lo", node.attrs.get("min", node.attrs.get("arg_1"))) - - if hi is None: - hi = inputs[2] if len(inputs) > 2 else node.attrs.get("hi", node.attrs.get("max", node.attrs.get("arg_2"))) - - if lo is None and hi is None: - return inputs[0] - - result = inputs[0] - - if isinstance(lo, (int, float)) and isinstance(hi, (int, float)): - return context.graph.clamp(result, float(lo), float(hi)) - - if lo is not None: - if isinstance(lo, (int, float)): - result = context.graph.clamp(result, float(lo), math.inf) - else: - lo = cast_to_precision(context, lo, getattr(result, "dtype", context.graph.FP16)) - result = context.graph.where(context.graph.greater_equal(result, lo), result, lo) - - if hi is not None: - if isinstance(hi, (int, float)): - result = context.graph.clamp(result, -math.inf, float(hi)) - else: - hi = cast_to_precision(context, hi, getattr(result, "dtype", context.graph.FP16)) - result = context.graph.where(context.graph.less_equal(result, hi), result, hi) - - return result - - -def lower_norm(context: models.GenerationContext, node: IRModels.Node) -> Any: - inputs = context.inputs_for(node) - target = node.target - - if target == "cactus.rms_norm": - if len(inputs) not in {1, 2}: - raise unsupported_arity(node, len(inputs), "activation plus optional weight") - - weight = inputs[1] if len(inputs) == 2 else None - return lower_rms_norm(context, node, inputs[0], weight) - - if target in {"cactus.layernorm", "cactus.layer_norm", "aten.native_layer_norm.default", "aten.layer_norm.default"}: - require_len(node, inputs, 2) - bias = inputs[2] if len(inputs) > 2 else None - return context.graph.layernorm(inputs[0], inputs[1], bias=bias, eps=epsilon_attr(node)) - - if target in {"cactus.groupnorm", "cactus.group_norm", "aten.native_group_norm.default", "aten.group_norm.default"}: - require_len(node, inputs, 3) - return context.graph.groupnorm(inputs[0], inputs[1], inputs[2], node.attrs.get("num_groups", 1), eps=epsilon_attr(node)) - - if target in {"cactus.batchnorm", "cactus.batch_norm", "aten.native_batch_norm.default", "aten.batch_norm.default"}: - require_len(node, inputs, 5) - return context.graph.batchnorm(inputs[0], inputs[1], inputs[2], inputs[3], inputs[4], axis=int(node.attrs.get("axis", 1)), eps=epsilon_attr(node)) - - raise UnsupportedLoweringError(f"{node.name}: unsupported norm target {target}") - - -def lower_rms_norm(context: models.GenerationContext, node: IRModels.Node, x: Any, weight: Any | None = None) -> Any: - x = cast_to_precision(context, x, context.graph.FP16) - declared_input_shape = meta_shape(node.parents[0]) if node.parents else meta_shape(node) - input_shape = concrete_shape(declared_input_shape) - - if input_shape is None: - actual_shape = tuple(getattr(x, "shape", ())) - - if actual_shape and concrete_dim(declared_input_shape[-1] if declared_input_shape else None) is not None: - input_shape = actual_shape - - if input_shape is None: - raise UnsupportedLoweringError(f"{node.name}: cactus.rms_norm requires concrete input shape") - - if weight is None: - weight = rms_norm_unit_weight(context, int(input_shape[-1])) - else: - weight = cast_to_precision(context, weight, context.graph.FP16) - - if len(input_shape) == 2: - return context.graph.rms_norm(x, weight, eps=epsilon_attr(node)) - - if len(input_shape) == 1: - return context.graph.reshape( - context.graph.rms_norm(context.graph.reshape(x, (1, input_shape[0])), weight, eps=epsilon_attr(node)), - input_shape, - ) - - rows = element_count(input_shape[:-1]) - - if rows is None: - raise UnsupportedLoweringError(f"{node.name}: cactus.rms_norm requires concrete leading dimensions") - - normalized = context.graph.rms_norm(context.graph.reshape(x, (rows, input_shape[-1])), weight, eps=epsilon_attr(node)) - return context.graph.reshape(normalized, input_shape) - - -def rms_norm_unit_weight(context: models.GenerationContext, hidden_dim: int) -> Any: - key = f"__embedded_rms_norm_unit_weight_{hidden_dim}" - existing = context.lookup(key) - - if existing is not None: - return existing - - weight = context.graph.input((hidden_dim,), dtype=context.graph.FP16) - context.graph.set_input(weight, np.ones((hidden_dim,), dtype=np.float16)) - context.graph.mark_embedded_input(weight) - context.values[key] = weight - return weight - - -def lower_conv(context: models.GenerationContext, node: IRModels.Node) -> Any: - inputs = context.inputs_for(node) - target = node.target - - if target.startswith(constants.CACTUS_TARGET_PREFIX): - method = target.removeprefix(constants.CACTUS_TARGET_PREFIX) - return lower_named_conv(context, node, method, inputs) - - require_len(node, inputs, 2) - bias = inputs[2] if len(inputs) > 2 else None - rank = output_rank(node) - - if target == "aten.conv1d.default" or rank == 3: - return lower_generic_conv1d(context, node, inputs, bias) - - if target == "aten.conv2d.default" or rank == 4: - specialized = lower_specialized_conv2d(context, node, inputs, bias) - - if specialized is not None: - return specialized - - return context.graph.conv2d( - inputs[0], - inputs[1], - bias=bias, - stride=first_int(node.attrs.get("stride"), 1), - padding=first_int(node.attrs.get("padding"), 0), - dilation=first_int(node.attrs.get("dilation"), 1), - groups=int(node.attrs.get("groups", 1)), - ) - - raise UnsupportedLoweringError(f"{node.name}: cannot infer conv rank for {target}") - - -def lower_named_conv(context: models.GenerationContext, node: IRModels.Node, method: str, inputs: tuple[Any, ...]) -> Any: - require_len(node, inputs, 2) - bias = inputs[2] if len(inputs) > 2 else None - - if method == "conv1d": - return lower_generic_conv1d(context, node, inputs, bias) - - if method == "conv1d_causal": - x = inputs[0] - - if node.attrs.get("layout") == "batch_hidden_sequence": - x = context.graph.permute(x, (0, 2, 1)) - - output = context.graph.conv1d_causal(x, inputs[1], kernel_size=int(node.attrs.get("kernel_size", 0)), dilation=first_int(node.attrs.get("dilation"), 1)) - - if node.attrs.get("layout") == "batch_hidden_sequence": - return context.graph.permute(output, (0, 2, 1)) - - return output - - if method == "conv1d_k3": - require_plain_conv1d_attrs(node, method) - return context.graph.conv1d_k3(inputs[0], inputs[1], stride=first_int(node.attrs.get("stride"), 1)) - - if method == "conv1d_k7s3": - require_len(node, inputs, 3) - return context.graph.conv1d_k7s3(inputs[0], inputs[1], inputs[2]) - - if method in {"conv1d_same_depthwise_k9", "conv1d_pointwise", "conv2d_k3s2p1", "conv2d_depthwise_k3s2p1", "conv2d_pointwise_1x1", "conv2d_k3s1p1"}: - return getattr(context.graph, method)(inputs[0], inputs[1], bias=bias) - - if method == "conv2d": - specialized = lower_specialized_conv2d(context, node, inputs, bias) - - if specialized is not None: - return specialized - - return context.graph.conv2d( - inputs[0], - inputs[1], - bias=bias, - stride=first_int(node.attrs.get("stride"), 1), - padding=first_int(node.attrs.get("padding"), 0), - dilation=first_int(node.attrs.get("dilation"), 1), - groups=int(node.attrs.get("groups", 1)), - ) - - raise UnsupportedLoweringError(f"{node.name}: unsupported Cactus conv method {method}") - - -def lower_generic_conv1d(context: models.GenerationContext, node: IRModels.Node, inputs: tuple[Any, ...], bias: Any | None) -> Any: - stride = first_int(node.attrs.get("stride"), 1) - padding = first_int(node.attrs.get("padding"), 0) - dilation = first_int(node.attrs.get("dilation"), 1) - groups = first_int(node.attrs.get("groups"), 1) - input_shape = meta_shape(node.parents[0]) if node.parents else () - weight_shape = meta_shape(node.parents[1]) if len(node.parents) > 1 else () - - if ( - padding == 0 - and dilation == 1 - and len(input_shape) == 3 - and len(weight_shape) == 3 - and groups == input_shape[1] - and weight_shape[0] == input_shape[1] - and weight_shape[1] == 1 - ): - return context.graph.conv1d_depthwise(inputs[0], inputs[1], bias=bias, stride=stride) - - require_plain_conv1d_attrs(node, "conv1d") - return context.graph.conv1d(inputs[0], inputs[1], bias=bias, stride=stride) - - -def lower_specialized_conv2d(context: models.GenerationContext, node: IRModels.Node, inputs: tuple[Any, ...], bias: Any | None) -> Any | None: - weight_shape = meta_shape(node.parents[1]) if len(node.parents) > 1 else () - stride = tuple_int_values(node.attrs.get("stride"), 1) - padding = tuple_int_values(node.attrs.get("padding"), 0) - dilation = tuple_int_values(node.attrs.get("dilation"), 1) - groups = first_int(node.attrs.get("groups"), 1) - - if len(weight_shape) != 4 or dilation != (1, 1): - return None - - kernel = tuple(int(dim) for dim in weight_shape[2:]) - - if kernel == (3, 3) and stride == (2, 2) and padding == (1, 1): - if groups != 1 and weight_shape[1] == 1: - return context.graph.conv2d_depthwise_k3s2p1(inputs[0], inputs[1], bias=bias) - - return context.graph.conv2d_k3s2p1(inputs[0], inputs[1], bias=bias) - - if kernel == (3, 3) and stride == (1, 1) and padding == (1, 1) and groups == 1: - return context.graph.conv2d_k3s1p1(inputs[0], inputs[1], bias=bias) - - if kernel == (1, 1) and stride == (1, 1) and padding == (0, 0) and groups == 1: - return context.graph.conv2d_pointwise_1x1(inputs[0], inputs[1], bias=bias) - - return None - - -def require_plain_conv1d_attrs(node: IRModels.Node, method: str) -> None: - padding = first_int(node.attrs.get("padding"), 0) - dilation = first_int(node.attrs.get("dilation"), 1) - groups = first_int(node.attrs.get("groups"), 1) - - if padding != 0 or dilation != 1 or groups != 1: - raise UnsupportedLoweringError( - f"{node.name}: {method} lowering only supports padding=0, dilation=1, groups=1; " - f"got padding={padding}, dilation={dilation}, groups={groups}" - ) - - -def lower_attention(context: models.GenerationContext, node: IRModels.Node) -> Any: - if node.target == "cactus.attention" or node.target == "aten.scaled_dot_product_attention.default": - cached = lower_cache_backed_attention(context, node) - - if cached is not None: - return cached - - inputs = context.inputs_for(node) - require_len(node, inputs, 3) - query, key, value, mask = attention_inputs_for_layout(context, node, inputs) - query = cast_to_precision(context, query, context.graph.FP16) - key = cast_to_precision(context, key, context.graph.FP16) - value = cast_to_precision(context, value, context.graph.FP16) - mask = cast_to_precision(context, mask, context.graph.FP16) if mask is not None else None - output = context.graph.attention( - query, - key, - value, - scale=float(node.attrs.get("scale", 1.0)), - is_causal=bool(node.attrs.get("is_causal", True)), - position_offset=int(node.attrs.get("position_offset", 0)), - window_size=int(node.attrs.get("window_size", 0)), - mask=mask, - additive_mask=bool(node.attrs.get("additive_mask", False)), - ) - return attention_output_for_layout(context, node, output) - - if node.target == "cactus.attention_cached": - inputs = context.inputs_for(node) - require_len(node, inputs, 5) - return context.graph.attention_cached( - inputs[0], - inputs[1], - inputs[2], - inputs[3], - inputs[4], - scale=float(node.attrs.get("scale", 1.0)), - position_offset=cached_attention_position_offset(context, node), - window_size=int(node.attrs.get("window_size", 0)), - v_head_dim=int(node.attrs.get("v_head_dim", 0)), - ) - - raise UnsupportedLoweringError(f"{node.name}: unsupported attention target {node.target}") - - -def lower_cache_backed_attention(context: models.GenerationContext, node: IRModels.Node) -> Any | None: - if node.target != "cactus.attention" or len(node.parents) < 3: - return None - - key_cat = find_attention_cache_cat_ancestor(context, node.parents[1], FModels.CacheTensorRole.KEY) - value_cat = find_attention_cache_cat_ancestor(context, node.parents[2], FModels.CacheTensorRole.VALUE) - - if key_cat is None or value_cat is None: - return None - - key_cache = key_cat.parents[0] - value_cache = value_cat.parents[0] - key_new_node = key_cat.parents[1] - value_new_node = value_cat.parents[1] - - query = context.require(node.parents[0].name) - input_layout = str(node.attrs.get("input_layout", "")) - key_new = cache_new_kv_for_native(context, key_new_node, context.require(key_new_node.name), FModels.CacheTensorRole.KEY, input_layout) - value_new = cache_new_kv_for_native(context, value_new_node, context.require(value_new_node.name), FModels.CacheTensorRole.VALUE, input_layout) - key_cache_state = prefill_cache_state_for_cat(context, key_cat) - value_cache_state = prefill_cache_state_for_cat(context, value_cat) - - if key_cache_state is None: - key_cache_state = context.require(key_cache.name) - - if value_cache_state is None: - value_cache_state = context.require(value_cache.name) - - if node.attrs.get("input_layout") == "bhqd_bhds_bhsd": - query = context.graph.permute(query, (0, 2, 1, 3)) - - query = cast_to_precision(context, query, context.graph.FP16) - key_new = cast_to_precision(context, key_new, context.graph.FP16) - value_new = cast_to_precision(context, value_new, context.graph.FP16) - - cache_pair = (key_cat.name, value_cat.name) - prefill_pair_already_appended = key_cat.name in context.prefill_cache_cat_states and value_cat.name in context.prefill_cache_cat_states - - if cache_pair not in context.appended_cache_pairs and not prefill_pair_already_appended: - context.graph.kv_cache_append(key_new, key_cache_state, window_size=int(node.attrs.get("window_size", 0)), sink_size=0) - context.graph.kv_cache_append(value_new, value_cache_state, window_size=int(node.attrs.get("window_size", 0)), sink_size=0) - context.appended_cache_pairs.add(cache_pair) - - context.values.setdefault(key_cat.name, key_cache_state) - context.values.setdefault(value_cat.name, value_cache_state) - - output = context.graph.attention_cached( - query, - key_new, - value_new, - key_cache_state, - value_cache_state, - scale=float(node.attrs.get("scale", 1.0)), - position_offset=cached_attention_position_offset(context, node), - window_size=int(node.attrs.get("window_size", 0)), - v_head_dim=last_dim(value_new_node), - ) - return attention_output_for_layout(context, node, output) - - -def find_attention_cache_cat_ancestor(context: models.GenerationContext, node: IRModels.Node, role: str) -> IRModels.Node | None: - cache_cat = find_cache_cat_ancestor(node, role) - - if cache_cat is not None: - return cache_cat - - return find_prefill_cache_cat_ancestor(context, node) - - -def find_prefill_cache_cat_ancestor(context: models.GenerationContext, node: IRModels.Node, max_depth: int = 16) -> IRModels.Node | None: - current = node - - for _ in range(max_depth): - if current.name in context.prefill_cache_cat_annotations: - return current - - if len(current.parents) != 1: - return None - - current = current.parents[0] - - return None - - -def prefill_cache_state_for_cat(context: models.GenerationContext, node: IRModels.Node) -> Any | None: - return context.prefill_cache_cat_states.get(node.name) - - -def cached_attention_position_offset(context: models.GenerationContext, node: IRModels.Node) -> int: - if context.component.name == "decoder_step": - return SIZE_T_MAX - - return int(node.attrs.get("position_offset", 0)) - - -def cache_new_kv_for_native(context: models.GenerationContext, node: IRModels.Node, value: Any, role: str, input_layout: str = "") -> Any: - if len(meta_shape(node)) == 4: - return context.graph.permute(value, (0, 2, 1, 3)) - - return value - - -def last_dim(node: IRModels.Node) -> int: - shape = concrete_shape(meta_shape(node)) - - if not shape: - return 0 - - return int(shape[-1]) - - -def attention_inputs_for_layout(context: models.GenerationContext, node: IRModels.Node, inputs: tuple[Any, ...]) -> tuple[Any, Any, Any, Any | None]: - mask = inputs[3] if len(inputs) > 3 else None - - if node.attrs.get("input_layout") == "bhqd_bhds_bhsd": - return ( - context.graph.permute(inputs[0], (0, 2, 1, 3)), - context.graph.permute(inputs[1], (0, 3, 1, 2)), - context.graph.permute(inputs[2], (0, 2, 1, 3)), - mask, - ) - - if node.attrs.get("input_layout") == "bqhd_bhds_bhsd": - return ( - inputs[0], - context.graph.permute(inputs[1], (0, 3, 1, 2)), - context.graph.permute(inputs[2], (0, 2, 1, 3)), - mask, - ) - - return inputs[0], inputs[1], inputs[2], mask - - -def attention_output_for_layout(context: models.GenerationContext, node: IRModels.Node, output: Any) -> Any: - output_layout = node.attrs.get("output_layout") - - if output_layout == "bthd_flat": - return context.graph.reshape(output, output_shape(node)) - - if output_layout == "bhqd": - return context.graph.permute(output, (0, 2, 1, 3)) - - return output - - -def lower_cache(context: models.GenerationContext, node: IRModels.Node) -> Any: - inputs = context.inputs_for(node) - target = node.target - - if target == "cactus.kv_cache_append": - if len(inputs) == 1: - cache_state = create_cache_state_for_cache_output(context, node) - context.graph.kv_cache_append(inputs[0], cache_state, window_size=int(node.attrs.get("window_size", 0)), sink_size=int(node.attrs.get("sink_size", 0))) - return cache_state - - require_len(node, inputs, 2) - require_cache_state_parent(node, 1, "cactus.kv_cache_state") - context.graph.kv_cache_append(inputs[0], inputs[1], window_size=int(node.attrs.get("window_size", 0)), sink_size=int(node.attrs.get("sink_size", 0))) - return inputs[1] - - if target == "cactus.conv_cache_append": - require_len(node, inputs, 2) - require_cache_state_parent(node, 1, "cactus.conv_cache_state") - window = context.graph.conv_cache_append(inputs[0], inputs[1]) - return conv_cache_window_to_ir_layout(context, node, window) - - if target == "cactus.conv_cache_initialize": - if len(inputs) == 1: - cache_state = create_cache_state_for_cache_output(context, node) - rows = conv_cache_rows_for_native(context, node, inputs[0]) - context.graph.conv_cache_initialize(rows, cache_state) - return cache_state - - require_len(node, inputs, 2) - rows = conv_cache_rows_for_native(context, node, inputs[0]) - context.graph.conv_cache_initialize(rows, inputs[1]) - return inputs[1] - - if target == "cactus.recurrent_cache_write": - require_len(node, inputs, 2) - return context.graph.recurrent_cache_write(inputs[0], inputs[1]) - - if target == "cactus.kv_cache_state": - return context.graph.kv_cache_state( - required_int_attr(node, "max_seq_len"), - required_int_attr(node, "num_kv_heads"), - required_int_attr(node, "head_dim"), - window_size=int(node.attrs.get("window_size", 0)), - sink_size=int(node.attrs.get("sink_size", 0)), - num_slots=int(node.attrs.get("num_slots", 1)), - ) - - if target == "cactus.conv_cache_state": - return context.graph.conv_cache_state(required_int_attr(node, "window_size"), required_int_attr(node, "hidden_dim")) - - if target == "cactus.recurrent_cache_state": - return context.graph.recurrent_cache_state(shape_attr(node), dtype=cactus_precision(context.graph, tensor_dtype(node))) - - raise UnsupportedLoweringError(f"{node.name}: unsupported cache target {target}") - - -def require_cache_state_parent(node: IRModels.Node, parent_index: int, expected_target: str) -> None: - if parent_index >= len(node.parents): - raise UnsupportedLoweringError(f"{node.name}: {node.target} missing cache parent {parent_index}") - - parent = node.parents[parent_index] - - if parent.target == expected_target: - return - - if expected_target == "cactus.conv_cache_state" and parent.cache is not None and parent.cache.kind == FModels.CacheKind.CONV: - return - - if expected_target == "cactus.kv_cache_state" and parent.cache is not None and parent.cache.kind == FModels.CacheKind.KV: - return - - raise UnsupportedLoweringError( - f"{node.name}: {node.target} requires parent {parent_index} to be {expected_target}; got {parent.target}. " - "Raw HF cache tensors need a cache-state bridge before native cache append lowering." - ) - - -def require_cache_annotation(node: IRModels.Node) -> IRModels.CacheAnnotation: - if node.cache is None: - raise UnsupportedLoweringError(f"{node.name}: cache lowering requires cache annotation metadata") - - return node.cache - - -def required_cache_int(value: int | None, node: IRModels.Node, name: str) -> int: - if value is None: - raise UnsupportedLoweringError(f"{node.name}: cache lowering missing {name}") - - return int(value) - - -def kv_cache_capacity(context: models.GenerationContext, annotation: IRModels.CacheAnnotation) -> int: - capacity = int(annotation.sequence_length or 1) - - if context.component.ir_graph.task == "decode_with_cache": - capacity += 1 - - return max(capacity, 1) - - -def create_cache_state_for_cache_output(context: models.GenerationContext, node: IRModels.Node) -> Any: - annotation = require_cache_annotation(node) - - if annotation.kind == FModels.CacheKind.KV: - cache_state = context.graph.kv_cache_state( - kv_cache_capacity(context, annotation), - required_cache_int(annotation.num_kv_heads, node, "num_kv_heads"), - required_cache_int(annotation.head_dim, node, "head_dim"), - window_size=int(annotation.window_size or 0), - sink_size=0, - num_slots=1, - ) - record_cache_state_binding(context, node, cache_state, annotation) - return cache_state - - if annotation.kind == FModels.CacheKind.CONV: - cache_state = context.graph.conv_cache_state( - required_cache_int(annotation.window_size, node, "window_size"), - required_cache_int(annotation.hidden_dim, node, "hidden_dim"), - ) - record_cache_state_binding(context, node, cache_state, annotation) - return cache_state - - raise UnsupportedLoweringError(f"{node.name}: unsupported cache output kind {annotation.kind}") - - -def record_cache_state_binding(context: models.GenerationContext, node: IRModels.Node, cache_state: Any, annotation: IRModels.CacheAnnotation) -> None: - node_id = models.tensor_node_id(cache_state) - - if node_id is None: - raise UnsupportedLoweringError(f"{node.name}: cache state tensor does not expose a Cactus node id") - - context.component.add_cache_state_binding(RPModels.cache_state_binding_from_annotation(annotation, node_id)) - - -def conv_cache_rows_for_native(context: models.GenerationContext, node: IRModels.Node, rows: Any) -> Any: - rows_shape = meta_shape(node.parents[0]) if node.parents else output_shape(node) - - if len(rows_shape) == 3: - batch, hidden_dim, window_size = rows_shape - - if batch != 1: - raise UnsupportedLoweringError(f"{node.name}: conv cache initialize currently supports batch size 1") - - return context.graph.reshape(context.graph.permute(rows, (0, 2, 1)), (int(window_size), int(hidden_dim))) - - if len(rows_shape) == 2: - return rows - - raise UnsupportedLoweringError(f"{node.name}: conv cache initialize rows must be rank 2 or 3, got shape {rows_shape}") - - -def conv_cache_window_to_ir_layout(context: models.GenerationContext, node: IRModels.Node, window: Any) -> Any: - target_shape = output_shape(node) - - if len(target_shape) != 3: - return window - - batch, hidden_dim, window_size = target_shape - - if batch != 1: - raise UnsupportedLoweringError(f"{node.name}: conv cache append layout bridge currently supports batch size 1") - - return context.graph.reshape(context.graph.transpose(window), (int(batch), int(hidden_dim), int(window_size))) - - -def lower_moe(context: models.GenerationContext, node: IRModels.Node) -> Any: - inputs = context.inputs_for(node) - - if node.target == "cactus.dense_mlp_tq_fused": - require_len(node, inputs, 4) - hidden = cast_to_precision(context, inputs[0], context.graph.FP16) - return context.graph.dense_mlp_tq_fused( - hidden, - inputs[1], - inputs[2], - inputs[3], - product_scale=float(node.attrs.get("product_scale", 1.0)), - ) - - if node.target == "cactus.moe_layer_gated": - packed = lower_packed_lfm_moe_layer_gated(context, node, inputs) - if packed is not None: - return packed - - if len(inputs) < 6: - raise UnsupportedLoweringError( - f"{node.name}: cactus.moe_layer_gated requires hidden, routing_probs, topk_indices, w1, w3, and w2 inputs; " - "packed/grouped MoE lowering still needs a weight-unpacking policy" - ) - - return context.graph.moe_layer_gated( - inputs[0], - inputs[1], - inputs[2], - ensure_tensor_sequence(inputs[3]), - ensure_tensor_sequence(inputs[4]), - ensure_tensor_sequence(inputs[5]), - required_int_attr(node, "num_experts"), - required_int_attr(node, "num_experts_per_tok"), - normalize_routing=bool(node.attrs.get("normalize_routing", True)), - epsilon=float(node.attrs.get("epsilon", 1e-6)), - routed_scaling_factor=float(node.attrs.get("routed_scaling_factor", 1.0)), - ) - - if node.target == "cactus.moe_layer_ungated": - if len(inputs) < 5: - raise unsupported_arity(node, len(inputs), "hidden, routing_probs, topk_indices, w1, and w2") - - return context.graph.moe_layer_ungated( - inputs[0], - inputs[1], - inputs[2], - ensure_tensor_sequence(inputs[3]), - ensure_tensor_sequence(inputs[4]), - required_int_attr(node, "num_experts"), - required_int_attr(node, "num_experts_per_tok"), - normalize_routing=bool(node.attrs.get("normalize_routing", True)), - epsilon=float(node.attrs.get("epsilon", 1e-6)), - routed_scaling_factor=float(node.attrs.get("routed_scaling_factor", 1.0)), - ) - - raise UnsupportedLoweringError(f"{node.name}: unsupported MoE target {node.target}") - - -def lower_packed_lfm_moe_layer_gated( - context: models.GenerationContext, - node: IRModels.Node, - inputs: tuple[Any, ...], -) -> Any | None: - if len(inputs) != 5 or len(node.parents) != 5: - return None - - hidden_node, router_node, bias_node, gate_up_node, down_node = node.parents - gate_up_shape = concrete_shape(meta_shape(gate_up_node)) - down_shape = concrete_shape(meta_shape(down_node)) - num_experts = required_int_attr(node, "num_experts") - top_k = required_int_attr(node, "num_experts_per_tok") - - if gate_up_shape is None or down_shape is None: - return None - - if len(gate_up_shape) != 3 or len(down_shape) != 3: - return None - - if gate_up_shape[0] != num_experts or down_shape[0] != num_experts: - return None - - packed_intermediate, hidden_dim = gate_up_shape[1], gate_up_shape[2] - if packed_intermediate % 2 != 0: - raise UnsupportedLoweringError(f"{node.name}: packed gated MoE gate/up dimension must be even") - - intermediate_dim = packed_intermediate // 2 - - if down_shape[1:] != (hidden_dim, intermediate_dim): - raise UnsupportedLoweringError( - f"{node.name}: expected LFM packed down_proj shape " - f"[experts, hidden, intermediate], got {down_shape}" - ) - - hidden = moe_hidden_2d(context, node, inputs[0], hidden_dim) - router_weight = inputs[1] - expert_bias = cast_to_precision(context, inputs[2], context.graph.FP16) - gate_up_weight = inputs[3] - down_weight = inputs[4] - - router_logits = context.graph.matmul(hidden, router_weight, pretransposed_rhs=True) - router_logits = context.graph.add(router_logits, expert_bias) - routing_probs = context.graph.softmax(cast_to_precision(context, router_logits, context.graph.FP16), axis=-1) - topk_indices = context.graph.index(context.graph.topk(routing_probs, top_k), 0, axis=0) - w1_weights, w3_weights, w2_weights = split_lfm_packed_moe_weights( - context, - gate_up_weight, - down_weight, - num_experts, - intermediate_dim, - hidden_dim, - ) - - return context.graph.moe_layer_gated( - hidden, - routing_probs, - topk_indices, - w1_weights, - w3_weights, - w2_weights, - num_experts, - top_k, - normalize_routing=bool(node.attrs.get("normalize_routing", True)), - epsilon=float(node.attrs.get("epsilon", 1e-6)), - routed_scaling_factor=float(node.attrs.get("routed_scaling_factor", 1.0)), - ) - - -def moe_hidden_2d(context: models.GenerationContext, node: IRModels.Node, hidden: Any, hidden_dim: int) -> Any: - hidden_shape = concrete_shape(meta_shape(node.parents[0])) if node.parents else None - - if hidden_shape is None: - return cast_to_precision(context, hidden, context.graph.FP16) - - if len(hidden_shape) == 2: - return cast_to_precision(context, hidden, context.graph.FP16) - - if len(hidden_shape) == 3 and hidden_shape[0] == 1 and hidden_shape[2] == hidden_dim: - return cast_to_precision(context, context.graph.reshape(hidden, (hidden_shape[1], hidden_dim)), context.graph.FP16) - - raise UnsupportedLoweringError(f"{node.name}: packed MoE hidden input must be [tokens, hidden] or [1, tokens, hidden]") - - -def split_lfm_packed_moe_weights( - context: models.GenerationContext, - gate_up_weight: Any, - down_weight: Any, - num_experts: int, - intermediate_dim: int, - hidden_dim: int, -) -> tuple[tuple[Any, ...], tuple[Any, ...], tuple[Any, ...]]: - w1_weights = [] - w3_weights = [] - w2_weights = [] - - for expert_index in range(num_experts): - expert_gate_up = context.graph.reshape( - context.graph.slice(gate_up_weight, 0, expert_index, length=1), - (2 * intermediate_dim, hidden_dim), - ) - w1_weights.append(context.graph.slice(expert_gate_up, 0, 0, length=intermediate_dim)) - w3_weights.append(context.graph.slice(expert_gate_up, 0, intermediate_dim, length=intermediate_dim)) - - expert_down = context.graph.reshape( - context.graph.slice(down_weight, 0, expert_index, length=1), - (hidden_dim, intermediate_dim), - ) - w2_weights.append(context.graph.transpose(expert_down)) - - return tuple(w1_weights), tuple(w3_weights), tuple(w2_weights) - - -def lower_special_cactus(context: models.GenerationContext, node: IRModels.Node) -> Any: - inputs = context.inputs_for(node) - target = node.target - - if target == "cactus.rope": - require_len(node, inputs, 1) - return context.graph.rope( - inputs[0], - theta=float(attr_value(node, "theta", 10_000.0)), - position_offset=int(attr_value(node, "position_offset", 0)), - ) - - if target == "cactus.rope_gptj": - require_len(node, inputs, 1) - return context.graph.rope_gptj( - inputs[0], - theta=float(attr_value(node, "theta", 10_000.0)), - position_offset=int(attr_value(node, "position_offset", 0)), - rot_dim=int(attr_value(node, "rot_dim", 0)), - ) - - if target == "cactus.glu": - require_len(node, inputs, 1) - return context.graph.glu(inputs[0], axis=axis_attr(node, default=-1)) - - if target == "cactus.lstm_cell": - require_len(node, inputs, 7) - return context.graph.lstm_cell(*inputs[:7]) - - if target == "cactus.gated_deltanet_decode": - require_len(node, inputs, 6) - return context.graph.gated_deltanet_decode(*inputs[:6], scale=float(node.attrs.get("scale", 1.0))) - - if target == "cactus.gated_deltanet_prefill": - require_len(node, inputs, 6) - return context.graph.gated_deltanet_prefill(*inputs[:6], chunk_size=int(node.attrs.get("chunk_size", 1)), scale=float(node.attrs.get("scale", 1.0))) - - if target == "cactus.lfm_short_conv_decode": - require_len(node, inputs, 2) - weight = lfm_short_conv_decode_weight(context, node, inputs[1]) - product = context.graph.multiply(inputs[0], weight) - return context.graph.sum(product, axis=normalize_dim(-1, len(meta_shape(node.parents[0])))) - - if target == "cactus.rel_pos_bias": - require_len(node, inputs, 2) - return context.graph.rel_pos_bias(inputs[0], inputs[1], scale=float(node.attrs.get("scale") or 1.0)) - - if target == "cactus.sample": - require_len(node, inputs, 1) - return context.graph.sample(inputs[0], temperature=float(node.attrs.get("temperature", 0.6)), top_p=float(node.attrs.get("top_p", 0.95)), top_k=int(node.attrs.get("top_k", 20))) - - if target == "cactus.scatter_topk": - require_len(node, inputs, 2) - return context.graph.scatter_topk(inputs[0], inputs[1], num_classes=required_int_attr(node, "num_classes")) - - if target == "cactus.gaussian_topk": - require_len(node, inputs, 1) - return context.graph.gaussian_topk(inputs[0], ppf=float(node.attrs.get("ppf", 0.0))) - - if target == "cactus.altup_predict": - require_len(node, inputs, 2) - return context.graph.altup_predict(inputs[0], list(inputs[1:])) - - if target == "cactus.altup_correct": - require_len(node, inputs, 3) - return context.graph.altup_correct(inputs[0], inputs[1], list(inputs[2:])) - - if target == "cactus.stft": - require_len(node, inputs, 2) - return context.graph.stft(inputs[0], inputs[1], stride=required_int_attr(node, "stride"), num_fft_bins=required_int_attr(node, "num_fft_bins")) - - if target == "cactus.rfft": - require_len(node, inputs, 1) - return context.graph.rfft(inputs[0]) - - if target == "cactus.irfft": - require_len(node, inputs, 1) - return context.graph.irfft(inputs[0], output_length=required_int_attr(node, "output_length")) - - if target == "cactus.spectrogram": - require_len(node, inputs, 2) - return context.graph.spectrogram( - inputs[0], - inputs[1], - frame_length=required_int_attr(node, "frame_length"), - hop_length=required_int_attr(node, "hop_length"), - fft_length=required_int_attr(node, "fft_length"), - power=float(node.attrs.get("power", 2.0)), - ) - - if target == "cactus.image_preprocess": - require_len(node, inputs, 1) - if not has_attrs(node, ("src_width", "src_height", "target_width", "target_height", "patch_size", "channels")): - return context.graph.reshape(inputs[0], output_shape(node)) - - return context.graph.image_preprocess( - inputs[0], - src_width=required_int_attr(node, "src_width"), - src_height=required_int_attr(node, "src_height"), - target_width=required_int_attr(node, "target_width"), - target_height=required_int_attr(node, "target_height"), - patch_size=required_int_attr(node, "patch_size"), - channels=required_int_attr(node, "channels"), - rescale_factor=float(node.attrs.get("rescale_factor", 1.0)), - mean=node.attrs.get("mean", (0.0, 0.0, 0.0)), - std_dev=node.attrs.get("std_dev", (1.0, 1.0, 1.0)), - ) - - if target == "cactus.bilinear_interpolation": - require_len(node, inputs, 1) - return context.graph.bilinear_interpolation(inputs[0], required_int_attr(node, "dst_height"), required_int_attr(node, "dst_width")) - - raise UnsupportedLoweringError(f"{node.name}: unsupported special Cactus target {target}") - - -def lfm_short_conv_decode_weight(context: models.GenerationContext, node: IRModels.Node, weight: Any) -> Any: - cache_window_shape = meta_shape(node.parents[0]) if node.parents else () - weight_shape = meta_shape(node.parents[1]) if len(node.parents) > 1 else () - - if len(cache_window_shape) != 3: - return weight - - batch_size, hidden_dim, window_size = cache_window_shape - - if not all(isinstance(dim, int) for dim in cache_window_shape): - return weight - - if tuple(weight_shape) == tuple(cache_window_shape): - return weight - - if len(weight_shape) == 2 and tuple(weight_shape) == (hidden_dim, window_size): - return context.graph.reshape(weight, (batch_size, hidden_dim, window_size)) - - if len(weight_shape) == 3 and tuple(weight_shape) in { - (1, hidden_dim, window_size), - (hidden_dim, 1, window_size), - }: - return context.graph.reshape(weight, (batch_size, hidden_dim, window_size)) - - raise UnsupportedLoweringError( - f"{node.name}: lfm short conv decode weight shape {weight_shape} is incompatible with cache window {cache_window_shape}" - ) - - -def lower_unsupported_semantic(context: models.GenerationContext, node: IRModels.Node) -> Any: - raise UnsupportedLoweringError(f"{node.name}: {node.target} has no safe Cactus lowering yet") - - -def require_input_count(context: models.GenerationContext, node: IRModels.Node, min_count: int) -> tuple[Any, ...]: - inputs = context.inputs_for(node) - require_len(node, inputs, min_count) - return inputs - - -def require_at_least_one_input(context: models.GenerationContext, node: IRModels.Node) -> tuple[Any, ...]: - return require_input_count(context, node, 1) - - -def require_len(node: IRModels.Node, inputs: tuple[Any, ...], min_count: int) -> None: - if len(inputs) < min_count: - raise unsupported_arity(node, len(inputs), f"at least {min_count} inputs") - - -def unsupported_arity(node: IRModels.Node, actual: int, expected: str) -> UnsupportedLoweringError: - return UnsupportedLoweringError(f"{node.name}: {node.target} got {actual} lowered inputs; expected {expected}") - - -def tensor_dtype(node: IRModels.Node) -> str | None: - if isinstance(node.tensor_output_meta, dict): - dtype = node.tensor_output_meta.get("dtype") - - if dtype is not None: - return str(dtype) - - if isinstance(node.tensor_output_meta, list) and node.tensor_output_meta: - first_meta = node.tensor_output_meta[0] - - if isinstance(first_meta, dict): - dtype = first_meta.get("dtype") - - if dtype is not None: - return str(dtype) - - return None - - -def cactus_precision(graph: Any, dtype: str | None) -> int: - precision_name = constants.DTYPE_TO_PRECISION.get(str(dtype), constants.DEFAULT_INPUT_PRECISION) - return int(getattr(graph, precision_name)) - - -def cast_to_precision(context: models.GenerationContext, value: Any, precision: int) -> Any: - if getattr(value, "dtype", precision) == precision: - return value - - return context.graph.precision_cast(value, precision) - - -def graph_input_shape(node: IRModels.Node) -> tuple[tuple[int, ...], tuple[bool, ...]]: - shape = meta_shape(node) - - if not shape: - raw_shape = node.tensor_output_meta.get("shape") if isinstance(node.tensor_output_meta, dict) else None - if raw_shape == [] or raw_shape == (): - return (1,), (False,) - - return (), () - - dims: list[int] = [] - dynamic_dims: list[bool] = [] - - for dim in shape: - if isinstance(dim, int) and dim >= 0: - dims.append(dim) - dynamic_dims.append(False) - continue - - if isinstance(dim, str) and dim.isdigit(): - dims.append(int(dim)) - dynamic_dims.append(False) - continue - - dims.append(1) - dynamic_dims.append(True) - - return tuple(dims), tuple(dynamic_dims) - - -def meta_shape(node: IRModels.Node) -> tuple[Any, ...]: - if isinstance(node.tensor_output_meta, dict): - shape = node.tensor_output_meta.get("shape") - - if isinstance(shape, list): - return tuple(shape) - - if isinstance(shape, tuple): - return shape - - if isinstance(node.tensor_output_meta, list) and node.tensor_output_meta: - first_meta = node.tensor_output_meta[0] - - if isinstance(first_meta, dict): - shape = first_meta.get("shape") - - if isinstance(shape, list): - return tuple(shape) - - if isinstance(shape, tuple): - return shape - - return () - - -def output_shape(node: IRModels.Node) -> tuple[int, ...]: - shape, _ = graph_input_shape(node) - - if shape: - return shape - - raise UnsupportedLoweringError(f"{node.name}: missing concrete output shape") - - -def concrete_dim(dim: Any) -> int | None: - if isinstance(dim, int) and dim >= 0: - return dim - - if isinstance(dim, str) and dim.isdigit(): - return int(dim) - - return None - - -def concrete_shape(shape: tuple[Any, ...]) -> tuple[int, ...] | None: - dims = tuple(concrete_dim(dim) for dim in shape) - - if any(dim is None for dim in dims): - return None - - return tuple(int(dim) for dim in dims) - - -def reduction_dropped_shape(shape: tuple[Any, ...], axis: int) -> tuple[Any, ...]: - dropped = tuple(dim for index, dim in enumerate(shape) if index != axis) - return dropped or (1,) - - -def shape_matches_tensor(node: IRModels.Node, expected_shape: tuple[int, ...]) -> bool: - actual_shape = meta_shape(node) - return tuple(actual_shape) == tuple(expected_shape) - - -def element_count(shape: tuple[Any, ...]) -> int | None: - if not shape: - return None - - count = 1 - - for dim in shape: - if not isinstance(dim, int) or dim < 0: - return None - - count *= dim - - return count - - -def shape_attr(node: IRModels.Node) -> tuple[int, ...]: - raw_shape = node.attrs.get("shape") - - if raw_shape is None: - return output_shape(node) - - resolved_output_shape = output_shape(node) - dims: list[int] = [] - - for index, dim in enumerate(raw_shape): - if isinstance(dim, int) and dim >= 0: - dims.append(dim) - elif isinstance(dim, str) and dim.isdigit(): - dims.append(int(dim)) - elif index < len(resolved_output_shape): - dims.append(resolved_output_shape[index]) - else: - raise UnsupportedLoweringError(f"{node.name}: cannot resolve shape dim {dim!r}") - - return tuple(dims) - - -def axis_attr(node: IRModels.Node, default: int | None = None) -> int | None: - axis = node.attrs.get("axis", node.attrs.get("dim", node.attrs.get("arg_1", default))) - - if isinstance(axis, list): - if len(axis) != 1: - raise UnsupportedLoweringError(f"{node.name}: Cactus reduction lowering only supports one axis") - return int(axis[0]) - - if axis is None: - return None - - return int(axis) - - -def scalar_attr(node: IRModels.Node, name: str) -> Any | None: - value = node.attrs.get(name) - - if isinstance(value, dict) and "node" in value: - return None - - if isinstance(value, list) and len(value) == 1: - return value[0] - - return value - - -def scalar_weight_bound_value(context: models.GenerationContext, node: IRModels.Node) -> float | None: - resolver = context.component.weight_resolver - - if resolver is None: - return None - - record = resolver.resolve(node.name) - - if record is None or record.output_name is None: - return None - - if record.shape and math.prod(int(dim) for dim in record.shape) != 1: - return None - - value = read_scalar_cactus_weight(resolver.weights_dir / record.output_name) - - if value is None: - return None - - return float(value) * float(record.scale_factor) - - -def read_scalar_cactus_weight(path: Path) -> float | None: - if not path.exists(): - return None - - with path.open("rb") as f: - header = f.read(84) - - if len(header) != 84 or header[:4] != b"CACT": - return None - - fields = struct.unpack("= 2: - return float(struct.unpack("= 4: - return float(struct.unpack(" int: - if alignment <= 0: - alignment = 32 - - remainder = offset % alignment - return offset if remainder == 0 else offset + alignment - remainder - - -def attr_value(node: IRModels.Node, name: str, default: Any) -> Any: - value = node.attrs.get(name) - return default if value is None else value - - -def epsilon_attr(node: IRModels.Node) -> float: - return float(node.attrs.get("epsilon", node.attrs.get("eps", 1e-5))) - - -def required_int_attr(node: IRModels.Node, name: str) -> int: - value = node.attrs.get(name) - - if value is None: - raise UnsupportedLoweringError(f"{node.name}: missing required attr {name}") - - return int(value) - - -def has_attrs(node: IRModels.Node, names: tuple[str, ...]) -> bool: - return all(node.attrs.get(name) is not None for name in names) - - -def first_int(value: Any, default: int) -> int: - if value is None: - return default - - if isinstance(value, (list, tuple)): - if not value: - return default - return int(value[0]) - - return int(value) - - -def tuple_int_values(value: Any, default: int) -> tuple[int, int]: - if value is None: - return (default, default) - - if isinstance(value, (list, tuple)): - if not value: - return (default, default) - - if len(value) == 1: - item = int(value[0]) - return (item, item) - - return (int(value[0]), int(value[1])) - - item = int(value) - return (item, item) - - -def tuple_ints(value: Any) -> tuple[int, ...]: - if not isinstance(value, (list, tuple)): - raise TypeError(f"Expected list/tuple of ints, got {type(value).__name__}") - - return tuple(int(item) for item in value) - - -def parent_rank(node: IRModels.Node) -> int: - if not node.parents: - raise UnsupportedLoweringError(f"{node.name}: cannot infer parent rank without parents") - - rank = len(meta_shape(node.parents[0])) - - if rank == 0: - raise UnsupportedLoweringError(f"{node.name}: cannot infer parent rank without parent shape metadata") - - return rank - - -def output_rank(node: IRModels.Node) -> int: - return len(meta_shape(node)) - - -def swap_permutation(rank: int, dim0: int, dim1: int) -> tuple[int, ...]: - dim0 = normalize_dim(dim0, rank) - dim1 = normalize_dim(dim1, rank) - permutation = list(range(rank)) - permutation[dim0], permutation[dim1] = permutation[dim1], permutation[dim0] - return tuple(permutation) - - -def normalize_dim(dim: int, rank: int) -> int: - if dim < 0: - dim += rank - - if dim < 0 or dim >= rank: - raise ValueError(f"Dimension {dim} is outside rank {rank}") - - return dim - - -def slice_length(node: IRModels.Node, start: int) -> int: - if "length" in node.attrs: - return int(node.attrs["length"]) - - end = node.attrs.get("end") - - if end is None: - return open_slice_length(node, start) - - end = int(end) - - if end >= constants.OPEN_SLICE_END: - return open_slice_length(node, start) - - return max(end - start, 0) - - -def open_slice_length(node: IRModels.Node, start: int) -> int: - if not node.parents: - return 0 - - axis = int(node.attrs.get("axis", node.attrs.get("dim", 0))) - parent_shape = meta_shape(node.parents[0]) - - if axis < 0: - axis += len(parent_shape) - - if axis < 0 or axis >= len(parent_shape): - return 0 - - axis_size = parent_shape[axis] - - if not isinstance(axis_size, int): - return 0 - - return max(axis_size - start, 0) - - -def binary_method_to_scalar_method(method: str) -> str: - return { - "add": "scalar_add", - "subtract": "scalar_subtract", - "multiply": "scalar_multiply", - "divide": "scalar_divide", - "not_equal": "scalar_not_equal", - }[method] - - -def ensure_tensor_sequence(value: Any) -> list[Any]: - if isinstance(value, list): - return value - - if isinstance(value, tuple): - return list(value) - - return [value] - - -def dump_result_manifest(result: models.GenerationResult, output_path: str | Path) -> Path: - path = Path(output_path) - path.parent.mkdir(parents=True, exist_ok=True) - path.write_text( - json.dumps( - { - "component_paths": {name: str(path_) for name, path_ in result.component_paths.items()}, - "component_manifest_paths": {name: str(path_) for name, path_ in result.component_manifest_paths.items()}, - "engine_manifest_path": str(result.engine_manifest_path) if result.engine_manifest_path is not None else None, - "runtime_plan_path": str(result.runtime_plan_path) if result.runtime_plan_path is not None else None, - "unsupported_nodes": list(result.unsupported_nodes), - "warnings": list(result.warnings), - "ok": result.ok, - }, - indent=4, - ), - encoding="utf-8", - ) - return path diff --git a/python/cactus/transpiler/Generator/lowering_basic_ops.py b/python/cactus/transpiler/Generator/lowering_basic_ops.py new file mode 100644 index 000000000..90acda8ac --- /dev/null +++ b/python/cactus/transpiler/Generator/lowering_basic_ops.py @@ -0,0 +1,795 @@ +from __future__ import annotations + +from typing import Any + +from . import constants +from . import models +from .errors import UnsupportedLoweringError +from .lowering_cache import lower_decode_cache_cat, lower_prefill_cache_cat +from .lowering_utils import * +from ..IR import models as IRModels + + +def lower_binary(context: models.GenerationContext, node: IRModels.Node) -> Any: + inputs = context.inputs_for(node) + method = constants.BINARY_TARGETS[node.target] + + if len(inputs) == 2: + left, right = align_binary_inputs(context, node, inputs) + + if method == "add" and context.component.name in constants.GEMMA_ADD_CLIPPED_COMPONENTS and looks_like_gemma_residual_add(node): + return context.graph.add_clipped(left, right) + + if method == "multiply" and looks_like_gemma4_mlp_product(node): + if is_gemma4_mlp_gate_activation(node.parents[0]): + left = context.graph.scalar_multiply(left, constants.GEMMA4_MLP_PRODUCT_SCALE) + else: + right = context.graph.scalar_multiply(right, constants.GEMMA4_MLP_PRODUCT_SCALE) + + return getattr(context.graph, method)(left, right) + + scalar = scalar_attr(node, "other") + + if len(inputs) == 1 and scalar is not None: + scalar_method = binary_method_to_scalar_method(method) + return getattr(context.graph, scalar_method)(inputs[0], scalar) + + raise unsupported_arity(node, len(inputs), "two tensor inputs or one tensor plus scalar attr") + + +def looks_like_gemma_residual_add(node: IRModels.Node) -> bool: + if len(node.parents) != 2: + return False + + shape = meta_shape(node) + + if not shape or len(shape) > 4: + return False + + left, right = node.parents + return ( + is_gemma_residual_branch(left, right) + or is_gemma_residual_branch(right, left) + ) + + +def is_gemma_residual_branch(residual: IRModels.Node, branch: IRModels.Node) -> bool: + branch_source = strip_precision_passthrough(branch) + + if branch_source.target != "cactus.rms_norm" or not branch_source.parents: + return False + + return strip_precision_passthrough(residual) is not strip_precision_passthrough(branch_source.parents[0]) + + +def looks_like_gemma4_mlp_product(node: IRModels.Node) -> bool: + if len(node.parents) != 2: + return False + + text = node_context_text(node) + + if "gemma" not in text.lower() or "language_model" not in text or ".mlp" not in text: + return False + + left, right = node.parents + return ( + is_gemma4_mlp_gate_activation(left) and has_ancestor_text(right, "mlp_up_proj") + ) or ( + is_gemma4_mlp_gate_activation(right) and has_ancestor_text(left, "mlp_up_proj") + ) + + +def is_gemma4_mlp_gate_activation(node: IRModels.Node) -> bool: + source = strip_precision_passthrough(node) + + if source.target not in {"cactus.gelu", "aten.gelu.default"}: + return False + + return has_ancestor_text(source, "mlp_gate_proj") + + +def has_ancestor_text(node: IRModels.Node, pattern: str, max_depth: int = 8) -> bool: + pattern = pattern.lower() + worklist: list[tuple[IRModels.Node, int]] = [(node, 0)] + seen: set[str] = set() + + while worklist: + current, depth = worklist.pop(0) + + if current.name in seen: + continue + + seen.add(current.name) + + if pattern in node_context_text(current).lower(): + return True + + if depth >= max_depth: + continue + + for parent in current.parents: + worklist.append((parent, depth + 1)) + + return False + + +def node_context_text(node: IRModels.Node) -> str: + return f"{node.name} {node.target} {node.module_stack!r}" + + +def strip_precision_passthrough(node: IRModels.Node) -> IRModels.Node: + current = node + passthrough_targets = constants.PASS_THROUGH_TARGETS | constants.TO_COPY_TARGETS + + while current.target in passthrough_targets and len(current.parents) == 1: + current = current.parents[0] + + return current + + +def lower_scalar(context: models.GenerationContext, node: IRModels.Node) -> Any: + inputs = require_input_count(context, node, 1) + method, attr_name = constants.SCALAR_TARGETS[node.target] + value = scalar_attr(node, attr_name) + + if value is None: + value = scalar_attr(node, "arg_1") + + if value is None: + if method == "scalar_equal": + return context.graph.scalar_multiply(inputs[0], 0.0) + + if method == "scalar_not_equal": + return context.graph.scalar_add(context.graph.scalar_multiply(inputs[0], 0.0), 1.0) + + raise UnsupportedLoweringError(f"{node.name}: scalar lowering {node.target} missing attr {attr_name}") + + if method == "scalar_multiply" and float(value) == 1.0: + return inputs[0] + + return getattr(context.graph, method)(inputs[0], value) + + +def align_binary_inputs(context: models.GenerationContext, node: IRModels.Node, inputs: tuple[Any, ...]) -> tuple[Any, Any]: + target_shape = output_shape(node) + left = align_input_to_shape(context, inputs[0], node.parents[0], target_shape) + right = align_input_to_shape(context, inputs[1], node.parents[1], target_shape) + return align_binary_precision(context, left, right) + + +def align_binary_precision(context: models.GenerationContext, left: Any, right: Any) -> tuple[Any, Any]: + left_precision = getattr(left, "dtype", None) + right_precision = getattr(right, "dtype", None) + + if left_precision is None or right_precision is None or left_precision == right_precision: + return left, right + + return left, cast_to_precision(context, right, left_precision) + + +def align_input_to_shape(context: models.GenerationContext, value: Any, source_node: IRModels.Node, target_shape: tuple[Any, ...]) -> Any: + value = align_value_to_declared_shape(context, value, source_node) + source_shape = meta_shape(source_node) + + if shape_matches_tensor(source_node, target_shape): + return value + + if element_count(source_shape) == element_count(target_shape): + return context.graph.reshape(value, target_shape) + + return value + + +def align_value_to_declared_shape(context: models.GenerationContext, value: Any, source_node: IRModels.Node) -> Any: + actual_shape = tuple(getattr(value, "shape", ())) + declared_shape = meta_shape(source_node) + + if not actual_shape or not declared_shape or len(declared_shape) <= len(actual_shape): + return value + + resolved_shape = resolve_declared_shape_from_actual(declared_shape, actual_shape) + + if resolved_shape is None or tuple(resolved_shape) == actual_shape: + return value + + return context.graph.reshape(value, resolved_shape) + + +def resolve_declared_shape_from_actual(declared_shape: tuple[Any, ...], actual_shape: tuple[int, ...]) -> tuple[int, ...] | None: + actual_index = len(actual_shape) - 1 + resolved = [1 for _ in declared_shape] + + for declared_index in range(len(declared_shape) - 1, -1, -1): + declared_dim = concrete_dim(declared_shape[declared_index]) + + if actual_index >= 0 and (declared_dim is None or declared_dim == actual_shape[actual_index]): + resolved[declared_index] = actual_shape[actual_index] + actual_index -= 1 + continue + + if declared_dim == 1: + resolved[declared_index] = 1 + continue + + return None + + if actual_index >= 0: + return None + + return tuple(resolved) + + +def lower_unary(context: models.GenerationContext, node: IRModels.Node) -> Any: + inputs = require_input_count(context, node, 1) + method = constants.UNARY_TARGETS[node.target] + value = fp16_tensor(context, inputs[0]) if method in constants.FP16_UNARY_METHODS else inputs[0] + return getattr(context.graph, method)(value) + + +def lower_log1p(context: models.GenerationContext, node: IRModels.Node) -> Any: + inputs = require_input_count(context, node, 1) + + if node.parents and node.parents[0].target in {"aten.exp.default", "cactus.scalar_exp"} and node.parents[0].parents: + x = context.require(node.parents[0].parents[0].name) + return lower_stable_softplus(context, x) + + return context.graph.scalar_log(context.graph.scalar_add(inputs[0], 1.0)) + + +def lower_stable_softplus(context: models.GenerationContext, x: Any) -> Any: + x = cast_to_precision(context, x, context.graph.FP16) + neg_abs = context.graph.scalar_multiply(context.graph.abs(x), -1.0) + log_term = context.graph.scalar_log(context.graph.scalar_add(context.graph.scalar_exp(neg_abs), 1.0)) + positive_term = context.graph.clamp(x, 0.0, math.inf) + return context.graph.add(log_term, positive_term) + + +def lower_pow(context: models.GenerationContext, node: IRModels.Node) -> Any: + inputs = require_input_count(context, node, 1) + exponent = scalar_attr(node, "exponent") + + if exponent is None: + raise UnsupportedLoweringError(f"{node.name}: pow lowering missing exponent attr") + + return context.graph.pow(inputs[0], exponent) + + +def lower_reduce(context: models.GenerationContext, node: IRModels.Node) -> Any: + inputs = require_input_count(context, node, 1) + method = constants.REDUCE_TARGETS[node.target] + axis = axis_attr(node) + + if axis is None: + raise UnsupportedLoweringError(f"{node.name}: reduce lowering missing axis/dim attr") + + normalized_axis = normalize_dim(axis, len(meta_shape(node.parents[0]))) + reduced = getattr(context.graph, method)(inputs[0], normalized_axis) + target_shape = output_shape(node) + dropped_shape = reduction_dropped_shape(meta_shape(node.parents[0]), normalized_axis) + + if bool(node.attrs.get("keepdim", False)) or ( + tuple(dropped_shape) != tuple(target_shape) + and element_count(dropped_shape) == element_count(target_shape) + ): + return context.graph.reshape(reduced, target_shape) + + return reduced + + +def lower_shape(context: models.GenerationContext, node: IRModels.Node) -> Any: + inputs = require_input_count(context, node, 1) + method = constants.SHAPE_TARGETS[node.target] + shape = decoder_prefill_logits_shape(context, node, shape_attr(node)) + + return getattr(context.graph, method)(inputs[0], shape) + + +def lower_expand(context: models.GenerationContext, node: IRModels.Node) -> Any: + inputs = require_input_count(context, node, 1) + target_shape = shape_attr(node) + + if shape_matches_tensor(node.parents[0], target_shape): + return inputs[0] + + if element_count(meta_shape(node.parents[0])) == element_count(target_shape): + return context.graph.reshape(inputs[0], target_shape) + + try: + return context.graph.expand(inputs[0], target_shape) + except AttributeError as e: + raise UnsupportedLoweringError( + f"{node.name}: native Cactus expand is unavailable for broadcast shape " + f"{meta_shape(node.parents[0])} -> {target_shape}" + ) from e + + +def lower_unsqueeze(context: models.GenerationContext, node: IRModels.Node) -> Any: + inputs = require_input_count(context, node, 1) + return context.graph.reshape(inputs[0], output_shape(node)) + + +def lower_flatten(context: models.GenerationContext, node: IRModels.Node) -> Any: + inputs = require_input_count(context, node, 1) + return context.graph.flatten( + inputs[0], + start_dim=int(node.attrs.get("start_dim", 0)), + end_dim=int(node.attrs.get("end_dim", -1)), + ) + + +def lower_repeat(context: models.GenerationContext, node: IRModels.Node) -> Any: + inputs = require_input_count(context, node, 1) + repeats = node.attrs.get("repeats", node.attrs.get("arg_1")) + + if isinstance(repeats, (list, tuple)) and all(int(value) == 1 for value in repeats): + return inputs[0] + + raise UnsupportedLoweringError(f"{node.name}: repeat lowering only supports no-op repeat factors, got {repeats}") + + +def lower_transpose(context: models.GenerationContext, node: IRModels.Node) -> Any: + inputs = require_input_count(context, node, 1) + + if can_alias_quantized_weight_transpose(context, node, inputs[0]): + return inputs[0] + + permutation = node.attrs.get("permutation") + + if permutation is not None: + return context.graph.permute(inputs[0], tuple_ints(permutation)) + + if node.target == "aten.t.default": + return context.graph.transpose(inputs[0]) + + dim0 = node.attrs.get("dim0") + dim1 = node.attrs.get("dim1") + + if dim0 is None or dim1 is None: + raise UnsupportedLoweringError(f"{node.name}: transpose lowering missing permutation or dim attrs") + + return context.graph.permute(inputs[0], swap_permutation(parent_rank(node), int(dim0), int(dim1))) + + +def can_alias_quantized_weight_transpose(context: models.GenerationContext, node: IRModels.Node, value: Any) -> bool: + return ( + is_weight_transpose_node(node) + and is_cq_tensor(context, value) + and bool(node.children) + and all(child.target in constants.MATMUL_TARGETS or child.target in constants.ADDM_CONST_TARGETS for child in node.children) + ) + + +def is_weight_transpose_node(node: IRModels.Node) -> bool: + if not node.parents: + return False + + parent = node.parents[0] + if not parent.is_placeholder or parent.value_kind not in constants.WEIGHT_VALUE_KINDS: + return False + + if node.target == "aten.t.default": + return True + + rank = len(meta_shape(parent)) + if rank < 2: + return False + + permutation = node.attrs.get("permutation") + if permutation is not None: + return tuple_ints(permutation) == swap_permutation(rank, rank - 2, rank - 1) + + dim0 = node.attrs.get("dim0") + dim1 = node.attrs.get("dim1") + if dim0 is None or dim1 is None: + return False + + return { + normalize_dim(int(dim0), rank), + normalize_dim(int(dim1), rank), + } == {rank - 2, rank - 1} + + +def is_cq_tensor(context: models.GenerationContext, value: Any) -> bool: + return getattr(value, "dtype", None) in { + context.graph.CQ1, + context.graph.CQ2, + context.graph.CQ3, + context.graph.CQ4, + } + + +def lower_slice(context: models.GenerationContext, node: IRModels.Node) -> Any: + inputs = require_input_count(context, node, 1) + axis_value = node.attrs.get("axis") + + if axis_value is None: + axis_value = node.attrs.get("dim") + + if axis_value is None: + if shape_matches_tensor(node.parents[0], output_shape(node)): + return inputs[0] + + raise UnsupportedLoweringError(f"{node.name}: slice lowering missing axis/dim attr") + + axis = normalize_dim(int(axis_value), len(meta_shape(node.parents[0]))) + + if node.target == "aten.select.int" or "index_value" in node.attrs: + return context.graph.index(inputs[0], int(node.attrs.get("index_value", node.attrs.get("index", 0))), axis=axis) + + start = int(node.attrs.get("start", 0) or 0) + step = node.attrs.get("step", 1) + + if is_full_axis_slice(node, axis, start, step): + return inputs[0] + + if step is not None and int(step) != 1: + target_shape = output_shape(node) + length = target_shape[axis] if axis < len(target_shape) else slice_length(node, start) + return context.graph.strided_slice(inputs[0], axis, start, length, int(step)) + + length = slice_length(node, start) + return context.graph.slice(inputs[0], axis, start, length=length) + + +def is_full_axis_slice(node: IRModels.Node, axis: int, start: int, step: Any) -> bool: + if start != 0: + return False + + if step is not None and int(step) != 1: + return False + + if not node.parents: + return False + + return shape_matches_tensor(node.parents[0], output_shape(node)) + + +def lower_index(context: models.GenerationContext, node: IRModels.Node) -> Any: + if node.target == "aten.index.Tensor": + return lower_index_tensor(context, node) + + inputs = require_input_count(context, node, 1) + index_value = node.attrs.get("index_value", node.attrs.get("index")) + + if index_value is None: + raise UnsupportedLoweringError(f"{node.name}: index lowering missing index_value/index attr") + + return context.graph.index(inputs[0], int(index_value), axis=int(node.attrs.get("axis", node.attrs.get("dim", 0)))) + + +def lower_index_tensor(context: models.GenerationContext, node: IRModels.Node) -> Any: + inputs = context.inputs_for(node) + + if len(inputs) < 2: + raise unsupported_arity(node, len(inputs), "source tensor plus at least one index tensor") + + source = inputs[0] + source_shape = meta_shape(node.parents[0]) + target_shape = meta_shape(node) + + if len(inputs) >= 2 and len(source_shape) == 2 and len(target_shape) == 3 and target_shape[-1] == source_shape[-1]: + return context.graph.embedding_from_tensor(source, inputs[1]) + + if len(source_shape) == 3 and len(target_shape) == 2 and concrete_dim(source_shape[0]) == 1: + feature_count = concrete_dim(source_shape[1]) + feature_dim = concrete_dim(source_shape[2]) + + if feature_count is not None and feature_dim is not None: + return context.graph.reshape(source, (feature_count, feature_dim)) + + if len(source_shape) == 2 and len(target_shape) == 4 and concrete_dim(source_shape[0]) == 1: + source_length = concrete_dim(source_shape[1]) + target_length = concrete_dim(target_shape[-1]) + + if source_length is not None and target_length is not None: + reshaped = context.graph.reshape(source, (1, 1, 1, source_length)) + + if target_length < source_length: + return context.graph.slice(reshaped, 3, 0, length=target_length) + + if target_length == source_length: + return reshaped + + concrete_target_shape = concrete_shape(target_shape) + + if concrete_target_shape is not None and element_count(source_shape) == element_count(concrete_target_shape): + return context.graph.reshape(source, concrete_target_shape) + + raise UnsupportedLoweringError(f"{node.name}: unsupported aten.index.Tensor pattern") + + +def lower_where(context: models.GenerationContext, node: IRModels.Node) -> Any: + inputs = require_input_count(context, node, 3) + output_precision = cactus_precision(context.graph, tensor_dtype(node)) + true_value = cast_to_precision(context, inputs[1], output_precision) + false_value = cast_to_precision(context, inputs[2], output_precision) + return context.graph.where(inputs[0], true_value, false_value) + + +def lower_unfold(context: models.GenerationContext, node: IRModels.Node) -> Any: + inputs = require_input_count(context, node, 1) + dimension = int(node.attrs.get("dimension", node.attrs.get("dim", node.attrs.get("arg_1", 0)))) + size = int(node.attrs.get("size", node.attrs.get("arg_2", 1))) + step = int(node.attrs.get("step", node.attrs.get("arg_3", 1))) + return context.graph.unfold(inputs[0], dimension, size, step) + + +def lower_cat(context: models.GenerationContext, node: IRModels.Node) -> Any: + if node.name in context.prefill_cache_cat_annotations: + prefill_cache = lower_prefill_cache_cat(context, node) + + if prefill_cache is not None: + return prefill_cache + + decode_cache = lower_decode_cache_cat(context, node) + + if decode_cache is not None: + return decode_cache + + inputs = context.inputs_for(node) + + if not inputs: + raise unsupported_arity(node, 0, "at least one input") + + passthrough = empty_cat_passthrough(node, inputs) + + if passthrough is not None: + return passthrough + + axis = normalize_dim(int(node.attrs.get("axis", node.attrs.get("dim", 0))), cat_rank(node)) + return context.graph.cat(fp16_cat_inputs(context, inputs), axis=axis) + + +def fp16_cat_inputs(context: models.GenerationContext, inputs: tuple[Any, ...]) -> tuple[Any, ...]: + fp16 = int(context.graph.FP16) + return tuple(value if getattr(value, "dtype", None) == fp16 else context.graph.precision_cast(value, fp16) for value in inputs) + + +def lower_matmul(context: models.GenerationContext, node: IRModels.Node) -> Any: + inputs = require_input_count(context, node, 2) + + if node.target == "aten.bmm.default": + return lower_bmm(context, node, inputs) + + lhs = matmul_activation_operand(context, inputs[0]) + rhs = inputs[1] + lhs = slice_decoder_prefill_logits_lhs(context, node, lhs) + + if len(node.parents) >= 2 and is_weight_transpose_node(node.parents[1]): + original_weight_node = node.parents[1].parents[0] + original_weight = context.lookup(original_weight_node.name) + + if original_weight is not None and is_cq_tensor(context, original_weight): + return context.graph.matmul(lhs, original_weight, pretransposed_rhs=True) + + if not is_cq_tensor(context, rhs): + rhs = matmul_activation_operand(context, rhs) + + return context.graph.matmul( + lhs, + rhs, + pretransposed_rhs=bool(node.attrs.get("pretransposed_rhs", False)), + ) + + +def matmul_activation_operand(context: models.GenerationContext, value: Any) -> Any: + if getattr(value, "dtype", None) == context.graph.FP32: + return cast_to_precision(context, value, context.graph.FP16) + + return value + + +def slice_decoder_prefill_logits_lhs(context: models.GenerationContext, node: IRModels.Node, lhs: Any) -> Any: + if not is_decoder_prefill_logits_matmul(context, node): + return lhs + + lhs_shape = meta_shape(node.parents[0]) if node.parents else () + + if len(lhs_shape) != 2 or not isinstance(lhs_shape[0], int) or lhs_shape[0] <= 1: + return lhs + + return context.graph.slice(lhs, 0, int(lhs_shape[0]) - 1, length=1) + + +def is_decoder_prefill_logits_matmul(context: models.GenerationContext, node: IRModels.Node) -> bool: + shape = meta_shape(node) + + return ( + context.component.name == "decoder_prefill_chunk" + and node.target in constants.MATMUL_TARGETS + and len(shape) == 2 + and len(node.parents) >= 2 + and isinstance(shape[0], int) + and shape[0] > 1 + and isinstance(shape[-1], int) + and shape[-1] >= 100000 + ) + + +def decoder_prefill_logits_shape(context: models.GenerationContext, node: IRModels.Node, shape: tuple[int, ...]) -> tuple[int, ...]: + if ( + context.component.name == "decoder_prefill_chunk" + and len(shape) == 3 + and shape[0] == 1 + and isinstance(shape[1], int) + and shape[1] > 1 + and shape[2] >= 100000 + ): + return (1, 1, shape[2]) + + return shape + + +def empty_cat_passthrough(node: IRModels.Node, inputs: tuple[Any, ...]) -> Any | None: + target_shape = meta_shape(node) + non_empty_inputs = [] + + for input_value, parent in zip(inputs, node.parents): + parent_shape = meta_shape(parent) + parent_elements = element_count(parent_shape) + + if parent_elements == 0: + continue + + non_empty_inputs.append((input_value, parent_shape)) + + if len(non_empty_inputs) != 1: + return None + + input_value, input_shape = non_empty_inputs[0] + + if tuple(input_shape) == tuple(target_shape): + return input_value + + return None + + +def cat_rank(node: IRModels.Node) -> int: + target_shape = meta_shape(node) + + if target_shape: + return len(target_shape) + + return len(meta_shape(node.parents[0])) if node.parents else 0 + + +def lower_bmm(context: models.GenerationContext, node: IRModels.Node, inputs: tuple[Any, ...]) -> Any: + left_shape = meta_shape(node.parents[0]) + right_shape = meta_shape(node.parents[1]) + + if len(left_shape) != 3 or len(right_shape) != 3: + raise UnsupportedLoweringError(f"{node.name}: bmm lowering requires rank-3 inputs") + + if left_shape[0] != right_shape[0]: + raise UnsupportedLoweringError(f"{node.name}: bmm lowering requires matching batch dimensions") + + batch_size = concrete_dim(left_shape[0]) + + if batch_size is None: + raise UnsupportedLoweringError(f"{node.name}: bmm lowering requires a concrete batch size") + + if batch_size == 1: + left_2d = matmul_activation_operand(context, context.graph.reshape(inputs[0], (left_shape[1], left_shape[2]))) + right_2d = matmul_activation_operand(context, context.graph.reshape(inputs[1], (right_shape[1], right_shape[2]))) + product = context.graph.matmul(left_2d, right_2d) + return context.graph.reshape(product, output_shape(node)) + + batch_outputs = [] + + for batch_index in range(batch_size): + left_2d = matmul_activation_operand(context, context.graph.index(inputs[0], batch_index, axis=0)) + right_2d = matmul_activation_operand(context, context.graph.index(inputs[1], batch_index, axis=0)) + product = context.graph.matmul(left_2d, right_2d) + batch_outputs.append(context.graph.reshape(product, (1, left_shape[1], right_shape[2]))) + + return context.graph.cat(batch_outputs, axis=0) + + +def lower_addmm(context: models.GenerationContext, node: IRModels.Node) -> Any: + inputs = require_input_count(context, node, 3) + pretransposed_rhs = False + rhs = inputs[2] + + if len(node.parents) >= 3 and is_weight_transpose_node(node.parents[2]): + original_weight_node = node.parents[2].parents[0] + original_weight = context.lookup(original_weight_node.name) + + if original_weight is not None and is_cq_tensor(context, original_weight): + rhs = original_weight + pretransposed_rhs = True + + product = context.graph.matmul( + matmul_activation_operand(context, inputs[1]), + rhs if is_cq_tensor(context, rhs) else matmul_activation_operand(context, rhs), + pretransposed_rhs=pretransposed_rhs, + ) + return context.graph.add(inputs[0], product) + + +def lower_split(context: models.GenerationContext, node: IRModels.Node) -> Any: + inputs = require_input_count(context, node, 1) + split_sizes = node.attrs.get("split_sizes") + + if not isinstance(split_sizes, list): + raise UnsupportedLoweringError(f"{node.name}: split lowering requires split_sizes list") + + axis = normalize_dim(int(node.attrs.get("axis", node.attrs.get("dim", 0))), len(meta_shape(node.parents[0]))) + offset = 0 + outputs = [] + + for size in split_sizes: + size_int = int(size) + outputs.append(context.graph.slice(inputs[0], axis, offset, length=size_int)) + offset += size_int + + return tuple(outputs) + + +def lower_getitem(context: models.GenerationContext, node: IRModels.Node) -> Any: + inputs = require_input_count(context, node, 1) + index = int(node.attrs.get("index", 0)) + source = inputs[0] + + if isinstance(source, (tuple, list)): + return source[index] + + if index == 0: + return source + + raise UnsupportedLoweringError(f"{node.name}: getitem index {index} requires tuple-producing parent") + + +def lower_to_copy(context: models.GenerationContext, node: IRModels.Node) -> Any: + inputs = require_input_count(context, node, 1) + dtype = node.attrs.get("dtype") + + if dtype is None or dtype == "torch.bool": + return inputs[0] + + if str(dtype) == "torch.float32" and can_skip_float32_copy(node): + return inputs[0] + + precision_name = constants.DTYPE_TO_PRECISION.get(str(dtype)) + + if precision_name is None: + return inputs[0] + + return cast_to_precision(context, inputs[0], cactus_precision(context.graph, str(dtype))) + + +def can_skip_float32_copy(node: IRModels.Node) -> bool: + if not node.children or not node.parents: + return False + + source_dtype = tensor_dtype(node.parents[0]) + + if constants.DTYPE_TO_PRECISION.get(str(source_dtype)) != "FP16": + return False + + safe_consumers = { + "aten.bmm.default", + "cactus.attention", + "cactus.attention_cached", + "cactus.add", + "cactus.dense_mlp_tq_fused", + "cactus.multiply", + "cactus.pow", + "cactus.rope", + "cactus.rms_norm", + "cactus.scalar_multiply", + "cactus.silu", + "cactus.slice", + "cactus.transpose", + "cactus.view", + } + return all(child.target in safe_consumers for child in node.children) + + +def lower_pass_through(context: models.GenerationContext, node: IRModels.Node) -> Any: + inputs = require_at_least_one_input(context, node) + return inputs[0] + + +def lower_copy(context: models.GenerationContext, node: IRModels.Node) -> Any: + inputs = require_at_least_one_input(context, node) + return inputs[-1] diff --git a/python/cactus/transpiler/Generator/lowering_cache.py b/python/cactus/transpiler/Generator/lowering_cache.py new file mode 100644 index 000000000..96cdba351 --- /dev/null +++ b/python/cactus/transpiler/Generator/lowering_cache.py @@ -0,0 +1,580 @@ +from __future__ import annotations + +from typing import Any + +from . import models +from .errors import UnsupportedLoweringError +from .lowering_utils import * +from ..Fusions import models as FModels +from ..IR import models as IRModels +from ..RuntimePlan import models as RPModels + + +def cache_attention_generation_plan(graph: IRModels.Graph) -> tuple[frozenset[str], frozenset[str], dict[str, IRModels.CacheAnnotation]]: + skip_names: set[str] = set() + cache_state_names: set[str] = set() + prefill_annotations = prefill_cache_cat_annotations(graph) + + for node in graph.nodes: + if node.target != "cactus.attention" or len(node.parents) < 3: + continue + + key_cat = find_cache_cat_ancestor(node.parents[1], FModels.CacheTensorRole.KEY) + value_cat = find_cache_cat_ancestor(node.parents[2], FModels.CacheTensorRole.VALUE) + + if key_cat is None or value_cat is None: + continue + + key_cache = key_cat.parents[0] + value_cache = value_cat.parents[0] + + if key_cache.cache is None or value_cache.cache is None: + continue + + if key_cache.cache.layer_index != value_cache.cache.layer_index: + continue + + skip_names.update(cache_wrapper_path_names(node.parents[1], key_cat)) + skip_names.update(cache_wrapper_path_names(node.parents[2], value_cat)) + cache_state_names.update((key_cache.name, value_cache.name)) + + for cat_name in prefill_annotations: + cat_node = graph.nodes_map.get(cat_name) + + if cat_node is not None and cat_node.parents: + skip_names.add(cat_node.parents[0].name) + + return frozenset(skip_names), frozenset(cache_state_names), prefill_annotations + + +def prefill_cache_cat_annotations(graph: IRModels.Graph) -> dict[str, IRModels.CacheAnnotation]: + annotations: dict[str, IRModels.CacheAnnotation] = {} + candidates = [node for node in graph.nodes if is_empty_cache_cat(node)] + + for index, node in enumerate(candidates): + shape = concrete_shape(meta_shape(node)) + + if shape is None or len(shape) != 4: + continue + + _, num_kv_heads, sequence_length, head_dim = shape + window_size = prefill_cache_window_size(graph, index // 2) + requested_capacity = int(node.ir_metadata.get("prefill_cache_capacity", sequence_length)) + cache_capacity = prefill_cache_capacity(graph, requested_capacity, window_size) + layer_index = index // 2 + annotations[node.name] = IRModels.CacheAnnotation( + kind=FModels.CacheKind.KV, + role=FModels.CacheTensorRole.KEY if index % 2 == 0 else FModels.CacheTensorRole.VALUE, + tensor_index=index, + layer_index=layer_index, + shape=tuple(shape), + layout="batch_heads_sequence_head_dim", + sequence_length=cache_capacity, + window_size=window_size, + num_kv_heads=int(num_kv_heads), + head_dim=int(head_dim), + source="prefill_empty_cache_cat", + ) + + return annotations + + +def prefill_cache_window_size(graph: IRModels.Graph, layer_index: int) -> int | None: + model_name = str(getattr(graph, "model_name", "")).lower() + + if "gemma" not in model_name: + return None + + return 0 if layer_index % 5 == 4 else 512 + + +def prefill_cache_capacity(graph: IRModels.Graph, sequence_length: int, window_size: int | None) -> int: + model_name = str(getattr(graph, "model_name", "")).lower() + capacity = int(sequence_length) + + if "gemma" in model_name: + if window_size: + return int(window_size) + + return capacity + 1 + + return capacity + + +def is_empty_cache_cat(node: IRModels.Node) -> bool: + if node.target not in {"cactus.cat", "aten.cat.default"} or len(node.parents) < 2: + return False + + first_parent = node.parents[0] + return first_parent.value_kind == FModels.ValueKind.LIFTED_CONSTANT and element_count(meta_shape(first_parent)) == 0 + + +def find_cache_cat_ancestor(node: IRModels.Node, role: str, max_depth: int = 16) -> IRModels.Node | None: + current = node + + for _ in range(max_depth): + if current.target in {"cactus.cat", "aten.cat.default"} and len(current.parents) >= 2: + cache_parent = current.parents[0] + + if cache_parent.cache is not None and cache_parent.cache.kind == FModels.CacheKind.KV and cache_parent.cache.role == role: + return current + + if len(current.parents) != 1: + return None + + current = current.parents[0] + + return None + + +def cache_wrapper_path_names(start: IRModels.Node, stop: IRModels.Node) -> set[str]: + names: set[str] = set() + current = start + + while True: + names.add(current.name) + + if current is stop: + return names + + if len(current.parents) != 1: + return set() + + current = current.parents[0] + + +def should_lower_cache_placeholder_as_state(context: models.GenerationContext, node: IRModels.Node) -> bool: + if node.cache is None: + return False + + if node.name in context.cache_state_placeholder_names: + return True + + if node.cache.kind == FModels.CacheKind.CONV: + return any(child.target in {"cactus.conv_cache_append", "cactus.conv_cache_initialize"} for child in node.children) + + if node.cache.kind == FModels.CacheKind.KV: + if any(child.target in {"cactus.cat", "aten.cat.default"} for child in node.children): + return True + + return bool(node.children) and all(child.target in {"cactus.kv_cache_append", "cactus.attention_cached"} for child in node.children) + + return False + + +def lower_cache_placeholder(context: models.GenerationContext, node: IRModels.Node) -> Any: + annotation = require_cache_annotation(node) + + if annotation.kind == FModels.CacheKind.KV: + cache_state = context.graph.kv_cache_state( + kv_cache_capacity(context, annotation), + required_cache_int(annotation.num_kv_heads, node, "num_kv_heads"), + required_cache_int(annotation.head_dim, node, "head_dim"), + window_size=int(annotation.window_size or 0), + sink_size=0, + num_slots=1, + ) + record_cache_state_binding(context, node, cache_state, annotation) + return cache_state + + if annotation.kind == FModels.CacheKind.CONV: + cache_state = context.graph.conv_cache_state( + required_cache_int(annotation.window_size, node, "window_size"), + required_cache_int(annotation.hidden_dim, node, "hidden_dim"), + ) + record_cache_state_binding(context, node, cache_state, annotation) + return cache_state + + raise UnsupportedLoweringError(f"{node.name}: unsupported cache placeholder kind {annotation.kind}") + + +def lower_decode_cache_cat(context: models.GenerationContext, node: IRModels.Node) -> Any | None: + if len(node.parents) < 2: + return None + + cache_parent = node.parents[0] + + if cache_parent.cache is None or cache_parent.cache.kind != FModels.CacheKind.KV: + return None + + return context.require(cache_parent.name) + + +def lower_prefill_cache_cat(context: models.GenerationContext, node: IRModels.Node) -> Any | None: + annotation = context.prefill_cache_cat_annotations.get(node.name) + + if annotation is None or len(node.parents) < 2: + return None + + new_kv_node = node.parents[1] + original_new_kv = context.require(new_kv_node.name) + new_kv = cache_new_kv_for_native(context, new_kv_node, original_new_kv, annotation.role or "") + new_kv = cast_to_precision(context, new_kv, context.graph.FP16) + cache_state = context.graph.kv_cache_state( + kv_cache_capacity(context, annotation), + required_cache_int(annotation.num_kv_heads, node, "num_kv_heads"), + required_cache_int(annotation.head_dim, node, "head_dim"), + window_size=int(annotation.window_size or 0), + sink_size=0, + num_slots=1, + ) + record_cache_state_binding(context, node, cache_state, annotation) + context.graph.kv_cache_append(new_kv, cache_state, window_size=int(annotation.window_size or 0), sink_size=0) + context.prefill_cache_cat_states[node.name] = cache_state + return original_new_kv + + +def lower_attention(context: models.GenerationContext, node: IRModels.Node) -> Any: + if node.target == "cactus.attention" or node.target == "aten.scaled_dot_product_attention.default": + cached = lower_cache_backed_attention(context, node) + + if cached is not None: + return cached + + inputs = context.inputs_for(node) + require_len(node, inputs, 3) + query, key, value, mask = attention_inputs_for_layout(context, node, inputs) + query = cast_to_precision(context, query, context.graph.FP16) + key = cast_to_precision(context, key, context.graph.FP16) + value = cast_to_precision(context, value, context.graph.FP16) + mask = cast_to_precision(context, mask, context.graph.FP16) if mask is not None else None + output = context.graph.attention( + query, + key, + value, + scale=float(node.attrs.get("scale", 1.0)), + is_causal=bool(node.attrs.get("is_causal", True)), + position_offset=int(node.attrs.get("position_offset", 0)), + window_size=int(node.attrs.get("window_size", 0)), + mask=mask, + additive_mask=bool(node.attrs.get("additive_mask", False)), + ) + return attention_output_for_layout(context, node, output) + + if node.target == "cactus.attention_cached": + inputs = context.inputs_for(node) + require_len(node, inputs, 5) + return context.graph.attention_cached( + inputs[0], + inputs[1], + inputs[2], + inputs[3], + inputs[4], + scale=float(node.attrs.get("scale", 1.0)), + position_offset=cached_attention_position_offset(context, node), + window_size=int(node.attrs.get("window_size", 0)), + v_head_dim=int(node.attrs.get("v_head_dim", 0)), + ) + + raise UnsupportedLoweringError(f"{node.name}: unsupported attention target {node.target}") + + +def lower_cache_backed_attention(context: models.GenerationContext, node: IRModels.Node) -> Any | None: + if node.target != "cactus.attention" or len(node.parents) < 3: + return None + + key_cat = find_attention_cache_cat_ancestor(context, node.parents[1], FModels.CacheTensorRole.KEY) + value_cat = find_attention_cache_cat_ancestor(context, node.parents[2], FModels.CacheTensorRole.VALUE) + + if key_cat is None or value_cat is None: + return None + + key_cache = key_cat.parents[0] + value_cache = value_cat.parents[0] + key_new_node = key_cat.parents[1] + value_new_node = value_cat.parents[1] + + query = context.require(node.parents[0].name) + input_layout = str(node.attrs.get("input_layout", "")) + key_new = cache_new_kv_for_native(context, key_new_node, context.require(key_new_node.name), FModels.CacheTensorRole.KEY, input_layout) + value_new = cache_new_kv_for_native(context, value_new_node, context.require(value_new_node.name), FModels.CacheTensorRole.VALUE, input_layout) + key_cache_state = prefill_cache_state_for_cat(context, key_cat) + value_cache_state = prefill_cache_state_for_cat(context, value_cat) + + if key_cache_state is None: + key_cache_state = context.require(key_cache.name) + + if value_cache_state is None: + value_cache_state = context.require(value_cache.name) + + if node.attrs.get("input_layout") == "bhqd_bhds_bhsd": + query = context.graph.permute(query, (0, 2, 1, 3)) + + query = cast_to_precision(context, query, context.graph.FP16) + key_new = cast_to_precision(context, key_new, context.graph.FP16) + value_new = cast_to_precision(context, value_new, context.graph.FP16) + + cache_pair = (key_cat.name, value_cat.name) + prefill_pair_already_appended = key_cat.name in context.prefill_cache_cat_states and value_cat.name in context.prefill_cache_cat_states + + if cache_pair not in context.appended_cache_pairs and not prefill_pair_already_appended: + context.graph.kv_cache_append(key_new, key_cache_state, window_size=int(node.attrs.get("window_size", 0)), sink_size=0) + context.graph.kv_cache_append(value_new, value_cache_state, window_size=int(node.attrs.get("window_size", 0)), sink_size=0) + context.appended_cache_pairs.add(cache_pair) + + context.values.setdefault(key_cat.name, key_cache_state) + context.values.setdefault(value_cat.name, value_cache_state) + + output = context.graph.attention_cached( + query, + key_new, + value_new, + key_cache_state, + value_cache_state, + scale=float(node.attrs.get("scale", 1.0)), + position_offset=cached_attention_position_offset(context, node), + window_size=int(node.attrs.get("window_size", 0)), + v_head_dim=last_dim(value_new_node), + ) + return attention_output_for_layout(context, node, output) + + +def find_attention_cache_cat_ancestor(context: models.GenerationContext, node: IRModels.Node, role: str) -> IRModels.Node | None: + cache_cat = find_cache_cat_ancestor(node, role) + + if cache_cat is not None: + return cache_cat + + return find_prefill_cache_cat_ancestor(context, node) + + +def find_prefill_cache_cat_ancestor(context: models.GenerationContext, node: IRModels.Node, max_depth: int = 16) -> IRModels.Node | None: + current = node + + for _ in range(max_depth): + if current.name in context.prefill_cache_cat_annotations: + return current + + if len(current.parents) != 1: + return None + + current = current.parents[0] + + return None + + +def prefill_cache_state_for_cat(context: models.GenerationContext, node: IRModels.Node) -> Any | None: + return context.prefill_cache_cat_states.get(node.name) + + +def cached_attention_position_offset(context: models.GenerationContext, node: IRModels.Node) -> int: + if context.component.name == "decoder_step": + return SIZE_T_MAX + + return int(node.attrs.get("position_offset", 0)) + + +def cache_new_kv_for_native(context: models.GenerationContext, node: IRModels.Node, value: Any, role: str, input_layout: str = "") -> Any: + if len(meta_shape(node)) == 4: + return context.graph.permute(value, (0, 2, 1, 3)) + + return value + + +def last_dim(node: IRModels.Node) -> int: + shape = concrete_shape(meta_shape(node)) + + if not shape: + return 0 + + return int(shape[-1]) + + +def attention_inputs_for_layout(context: models.GenerationContext, node: IRModels.Node, inputs: tuple[Any, ...]) -> tuple[Any, Any, Any, Any | None]: + mask = inputs[3] if len(inputs) > 3 else None + + if node.attrs.get("input_layout") == "bhqd_bhds_bhsd": + return ( + context.graph.permute(inputs[0], (0, 2, 1, 3)), + context.graph.permute(inputs[1], (0, 3, 1, 2)), + context.graph.permute(inputs[2], (0, 2, 1, 3)), + mask, + ) + + if node.attrs.get("input_layout") == "bqhd_bhds_bhsd": + return ( + inputs[0], + context.graph.permute(inputs[1], (0, 3, 1, 2)), + context.graph.permute(inputs[2], (0, 2, 1, 3)), + mask, + ) + + return inputs[0], inputs[1], inputs[2], mask + + +def attention_output_for_layout(context: models.GenerationContext, node: IRModels.Node, output: Any) -> Any: + output_layout = node.attrs.get("output_layout") + + if output_layout == "bthd_flat": + return context.graph.reshape(output, output_shape(node)) + + if output_layout == "bhqd": + return context.graph.permute(output, (0, 2, 1, 3)) + + return output + + +def lower_cache(context: models.GenerationContext, node: IRModels.Node) -> Any: + inputs = context.inputs_for(node) + target = node.target + + if target == "cactus.kv_cache_append": + if len(inputs) == 1: + cache_state = create_cache_state_for_cache_output(context, node) + context.graph.kv_cache_append(inputs[0], cache_state, window_size=int(node.attrs.get("window_size", 0)), sink_size=int(node.attrs.get("sink_size", 0))) + return cache_state + + require_len(node, inputs, 2) + require_cache_state_parent(node, 1, "cactus.kv_cache_state") + context.graph.kv_cache_append(inputs[0], inputs[1], window_size=int(node.attrs.get("window_size", 0)), sink_size=int(node.attrs.get("sink_size", 0))) + return inputs[1] + + if target == "cactus.conv_cache_append": + require_len(node, inputs, 2) + require_cache_state_parent(node, 1, "cactus.conv_cache_state") + window = context.graph.conv_cache_append(inputs[0], inputs[1]) + return conv_cache_window_to_ir_layout(context, node, window) + + if target == "cactus.conv_cache_initialize": + if len(inputs) == 1: + cache_state = create_cache_state_for_cache_output(context, node) + rows = conv_cache_rows_for_native(context, node, inputs[0]) + context.graph.conv_cache_initialize(rows, cache_state) + return cache_state + + require_len(node, inputs, 2) + rows = conv_cache_rows_for_native(context, node, inputs[0]) + context.graph.conv_cache_initialize(rows, inputs[1]) + return inputs[1] + + if target == "cactus.recurrent_cache_write": + require_len(node, inputs, 2) + return context.graph.recurrent_cache_write(inputs[0], inputs[1]) + + if target == "cactus.kv_cache_state": + return context.graph.kv_cache_state( + required_int_attr(node, "max_seq_len"), + required_int_attr(node, "num_kv_heads"), + required_int_attr(node, "head_dim"), + window_size=int(node.attrs.get("window_size", 0)), + sink_size=int(node.attrs.get("sink_size", 0)), + num_slots=int(node.attrs.get("num_slots", 1)), + ) + + if target == "cactus.conv_cache_state": + return context.graph.conv_cache_state(required_int_attr(node, "window_size"), required_int_attr(node, "hidden_dim")) + + if target == "cactus.recurrent_cache_state": + return context.graph.recurrent_cache_state(shape_attr(node), dtype=cactus_precision(context.graph, tensor_dtype(node))) + + raise UnsupportedLoweringError(f"{node.name}: unsupported cache target {target}") + + +def require_cache_state_parent(node: IRModels.Node, parent_index: int, expected_target: str) -> None: + if parent_index >= len(node.parents): + raise UnsupportedLoweringError(f"{node.name}: {node.target} missing cache parent {parent_index}") + + parent = node.parents[parent_index] + + if parent.target == expected_target: + return + + if expected_target == "cactus.conv_cache_state" and parent.cache is not None and parent.cache.kind == FModels.CacheKind.CONV: + return + + if expected_target == "cactus.kv_cache_state" and parent.cache is not None and parent.cache.kind == FModels.CacheKind.KV: + return + + raise UnsupportedLoweringError( + f"{node.name}: {node.target} requires parent {parent_index} to be {expected_target}; got {parent.target}. " + "Raw HF cache tensors need a cache-state bridge before native cache append lowering." + ) + + +def require_cache_annotation(node: IRModels.Node) -> IRModels.CacheAnnotation: + if node.cache is None: + raise UnsupportedLoweringError(f"{node.name}: cache lowering requires cache annotation metadata") + + return node.cache + + +def required_cache_int(value: int | None, node: IRModels.Node, name: str) -> int: + if value is None: + raise UnsupportedLoweringError(f"{node.name}: cache lowering missing {name}") + + return int(value) + + +def kv_cache_capacity(context: models.GenerationContext, annotation: IRModels.CacheAnnotation) -> int: + capacity = int(annotation.sequence_length or 1) + + if context.component.ir_graph.task == "decode_with_cache": + capacity += 1 + + return max(capacity, 1) + + +def create_cache_state_for_cache_output(context: models.GenerationContext, node: IRModels.Node) -> Any: + annotation = require_cache_annotation(node) + + if annotation.kind == FModels.CacheKind.KV: + cache_state = context.graph.kv_cache_state( + kv_cache_capacity(context, annotation), + required_cache_int(annotation.num_kv_heads, node, "num_kv_heads"), + required_cache_int(annotation.head_dim, node, "head_dim"), + window_size=int(annotation.window_size or 0), + sink_size=0, + num_slots=1, + ) + record_cache_state_binding(context, node, cache_state, annotation) + return cache_state + + if annotation.kind == FModels.CacheKind.CONV: + cache_state = context.graph.conv_cache_state( + required_cache_int(annotation.window_size, node, "window_size"), + required_cache_int(annotation.hidden_dim, node, "hidden_dim"), + ) + record_cache_state_binding(context, node, cache_state, annotation) + return cache_state + + raise UnsupportedLoweringError(f"{node.name}: unsupported cache output kind {annotation.kind}") + + +def record_cache_state_binding(context: models.GenerationContext, node: IRModels.Node, cache_state: Any, annotation: IRModels.CacheAnnotation) -> None: + node_id = models.tensor_node_id(cache_state) + + if node_id is None: + raise UnsupportedLoweringError(f"{node.name}: cache state tensor does not expose a Cactus node id") + + context.component.add_cache_state_binding(RPModels.cache_state_binding_from_annotation(annotation, node_id)) + + +def conv_cache_rows_for_native(context: models.GenerationContext, node: IRModels.Node, rows: Any) -> Any: + rows_shape = meta_shape(node.parents[0]) if node.parents else output_shape(node) + + if len(rows_shape) == 3: + batch, hidden_dim, window_size = rows_shape + + if batch != 1: + raise UnsupportedLoweringError(f"{node.name}: conv cache initialize currently supports batch size 1") + + return context.graph.reshape(context.graph.permute(rows, (0, 2, 1)), (int(window_size), int(hidden_dim))) + + if len(rows_shape) == 2: + return rows + + raise UnsupportedLoweringError(f"{node.name}: conv cache initialize rows must be rank 2 or 3, got shape {rows_shape}") + + +def conv_cache_window_to_ir_layout(context: models.GenerationContext, node: IRModels.Node, window: Any) -> Any: + target_shape = output_shape(node) + + if len(target_shape) != 3: + return window + + batch, hidden_dim, window_size = target_shape + + if batch != 1: + raise UnsupportedLoweringError(f"{node.name}: conv cache append layout bridge currently supports batch size 1") + + return context.graph.reshape(context.graph.transpose(window), (int(batch), int(hidden_dim), int(window_size))) diff --git a/python/cactus/transpiler/Generator/lowering_constant_ops.py b/python/cactus/transpiler/Generator/lowering_constant_ops.py new file mode 100644 index 000000000..2d9214f8e --- /dev/null +++ b/python/cactus/transpiler/Generator/lowering_constant_ops.py @@ -0,0 +1,253 @@ +from __future__ import annotations + +from typing import Any + +from . import models +from .errors import UnsupportedLoweringError +from .lowering_utils import * +from ..IR import models as IRModels + + +def lower_constant_pad_nd(context: models.GenerationContext, node: IRModels.Node) -> Any: + inputs = require_input_count(context, node, 1) + pads = node.attrs.get("pad", node.attrs.get("arg_1")) + value = scalar_attr(node, "value") + + if value is None: + value = scalar_attr(node, "arg_2") + + if value is None: + value = 0.0 + + if not isinstance(pads, (list, tuple)) or len(pads) % 2 != 0: + raise UnsupportedLoweringError(f"{node.name}: constant_pad_nd lowering requires an even pad list") + + result = inputs[0] + parent_shape = meta_shape(node.parents[0]) + current_shape = list(parent_shape) + rank = len(parent_shape) + positive_pads = [0 for _ in pads] + + for pair_index in range(0, len(pads), 2): + left_pad = int(pads[pair_index]) + right_pad = int(pads[pair_index + 1]) + positive_pads[pair_index] = max(left_pad, 0) + positive_pads[pair_index + 1] = max(right_pad, 0) + + if left_pad >= 0 and right_pad >= 0: + continue + + axis = rank - 1 - (pair_index // 2) + + if axis < 0: + raise UnsupportedLoweringError(f"{node.name}: pad rank exceeds input rank") + + axis_size = current_shape[axis] + + if not isinstance(axis_size, int): + raise UnsupportedLoweringError(f"{node.name}: cannot crop symbolic pad dimension {axis_size}") + + start = max(-left_pad, 0) + end_crop = max(-right_pad, 0) + length = axis_size - start - end_crop + + if length < 0: + raise UnsupportedLoweringError(f"{node.name}: pad crop removes more than the full axis") + + result = context.graph.slice(result, axis, start, length=length) + current_shape[axis] = length + + if any(positive_pads): + return context.graph.pad(result, positive_pads, float(value)) + + return result + + +def lower_constant_input(context: models.GenerationContext, node: IRModels.Node) -> Any: + if node.target == "aten.full_like.default": + return lower_full_like(context, node) + + constant = deterministic_constant_values(node) + + if constant is not None: + values, shape = constant + dtype = cactus_precision(context.graph, tensor_dtype(node)) + tensor = context.graph.input(shape, dtype=dtype) + path = write_constant_tensor(context, node, values, shape, dtype) + node_id = models.tensor_node_id(tensor) + + if node_id is None: + raise UnsupportedLoweringError(f"{node.name}: generated constant does not expose a Cactus node id") + + context.component.add_weight_binding( + models.WeightBinding( + placeholder=node.name, + source_target=node.target, + node_id=node_id, + path=path, + output_name=path, + source_name=node.target, + value_id=node.name, + precision=precision_name(context.graph, dtype), + component=context.component.name, + binding_kind="generated_constant", + ) + ) + return tensor + + context.component.warn(f"{node.name}: constant-producing op {node.target} lowered as graph input") + shape, dynamic_dims = graph_input_shape(node) + dtype = cactus_precision(context.graph, tensor_dtype(node)) + tensor = context.graph.input(shape, dtype=dtype, dynamic_dims=dynamic_dims if any(dynamic_dims) else None) + context.component.add_runtime_input(tensor, node.name) + return tensor + + +def lower_full_like(context: models.GenerationContext, node: IRModels.Node) -> Any: + inputs = require_input_count(context, node, 1) + fill_value = scalar_attr(node, "fill_value") + + if fill_value is None: + fill_value = scalar_attr(node, "arg_1") + + if fill_value is None: + fill_value = 0.0 + + output = context.graph.scalar_multiply(inputs[0], 0.0) + fill_float = float(fill_value) + + if fill_float != 0.0: + output = context.graph.scalar_add(output, fill_float) + + return cast_to_precision(context, output, cactus_precision(context.graph, tensor_dtype(node))) + + +def deterministic_constant_values(node: IRModels.Node) -> tuple[list[float], tuple[int, ...]] | None: + shape = concrete_shape(meta_shape(node)) + + if node.target.startswith("aten.arange"): + start = numeric_attr(node, "start", "arg_0", default=0) + end = numeric_attr(node, "end", "arg_1") + step = numeric_attr(node, "step", "arg_2", default=1) + + if end is None or step in {None, 0}: + return None + + values = arange_values(float(start), float(end), float(step)) + return values, (len(values),) + + if node.target == "aten.scalar_tensor.default": + value = numeric_attr(node, "value", "arg_0", default=0) + return [float(value or 0)], (1,) + + if node.target in {"aten.full.default", "aten.full_like.default"}: + if shape is None: + raw_size = node.attrs.get("size", node.attrs.get("arg_0")) + shape = tuple(int(dim) for dim in raw_size) if isinstance(raw_size, (list, tuple)) else None + + if shape is None: + return None + + fill_value = numeric_attr(node, "fill_value", "arg_1", default=0) + count = math.prod(shape) if shape else 1 + return [float(fill_value or 0)] * int(count), tuple(shape or (1,)) + + return None + + +def arange_values(start: float, end: float, step: float) -> list[float]: + values: list[float] = [] + current = start + + if step > 0: + while current < end: + values.append(current) + current += step + else: + while current > end: + values.append(current) + current += step + + return values + + +def lower_neg(context: models.GenerationContext, node: IRModels.Node) -> Any: + inputs = require_input_count(context, node, 1) + return context.graph.scalar_multiply(inputs[0], -1.0) + + +def lower_cumsum(context: models.GenerationContext, node: IRModels.Node) -> Any: + inputs = require_input_count(context, node, 1) + axis = axis_attr(node) + + if axis is None: + raise UnsupportedLoweringError(f"{node.name}: cumsum lowering missing axis/dim attr") + + return context.graph.cumsum(inputs[0], normalize_dim(axis, len(meta_shape(node.parents[0])))) + + +def lower_softmax(context: models.GenerationContext, node: IRModels.Node) -> Any: + inputs = require_input_count(context, node, 1) + axis = axis_attr(node, default=-1) + + if axis is None: + raise UnsupportedLoweringError(f"{node.name}: softmax lowering missing axis/dim attr") + + return context.graph.softmax(fp16_tensor(context, inputs[0]), axis=normalize_dim(axis, len(meta_shape(node.parents[0])))) + + +def lower_topk(context: models.GenerationContext, node: IRModels.Node) -> Any: + inputs = require_input_count(context, node, 1) + k = node.attrs.get("k") + + if k is None: + raise UnsupportedLoweringError(f"{node.name}: topk lowering missing k attr") + + return context.graph.topk(inputs[0], int(k)) + + +def lower_gather(context: models.GenerationContext, node: IRModels.Node) -> Any: + inputs = require_input_count(context, node, 2) + axis = axis_attr(node, default=0) + return context.graph.gather(inputs[0], inputs[1], axis=axis if axis is not None else 0) + + +def lower_embedding(context: models.GenerationContext, node: IRModels.Node) -> Any: + inputs = require_input_count(context, node, 2) + return context.graph.embedding_from_tensor(inputs[0], inputs[1]) + + +def lower_clamp(context: models.GenerationContext, node: IRModels.Node) -> Any: + inputs = require_input_count(context, node, 1) + lo = scalar_weight_bound_value(context, node.parents[1]) if len(node.parents) > 1 else None + hi = scalar_weight_bound_value(context, node.parents[2]) if len(node.parents) > 2 else None + + if lo is None: + lo = inputs[1] if len(inputs) > 1 else node.attrs.get("lo", node.attrs.get("min", node.attrs.get("arg_1"))) + + if hi is None: + hi = inputs[2] if len(inputs) > 2 else node.attrs.get("hi", node.attrs.get("max", node.attrs.get("arg_2"))) + + if lo is None and hi is None: + return inputs[0] + + result = inputs[0] + + if isinstance(lo, (int, float)) and isinstance(hi, (int, float)): + return context.graph.clamp(result, float(lo), float(hi)) + + if lo is not None: + if isinstance(lo, (int, float)): + result = context.graph.clamp(result, float(lo), math.inf) + else: + lo = cast_to_precision(context, lo, getattr(result, "dtype", context.graph.FP16)) + result = context.graph.where(context.graph.greater_equal(result, lo), result, lo) + + if hi is not None: + if isinstance(hi, (int, float)): + result = context.graph.clamp(result, -math.inf, float(hi)) + else: + hi = cast_to_precision(context, hi, getattr(result, "dtype", context.graph.FP16)) + result = context.graph.where(context.graph.less_equal(result, hi), result, hi) + + return result diff --git a/python/cactus/transpiler/Generator/lowering_core.py b/python/cactus/transpiler/Generator/lowering_core.py new file mode 100644 index 000000000..78e279fa6 --- /dev/null +++ b/python/cactus/transpiler/Generator/lowering_core.py @@ -0,0 +1,315 @@ +from __future__ import annotations + +from collections.abc import Mapping +from typing import Any + +from . import constants +from . import models +from .errors import UnsupportedLoweringError +from .lowering_basic_ops import * +from .lowering_cache import cache_attention_generation_plan, lower_cache, lower_cache_placeholder, should_lower_cache_placeholder_as_state, lower_attention +from .lowering_constant_ops import * +from .lowering_nn_ops import lower_conv, lower_norm +from .lowering_special_ops import lower_moe, lower_special_cactus, lower_unsupported_semantic +from .lowering_utils import * +from .lowering_weights import bind_weight_placeholder, lower_lfm_grouped_moe_placeholder +from ..IR import models as IRModels + + +def create_cactus_graph() -> Any: + try: + from cactus import Graph + except ModuleNotFoundError: + from bindings.cactus import Graph + + return Graph() + + +def component_weight_resolver( + component: models.ComponentGraph, + config: models.GeneratorConfig, +) -> models.WeightResolver | None: + if config.weights_dir is None: + return None + + resolver = models.WeightResolver.from_graph( + component.ir_graph, + config.weights_dir, + config.weights_manifest_path, + ) + + if not resolver.records: + component.warn(f"{component.name}: no converted weight records found in {config.weights_dir}") + + return resolver + + +def build_lowering_rules( + extra_lowerings: Mapping[str, models.LoweringRule | models.LoweringFn] | None = None, +) -> dict[str, models.LoweringRule]: + rules: dict[str, models.LoweringRule] = {} + + add_rules(rules, constants.BINARY_TARGETS, lower_binary) + add_rules(rules, constants.SCALAR_TARGETS, lower_scalar) + add_rules(rules, constants.UNARY_TARGETS, lower_unary) + add_rules(rules, constants.LOG1P_TARGETS, lower_log1p) + add_rules(rules, constants.POW_TARGETS, lower_pow) + add_rules(rules, constants.REDUCE_TARGETS, lower_reduce) + add_rules(rules, constants.SHAPE_TARGETS, lower_shape) + add_rules(rules, constants.EXPAND_TARGETS, lower_expand) + add_rules(rules, constants.UNSQUEEZE_TARGETS, lower_unsqueeze) + add_rules(rules, constants.FLATTEN_TARGETS, lower_flatten) + add_rules(rules, constants.REPEAT_TARGETS, lower_repeat) + add_rules(rules, constants.TRANSPOSE_TARGETS, lower_transpose) + add_rules(rules, constants.SLICE_TARGETS, lower_slice) + add_rules(rules, constants.INDEX_TARGETS, lower_index) + add_rules(rules, constants.WHERE_TARGETS, lower_where) + add_rules(rules, constants.UNFOLD_TARGETS, lower_unfold) + add_rules(rules, constants.CAT_TARGETS, lower_cat) + add_rules(rules, constants.MATMUL_TARGETS, lower_matmul) + add_rules(rules, constants.ADDM_CONST_TARGETS, lower_addmm) + add_rules(rules, constants.SPLIT_TARGETS, lower_split) + add_rules(rules, constants.GETITEM_TARGETS, lower_getitem) + add_rules(rules, constants.TO_COPY_TARGETS, lower_to_copy) + add_rules(rules, constants.PASS_THROUGH_TARGETS, lower_pass_through) + add_rules(rules, constants.COPY_TARGETS, lower_copy) + add_rules(rules, constants.PAD_TARGETS, lower_constant_pad_nd) + add_rules(rules, constants.CONSTANT_INPUT_TARGETS, lower_constant_input) + add_rules(rules, constants.NEG_TARGETS, lower_neg) + add_rules(rules, constants.CUMSUM_TARGETS, lower_cumsum) + add_rules(rules, constants.SOFTMAX_TARGETS, lower_softmax) + add_rules(rules, constants.TOPK_TARGETS, lower_topk) + add_rules(rules, constants.GATHER_TARGETS, lower_gather) + add_rules(rules, constants.EMBEDDING_TARGETS, lower_embedding) + add_rules(rules, constants.CLAMP_TARGETS, lower_clamp) + add_rules(rules, constants.NORM_TARGETS, lower_norm) + add_rules(rules, constants.CONV_TARGETS, lower_conv) + add_rules(rules, constants.ATTENTION_TARGETS, lower_attention) + add_rules(rules, constants.CACHE_TARGETS, lower_cache) + add_rules(rules, constants.MOE_TARGETS, lower_moe) + add_rules(rules, constants.SPECIAL_CACTUS_TARGETS, lower_special_cactus) + add_rules(rules, constants.UNSUPPORTED_SEMANTIC_TARGETS, lower_unsupported_semantic) + + if extra_lowerings: + for target, rule in extra_lowerings.items(): + if isinstance(rule, models.LoweringRule): + rules[target] = rule + else: + rules[target] = models.LoweringRule(target, rule) + + return rules + + +def add_rules(rules: dict[str, models.LoweringRule], targets: Any, lower: models.LoweringFn) -> None: + target_names = targets.keys() if isinstance(targets, dict) else targets + + for target in target_names: + rules[target] = models.LoweringRule(target, lower) + + +def lower_component( + component: models.ComponentGraph, + config: models.GeneratorConfig, + lowering_rules: dict[str, models.LoweringRule], +) -> models.ComponentGraph: + component.graph = create_cactus_graph() + component.weight_resolver = component_weight_resolver(component, config) + context = models.GenerationContext(component=component, config=config, lowerings=lowering_rules) + ( + context.skip_node_names, + context.cache_state_placeholder_names, + context.prefill_cache_cat_annotations, + ) = cache_attention_generation_plan(component.ir_graph) + + for node in IRModels.topological_sort(component.ir_graph): + try: + lower_node(context, node) + except UnsupportedLoweringError as e: + component.mark_unsupported(node) + component.warn(str(e)) + + if config.strict and not config.allow_unsupported_ops: + raise + except KeyError as e: + component.mark_unsupported(node) + component.warn(f"{node.name}: missing lowered input {e}") + + if config.strict and not config.allow_unsupported_ops: + raise UnsupportedLoweringError(f"{node.name}: missing lowered input {e}") from e + + component.save() + return component + + +def lower_node(context: models.GenerationContext, node: IRModels.Node) -> None: + if node.name in context.skip_node_names: + return + + if node.is_placeholder: + context.bind(node, lower_placeholder(context, node)) + return + + if node.is_output: + context.bind(node, lower_output(context, node)) + return + + rule = context.lowering_for(node) + + if rule is None: + raise UnsupportedLoweringError(f"{node.name}: no lowering rule for {node.target}") + + result = rule.lower(context, node) + + if result is not None: + context.bind(node, result) + + +def lower_placeholder(context: models.GenerationContext, node: IRModels.Node) -> Any: + if should_lower_cache_placeholder_as_state(context, node): + return lower_cache_placeholder(context, node) + + if node.value_kind in constants.WEIGHT_VALUE_KINDS: + grouped_moe_weights = lower_lfm_grouped_moe_placeholder(context, node) + if grouped_moe_weights is not None: + return grouped_moe_weights + + shape, dynamic_dims = graph_input_shape(node) + dtype = placeholder_precision(context, node) + tensor = context.graph.input(shape, dtype=dtype, dynamic_dims=dynamic_dims if any(dynamic_dims) else None) + logical_name = logical_input_name(context.component.ir_graph, node) + + if node.value_kind in constants.WEIGHT_VALUE_KINDS: + bind_weight_placeholder(context, node, tensor, logical_name) + return tensor + + context.component.add_runtime_input(tensor, logical_name) + return tensor + + +def placeholder_precision(context: models.GenerationContext, node: IRModels.Node) -> int: + if node.value_kind in constants.WEIGHT_VALUE_KINDS and context.component.weight_resolver is not None: + record = context.component.weight_resolver.resolve(node.name) + if record is not None and record.precision and hasattr(context.graph, record.precision): + return int(getattr(context.graph, record.precision)) + + logical_name = logical_input_name(context.component.ir_graph, node) + if is_fp16_runtime_input(logical_name) and str(tensor_dtype(node)) in {"torch.float32", "torch.float"}: + return int(context.graph.FP16) + + return cactus_precision(context.graph, tensor_dtype(node)) + + +def is_fp16_runtime_input(logical_name: str) -> bool: + return logical_name in constants.FP16_RUNTIME_INPUTS or logical_name.startswith(("cross_k_", "cross_v_")) + + +def lower_output(context: models.GenerationContext, node: IRModels.Node) -> Any: + refs = IRModels.extract_node_refs((node.args, node.kwargs)) + values = tuple(context.require(ref) for ref in refs) + output_index = 0 + + for ref, value in zip(refs, values): + if isinstance(value, (tuple, list)): + for item in value: + context.component.add_output(item, logical_output_name(context.component.ir_graph, output_index, ref)) + output_index += 1 + else: + context.component.add_output(value, logical_output_name(context.component.ir_graph, output_index, ref)) + output_index += 1 + + if len(values) == 1: + return values[0] + + return values + + +def logical_input_name(graph: IRModels.Graph, node: IRModels.Node) -> str: + metadata_name = logical_name_from_metadata(node.ir_metadata, ("logical_input", "logical_name", "runtime_input")) + + if metadata_name: + return metadata_name + + spec_name = logical_input_name_from_specs(graph, node.name) + + if spec_name: + return spec_name + + if node.cache is not None: + return cache_logical_name(node) + + target_name = clean_logical_name(node.target) + + if target_name is not None: + return target_name + + return node.name + + +def logical_output_name(graph: IRModels.Graph, output_index: int, source_ref: str | None = None) -> str: + source = graph.nodes_map.get(source_ref) if source_ref is not None else None + + if source is not None: + metadata_name = logical_name_from_metadata(source.ir_metadata, ("logical_output", "logical_name", "runtime_output")) + + if metadata_name: + return metadata_name + + if source.cache is not None: + return cache_logical_name(source) + + if "logits" in source.name or "lm_head" in source.target: + return "logits" + + spec_name = logical_output_name_from_specs(graph, output_index) + + if spec_name: + return spec_name + + return f"output_{output_index}" + + +def logical_input_name_from_specs(graph: IRModels.Graph, node_name: str) -> str | None: + for spec in graph.input_specs: + if spec.arg_name != node_name: + continue + + return clean_logical_name(spec.target) or clean_logical_name(spec.arg_name) + + return None + + +def logical_output_name_from_specs(graph: IRModels.Graph, output_index: int) -> str | None: + if output_index < 0 or output_index >= len(graph.output_specs): + return None + + spec = graph.output_specs[output_index] + return clean_logical_name(spec.target) or clean_logical_name(spec.arg_name) + + +def logical_name_from_metadata(metadata: dict[str, Any], keys: tuple[str, ...]) -> str | None: + for key in keys: + value = clean_logical_name(metadata.get(key)) + + if value is not None: + return value + + return None + + +def clean_logical_name(value: Any) -> str | None: + if value is None: + return None + + name = str(value) + + if not name or name == "None": + return None + + return name + + +def cache_logical_name(node: IRModels.Node) -> str: + if node.cache is not None and node.cache.tensor_index is not None: + return f"past_key_values_{node.cache.tensor_index}" + + return node.name diff --git a/python/cactus/transpiler/Generator/lowering_nn_ops.py b/python/cactus/transpiler/Generator/lowering_nn_ops.py new file mode 100644 index 000000000..303c5bebf --- /dev/null +++ b/python/cactus/transpiler/Generator/lowering_nn_ops.py @@ -0,0 +1,240 @@ +from __future__ import annotations + +from typing import Any + +from . import constants +from . import models +from .errors import UnsupportedLoweringError +from .lowering_utils import * +from ..IR import models as IRModels + + +def lower_norm(context: models.GenerationContext, node: IRModels.Node) -> Any: + inputs = context.inputs_for(node) + target = node.target + + if target == "cactus.rms_norm": + if len(inputs) not in {1, 2}: + raise unsupported_arity(node, len(inputs), "activation plus optional weight") + + weight = inputs[1] if len(inputs) == 2 else None + return lower_rms_norm(context, node, inputs[0], weight) + + if target in {"cactus.layernorm", "cactus.layer_norm", "aten.native_layer_norm.default", "aten.layer_norm.default"}: + require_len(node, inputs, 2) + bias = inputs[2] if len(inputs) > 2 else None + return context.graph.layernorm(inputs[0], inputs[1], bias=bias, eps=epsilon_attr(node)) + + if target in {"cactus.groupnorm", "cactus.group_norm", "aten.native_group_norm.default", "aten.group_norm.default"}: + require_len(node, inputs, 3) + return context.graph.groupnorm(inputs[0], inputs[1], inputs[2], node.attrs.get("num_groups", 1), eps=epsilon_attr(node)) + + if target in {"cactus.batchnorm", "cactus.batch_norm", "aten.native_batch_norm.default", "aten.batch_norm.default"}: + require_len(node, inputs, 5) + return context.graph.batchnorm(inputs[0], inputs[1], inputs[2], inputs[3], inputs[4], axis=int(node.attrs.get("axis", 1)), eps=epsilon_attr(node)) + + raise UnsupportedLoweringError(f"{node.name}: unsupported norm target {target}") + + +def lower_rms_norm(context: models.GenerationContext, node: IRModels.Node, x: Any, weight: Any | None = None) -> Any: + x = cast_to_precision(context, x, context.graph.FP16) + declared_input_shape = meta_shape(node.parents[0]) if node.parents else meta_shape(node) + input_shape = concrete_shape(declared_input_shape) + + if input_shape is None: + actual_shape = tuple(getattr(x, "shape", ())) + + if actual_shape and concrete_dim(declared_input_shape[-1] if declared_input_shape else None) is not None: + input_shape = actual_shape + + if input_shape is None: + raise UnsupportedLoweringError(f"{node.name}: cactus.rms_norm requires concrete input shape") + + if weight is None: + weight = rms_norm_unit_weight(context, int(input_shape[-1])) + else: + weight = cast_to_precision(context, weight, context.graph.FP16) + + if len(input_shape) == 2: + return context.graph.rms_norm(x, weight, eps=epsilon_attr(node)) + + if len(input_shape) == 1: + return context.graph.reshape( + context.graph.rms_norm(context.graph.reshape(x, (1, input_shape[0])), weight, eps=epsilon_attr(node)), + input_shape, + ) + + rows = element_count(input_shape[:-1]) + + if rows is None: + raise UnsupportedLoweringError(f"{node.name}: cactus.rms_norm requires concrete leading dimensions") + + normalized = context.graph.rms_norm(context.graph.reshape(x, (rows, input_shape[-1])), weight, eps=epsilon_attr(node)) + return context.graph.reshape(normalized, input_shape) + + +def rms_norm_unit_weight(context: models.GenerationContext, hidden_dim: int) -> Any: + key = f"__embedded_rms_norm_unit_weight_{hidden_dim}" + existing = context.lookup(key) + + if existing is not None: + return existing + + weight = context.graph.input((hidden_dim,), dtype=context.graph.FP16) + context.graph.set_input(weight, np.ones((hidden_dim,), dtype=np.float16)) + context.graph.mark_embedded_input(weight) + context.values[key] = weight + return weight + + +def lower_conv(context: models.GenerationContext, node: IRModels.Node) -> Any: + inputs = context.inputs_for(node) + target = node.target + + if target.startswith(constants.CACTUS_TARGET_PREFIX): + method = target.removeprefix(constants.CACTUS_TARGET_PREFIX) + return lower_named_conv(context, node, method, inputs) + + require_len(node, inputs, 2) + bias = inputs[2] if len(inputs) > 2 else None + rank = output_rank(node) + + if target == "aten.conv1d.default" or rank == 3: + return lower_generic_conv1d(context, node, inputs, bias) + + if target == "aten.conv2d.default" or rank == 4: + specialized = lower_specialized_conv2d(context, node, inputs, bias) + + if specialized is not None: + return specialized + + return context.graph.conv2d( + inputs[0], + inputs[1], + bias=bias, + stride=first_int(node.attrs.get("stride"), 1), + padding=first_int(node.attrs.get("padding"), 0), + dilation=first_int(node.attrs.get("dilation"), 1), + groups=int(node.attrs.get("groups", 1)), + ) + + raise UnsupportedLoweringError(f"{node.name}: cannot infer conv rank for {target}") + + +def lower_named_conv(context: models.GenerationContext, node: IRModels.Node, method: str, inputs: tuple[Any, ...]) -> Any: + require_len(node, inputs, 2) + bias = inputs[2] if len(inputs) > 2 else None + + if method == "conv1d": + return lower_generic_conv1d(context, node, inputs, bias) + + if method == "conv1d_causal": + x = inputs[0] + + if node.attrs.get("layout") == "batch_hidden_sequence": + x = context.graph.permute(x, (0, 2, 1)) + + output = context.graph.conv1d_causal(x, inputs[1], kernel_size=int(node.attrs.get("kernel_size", 0)), dilation=first_int(node.attrs.get("dilation"), 1)) + + if node.attrs.get("layout") == "batch_hidden_sequence": + return context.graph.permute(output, (0, 2, 1)) + + return output + + if method == "conv1d_k3": + require_plain_conv1d_attrs(node, method) + return context.graph.conv1d_k3(inputs[0], inputs[1], stride=first_int(node.attrs.get("stride"), 1)) + + if method == "conv1d_k7s3": + require_len(node, inputs, 3) + return context.graph.conv1d_k7s3(inputs[0], inputs[1], inputs[2]) + + if method in {"conv1d_same_depthwise_k9", "conv1d_pointwise", "conv2d_k3s2p1", "conv2d_depthwise_k3s2p1", "conv2d_pointwise_1x1", "conv2d_k3s1p1"}: + return getattr(context.graph, method)(inputs[0], inputs[1], bias=bias) + + if method == "conv2d": + specialized = lower_specialized_conv2d(context, node, inputs, bias) + + if specialized is not None: + return specialized + + return context.graph.conv2d( + inputs[0], + inputs[1], + bias=bias, + stride=first_int(node.attrs.get("stride"), 1), + padding=first_int(node.attrs.get("padding"), 0), + dilation=first_int(node.attrs.get("dilation"), 1), + groups=int(node.attrs.get("groups", 1)), + ) + + raise UnsupportedLoweringError(f"{node.name}: unsupported Cactus conv method {method}") + + +def lower_generic_conv1d(context: models.GenerationContext, node: IRModels.Node, inputs: tuple[Any, ...], bias: Any | None) -> Any: + stride = first_int(node.attrs.get("stride"), 1) + padding = first_int(node.attrs.get("padding"), 0) + dilation = first_int(node.attrs.get("dilation"), 1) + groups = first_int(node.attrs.get("groups"), 1) + input_shape = meta_shape(node.parents[0]) if node.parents else () + weight_shape = meta_shape(node.parents[1]) if len(node.parents) > 1 else () + + if ( + padding == 0 + and dilation == 1 + and len(input_shape) == 3 + and len(weight_shape) == 3 + and groups == input_shape[1] + and weight_shape[0] == input_shape[1] + and weight_shape[1] == 1 + ): + return context.graph.conv1d_depthwise(inputs[0], inputs[1], bias=bias, stride=stride) + + if dilation != 1 or groups != 1: + require_plain_conv1d_attrs(node, "conv1d") + + x = inputs[0] + + if padding > 0: + x = context.graph.pad(x, (padding, padding), 0.0) + + return context.graph.conv1d(x, inputs[1], bias=bias, stride=stride) + + +def lower_specialized_conv2d(context: models.GenerationContext, node: IRModels.Node, inputs: tuple[Any, ...], bias: Any | None) -> Any | None: + weight_shape = meta_shape(node.parents[1]) if len(node.parents) > 1 else () + stride = tuple_int_values(node.attrs.get("stride"), 1) + padding = tuple_int_values(node.attrs.get("padding"), 0) + dilation = tuple_int_values(node.attrs.get("dilation"), 1) + groups = first_int(node.attrs.get("groups"), 1) + + if len(weight_shape) != 4 or dilation != (1, 1): + return None + + kernel = tuple(int(dim) for dim in weight_shape[2:]) + + if kernel == (3, 3) and stride == (2, 2) and padding == (1, 1): + if groups != 1 and weight_shape[1] == 1: + return context.graph.conv2d_depthwise_k3s2p1(inputs[0], inputs[1], bias=bias) + + return context.graph.conv2d_k3s2p1(inputs[0], inputs[1], bias=bias) + + if kernel == (3, 3) and stride == (1, 1) and padding == (1, 1) and groups == 1: + return context.graph.conv2d_k3s1p1(inputs[0], inputs[1], bias=bias) + + if kernel == (1, 1) and stride == (1, 1) and padding == (0, 0) and groups == 1: + return context.graph.conv2d_pointwise_1x1(inputs[0], inputs[1], bias=bias) + + return None + + +def require_plain_conv1d_attrs(node: IRModels.Node, method: str) -> None: + padding = first_int(node.attrs.get("padding"), 0) + dilation = first_int(node.attrs.get("dilation"), 1) + groups = first_int(node.attrs.get("groups"), 1) + + if padding != 0 or dilation != 1 or groups != 1: + raise UnsupportedLoweringError( + f"{node.name}: {method} lowering only supports padding=0, dilation=1, groups=1; " + f"got padding={padding}, dilation={dilation}, groups={groups}" + ) diff --git a/python/cactus/transpiler/Generator/lowering_special_ops.py b/python/cactus/transpiler/Generator/lowering_special_ops.py new file mode 100644 index 000000000..a7de111a5 --- /dev/null +++ b/python/cactus/transpiler/Generator/lowering_special_ops.py @@ -0,0 +1,358 @@ +from __future__ import annotations + +from typing import Any + +from . import models +from .errors import UnsupportedLoweringError +from .lowering_utils import * +from ..IR import models as IRModels + + +def lower_moe(context: models.GenerationContext, node: IRModels.Node) -> Any: + inputs = context.inputs_for(node) + + if node.target == "cactus.dense_mlp_tq_fused": + require_len(node, inputs, 4) + hidden = cast_to_precision(context, inputs[0], context.graph.FP16) + return context.graph.dense_mlp_tq_fused( + hidden, + inputs[1], + inputs[2], + inputs[3], + product_scale=float(node.attrs.get("product_scale", 1.0)), + ) + + if node.target == "cactus.moe_layer_gated": + packed = lower_packed_lfm_moe_layer_gated(context, node, inputs) + if packed is not None: + return packed + + if len(inputs) < 6: + raise UnsupportedLoweringError( + f"{node.name}: cactus.moe_layer_gated requires hidden, routing_probs, topk_indices, w1, w3, and w2 inputs; " + "packed/grouped MoE lowering still needs a weight-unpacking policy" + ) + + return context.graph.moe_layer_gated( + inputs[0], + inputs[1], + cast_to_precision(context, inputs[2], context.graph.FP32), + ensure_tensor_sequence(inputs[3]), + ensure_tensor_sequence(inputs[4]), + ensure_tensor_sequence(inputs[5]), + required_int_attr(node, "num_experts"), + required_int_attr(node, "num_experts_per_tok"), + normalize_routing=bool(node.attrs.get("normalize_routing", True)), + epsilon=float(node.attrs.get("epsilon", 1e-6)), + routed_scaling_factor=float(node.attrs.get("routed_scaling_factor", 1.0)), + ) + + if node.target == "cactus.moe_layer_ungated": + if len(inputs) < 5: + raise unsupported_arity(node, len(inputs), "hidden, routing_probs, topk_indices, w1, and w2") + + return context.graph.moe_layer_ungated( + inputs[0], + inputs[1], + cast_to_precision(context, inputs[2], context.graph.FP32), + ensure_tensor_sequence(inputs[3]), + ensure_tensor_sequence(inputs[4]), + required_int_attr(node, "num_experts"), + required_int_attr(node, "num_experts_per_tok"), + normalize_routing=bool(node.attrs.get("normalize_routing", True)), + epsilon=float(node.attrs.get("epsilon", 1e-6)), + routed_scaling_factor=float(node.attrs.get("routed_scaling_factor", 1.0)), + ) + + raise UnsupportedLoweringError(f"{node.name}: unsupported MoE target {node.target}") + + +def lower_packed_lfm_moe_layer_gated( + context: models.GenerationContext, + node: IRModels.Node, + inputs: tuple[Any, ...], +) -> Any | None: + if len(inputs) != 5 or len(node.parents) != 5: + return None + + hidden_node, router_node, bias_node, gate_up_node, down_node = node.parents + gate_up_shape = concrete_shape(meta_shape(gate_up_node)) + down_shape = concrete_shape(meta_shape(down_node)) + num_experts = required_int_attr(node, "num_experts") + top_k = required_int_attr(node, "num_experts_per_tok") + + if gate_up_shape is None or down_shape is None: + return None + + if len(gate_up_shape) != 3 or len(down_shape) != 3: + return None + + if gate_up_shape[0] != num_experts or down_shape[0] != num_experts: + return None + + packed_intermediate, hidden_dim = gate_up_shape[1], gate_up_shape[2] + if packed_intermediate % 2 != 0: + raise UnsupportedLoweringError(f"{node.name}: packed gated MoE gate/up dimension must be even") + + intermediate_dim = packed_intermediate // 2 + + if down_shape[1:] != (hidden_dim, intermediate_dim): + raise UnsupportedLoweringError( + f"{node.name}: expected LFM packed down_proj shape " + f"[experts, hidden, intermediate], got {down_shape}" + ) + + hidden = moe_hidden_2d(context, node, inputs[0], hidden_dim) + router_weight = inputs[1] + expert_bias = cast_to_precision(context, inputs[2], context.graph.FP16) + gate_up_weight = inputs[3] + down_weight = inputs[4] + + router_logits = context.graph.matmul(hidden, router_weight, pretransposed_rhs=True) + router_logits = context.graph.add(router_logits, expert_bias) + routing_probs = context.graph.softmax(cast_to_precision(context, router_logits, context.graph.FP16), axis=-1) + topk_indices = context.graph.index(context.graph.topk(routing_probs, top_k), 0, axis=0) + topk_indices = cast_to_precision(context, topk_indices, context.graph.FP32) + bundled_weights = lfm_moe_weight_bundle_parts(gate_up_weight, down_weight, num_experts) + + if bundled_weights is None: + w1_weights, w3_weights, w2_weights = split_lfm_packed_moe_weights( + context, + gate_up_weight, + down_weight, + num_experts, + intermediate_dim, + hidden_dim, + ) + else: + w1_weights, w3_weights, w2_weights = bundled_weights + + return context.graph.moe_layer_gated( + hidden, + routing_probs, + topk_indices, + w1_weights, + w3_weights, + w2_weights, + num_experts, + top_k, + normalize_routing=bool(node.attrs.get("normalize_routing", True)), + epsilon=float(node.attrs.get("epsilon", 1e-6)), + routed_scaling_factor=float(node.attrs.get("routed_scaling_factor", 1.0)), + ) + + +def moe_hidden_2d(context: models.GenerationContext, node: IRModels.Node, hidden: Any, hidden_dim: int) -> Any: + hidden_shape = concrete_shape(meta_shape(node.parents[0])) if node.parents else None + + if hidden_shape is None: + return cast_to_precision(context, hidden, context.graph.FP16) + + if len(hidden_shape) == 2: + return cast_to_precision(context, hidden, context.graph.FP16) + + if len(hidden_shape) == 3 and hidden_shape[0] == 1 and hidden_shape[2] == hidden_dim: + return cast_to_precision(context, context.graph.reshape(hidden, (hidden_shape[1], hidden_dim)), context.graph.FP16) + + raise UnsupportedLoweringError(f"{node.name}: packed MoE hidden input must be [tokens, hidden] or [1, tokens, hidden]") + + +def lfm_moe_weight_bundle_parts( + gate_up_weight: Any, + down_weight: Any, + num_experts: int, +) -> tuple[tuple[Any, ...], tuple[Any, ...], tuple[Any, ...]] | None: + if not isinstance(gate_up_weight, models.LfmMoeWeightBundle): + return None + + if not isinstance(down_weight, models.LfmMoeWeightBundle): + return None + + if ( + len(gate_up_weight.w1_weights) != num_experts + or len(gate_up_weight.w3_weights) != num_experts + or len(down_weight.w2_weights) != num_experts + ): + return None + + return gate_up_weight.w1_weights, gate_up_weight.w3_weights, down_weight.w2_weights + + +def split_lfm_packed_moe_weights( + context: models.GenerationContext, + gate_up_weight: Any, + down_weight: Any, + num_experts: int, + intermediate_dim: int, + hidden_dim: int, +) -> tuple[tuple[Any, ...], tuple[Any, ...], tuple[Any, ...]]: + w1_weights = [] + w3_weights = [] + w2_weights = [] + + for expert_index in range(num_experts): + expert_gate_up = context.graph.reshape( + context.graph.slice(gate_up_weight, 0, expert_index, length=1), + (2 * intermediate_dim, hidden_dim), + ) + w1_weights.append(context.graph.slice(expert_gate_up, 0, 0, length=intermediate_dim)) + w3_weights.append(context.graph.slice(expert_gate_up, 0, intermediate_dim, length=intermediate_dim)) + + expert_down = context.graph.reshape( + context.graph.slice(down_weight, 0, expert_index, length=1), + (hidden_dim, intermediate_dim), + ) + w2_weights.append(context.graph.transpose(expert_down)) + + return tuple(w1_weights), tuple(w3_weights), tuple(w2_weights) + + +def lower_special_cactus(context: models.GenerationContext, node: IRModels.Node) -> Any: + inputs = context.inputs_for(node) + target = node.target + + if target == "cactus.rope": + require_len(node, inputs, 1) + return context.graph.rope( + inputs[0], + theta=float(attr_value(node, "theta", 10_000.0)), + position_offset=int(attr_value(node, "position_offset", 0)), + ) + + if target == "cactus.rope_gptj": + require_len(node, inputs, 1) + return context.graph.rope_gptj( + inputs[0], + theta=float(attr_value(node, "theta", 10_000.0)), + position_offset=int(attr_value(node, "position_offset", 0)), + rot_dim=int(attr_value(node, "rot_dim", 0)), + ) + + if target == "cactus.glu": + require_len(node, inputs, 1) + if meta_shape(node) == meta_shape(node.parents[0]): + return inputs[0] + return context.graph.glu(inputs[0], axis=axis_attr(node, default=-1)) + + if target == "cactus.lstm_cell": + require_len(node, inputs, 7) + return context.graph.lstm_cell(*inputs[:7]) + + if target == "cactus.gated_deltanet_decode": + require_len(node, inputs, 6) + return context.graph.gated_deltanet_decode(*inputs[:6], scale=float(node.attrs.get("scale", 1.0))) + + if target == "cactus.gated_deltanet_prefill": + require_len(node, inputs, 6) + return context.graph.gated_deltanet_prefill(*inputs[:6], chunk_size=int(node.attrs.get("chunk_size", 1)), scale=float(node.attrs.get("scale", 1.0))) + + if target == "cactus.lfm_short_conv_decode": + require_len(node, inputs, 2) + weight = lfm_short_conv_decode_weight(context, node, inputs[1]) + weight = cast_to_precision(context, weight, getattr(inputs[0], "dtype", context.graph.FP16)) + product = context.graph.multiply(inputs[0], weight) + return context.graph.sum(product, axis=normalize_dim(-1, len(meta_shape(node.parents[0])))) + + if target == "cactus.rel_pos_bias": + require_len(node, inputs, 2) + return context.graph.rel_pos_bias(inputs[0], inputs[1], scale=float(node.attrs.get("scale") or 1.0)) + + if target == "cactus.sample": + require_len(node, inputs, 1) + return context.graph.sample(inputs[0], temperature=float(node.attrs.get("temperature", 0.6)), top_p=float(node.attrs.get("top_p", 0.95)), top_k=int(node.attrs.get("top_k", 20))) + + if target == "cactus.scatter_topk": + require_len(node, inputs, 2) + return context.graph.scatter_topk(inputs[0], inputs[1], num_classes=required_int_attr(node, "num_classes")) + + if target == "cactus.gaussian_topk": + require_len(node, inputs, 1) + return context.graph.gaussian_topk(inputs[0], ppf=float(node.attrs.get("ppf", 0.0))) + + if target == "cactus.altup_predict": + require_len(node, inputs, 2) + return context.graph.altup_predict(inputs[0], list(inputs[1:])) + + if target == "cactus.altup_correct": + require_len(node, inputs, 3) + return context.graph.altup_correct(inputs[0], inputs[1], list(inputs[2:])) + + if target == "cactus.stft": + require_len(node, inputs, 2) + return context.graph.stft(inputs[0], inputs[1], stride=required_int_attr(node, "stride"), num_fft_bins=required_int_attr(node, "num_fft_bins")) + + if target == "cactus.rfft": + require_len(node, inputs, 1) + return context.graph.rfft(inputs[0]) + + if target == "cactus.irfft": + require_len(node, inputs, 1) + return context.graph.irfft(inputs[0], output_length=required_int_attr(node, "output_length")) + + if target == "cactus.spectrogram": + require_len(node, inputs, 2) + return context.graph.spectrogram( + inputs[0], + inputs[1], + frame_length=required_int_attr(node, "frame_length"), + hop_length=required_int_attr(node, "hop_length"), + fft_length=required_int_attr(node, "fft_length"), + power=float(node.attrs.get("power", 2.0)), + ) + + if target == "cactus.image_preprocess": + require_len(node, inputs, 1) + if not has_attrs(node, ("src_width", "src_height", "target_width", "target_height", "patch_size", "channels")): + return context.graph.reshape(inputs[0], output_shape(node)) + + return context.graph.image_preprocess( + inputs[0], + src_width=required_int_attr(node, "src_width"), + src_height=required_int_attr(node, "src_height"), + target_width=required_int_attr(node, "target_width"), + target_height=required_int_attr(node, "target_height"), + patch_size=required_int_attr(node, "patch_size"), + channels=required_int_attr(node, "channels"), + rescale_factor=float(node.attrs.get("rescale_factor", 1.0)), + mean=node.attrs.get("mean", (0.0, 0.0, 0.0)), + std_dev=node.attrs.get("std_dev", (1.0, 1.0, 1.0)), + ) + + if target == "cactus.bilinear_interpolation": + require_len(node, inputs, 1) + return context.graph.bilinear_interpolation(inputs[0], required_int_attr(node, "dst_height"), required_int_attr(node, "dst_width")) + + raise UnsupportedLoweringError(f"{node.name}: unsupported special Cactus target {target}") + + +def lfm_short_conv_decode_weight(context: models.GenerationContext, node: IRModels.Node, weight: Any) -> Any: + cache_window_shape = meta_shape(node.parents[0]) if node.parents else () + weight_shape = meta_shape(node.parents[1]) if len(node.parents) > 1 else () + + if len(cache_window_shape) != 3: + return weight + + batch_size, hidden_dim, window_size = cache_window_shape + + if not all(isinstance(dim, int) for dim in cache_window_shape): + return weight + + if tuple(weight_shape) == tuple(cache_window_shape): + return weight + + if len(weight_shape) == 2 and tuple(weight_shape) == (hidden_dim, window_size): + return context.graph.reshape(weight, (batch_size, hidden_dim, window_size)) + + if len(weight_shape) == 3 and tuple(weight_shape) in { + (1, hidden_dim, window_size), + (hidden_dim, 1, window_size), + }: + return context.graph.reshape(weight, (batch_size, hidden_dim, window_size)) + + raise UnsupportedLoweringError( + f"{node.name}: lfm short conv decode weight shape {weight_shape} is incompatible with cache window {cache_window_shape}" + ) + + +def lower_unsupported_semantic(context: models.GenerationContext, node: IRModels.Node) -> Any: + raise UnsupportedLoweringError(f"{node.name}: {node.target} has no safe Cactus lowering yet") diff --git a/python/cactus/transpiler/Generator/lowering_utils.py b/python/cactus/transpiler/Generator/lowering_utils.py new file mode 100644 index 000000000..478a82b47 --- /dev/null +++ b/python/cactus/transpiler/Generator/lowering_utils.py @@ -0,0 +1,534 @@ +from __future__ import annotations + +import json +import math +import struct +from pathlib import Path +from typing import Any + +from . import models +from .errors import UnsupportedLoweringError +from ..IR import models as IRModels + + +SIZE_T_MAX = (1 << 64) - 1 + + +def write_constant_tensor( + context: models.GenerationContext, + node: IRModels.Node, + values: list[float], + shape: tuple[int, ...], + precision: int, +) -> str: + constants_dir = context.component.output_path.parent / "constants" + constants_dir.mkdir(parents=True, exist_ok=True) + filename = f"{models.sanitize_component_name(context.component.name)}__{models.sanitize_component_name(node.name)}.weights" + path = constants_dir / filename + data = constant_data_bytes(values, precision) + write_cactus_tensor_file(path, shape, precision, data) + return constant_binding_path(context, path) + + +def constant_data_bytes(values: list[float], precision: int) -> bytes: + if precision == 1: + return b"".join(struct.pack(" None: + alignment = 32 + header_size = 84 + ndim = len(shape) + original_n = shape[0] if shape else 0 + + with path.open("wb") as f: + f.write(struct.pack(" str: + try: + return str(path.relative_to(context.config.output_dir.parent)) + except ValueError: + return str(path) + + +def precision_name(graph: Any, precision: int) -> str: + for name in ("INT8", "FP16", "FP32", "CQ1", "CQ2", "CQ3", "CQ4"): + if int(getattr(graph, name)) == int(precision): + return name + + return str(precision) + + +def fp16_tensor(context: models.GenerationContext, value: Any) -> Any: + fp16 = int(context.graph.FP16) + return value if getattr(value, "dtype", None) == fp16 else context.graph.precision_cast(value, fp16) + + +def numeric_attr(node: IRModels.Node, *names: str, default: Any | None = None) -> Any | None: + for name in names: + value = node.attrs.get(name) + + if value is not None: + return value + + return default + + +def require_input_count(context: models.GenerationContext, node: IRModels.Node, min_count: int) -> tuple[Any, ...]: + inputs = context.inputs_for(node) + require_len(node, inputs, min_count) + return inputs + + +def require_at_least_one_input(context: models.GenerationContext, node: IRModels.Node) -> tuple[Any, ...]: + return require_input_count(context, node, 1) + + +def require_len(node: IRModels.Node, inputs: tuple[Any, ...], min_count: int) -> None: + if len(inputs) < min_count: + raise unsupported_arity(node, len(inputs), f"at least {min_count} inputs") + + +def unsupported_arity(node: IRModels.Node, actual: int, expected: str) -> UnsupportedLoweringError: + return UnsupportedLoweringError(f"{node.name}: {node.target} got {actual} lowered inputs; expected {expected}") + + +def tensor_dtype(node: IRModels.Node) -> str | None: + if isinstance(node.tensor_output_meta, dict): + dtype = node.tensor_output_meta.get("dtype") + + if dtype is not None: + return str(dtype) + + if isinstance(node.tensor_output_meta, list) and node.tensor_output_meta: + first_meta = node.tensor_output_meta[0] + + if isinstance(first_meta, dict): + dtype = first_meta.get("dtype") + + if dtype is not None: + return str(dtype) + + return None + + +def cactus_precision(graph: Any, dtype: str | None) -> int: + precision_name = constants.DTYPE_TO_PRECISION.get(str(dtype), constants.DEFAULT_INPUT_PRECISION) + return int(getattr(graph, precision_name)) + + +def cast_to_precision(context: models.GenerationContext, value: Any, precision: int) -> Any: + if getattr(value, "dtype", precision) == precision: + return value + + return context.graph.precision_cast(value, precision) + + +def graph_input_shape(node: IRModels.Node) -> tuple[tuple[int, ...], tuple[bool, ...]]: + shape = meta_shape(node) + + if not shape: + raw_shape = node.tensor_output_meta.get("shape") if isinstance(node.tensor_output_meta, dict) else None + if raw_shape == [] or raw_shape == (): + return (1,), (False,) + + return (), () + + dims: list[int] = [] + dynamic_dims: list[bool] = [] + + for dim in shape: + if isinstance(dim, int) and dim >= 0: + dims.append(dim) + dynamic_dims.append(False) + continue + + if isinstance(dim, str) and dim.isdigit(): + dims.append(int(dim)) + dynamic_dims.append(False) + continue + + dims.append(1) + dynamic_dims.append(True) + + return tuple(dims), tuple(dynamic_dims) + + +def meta_shape(node: IRModels.Node) -> tuple[Any, ...]: + if isinstance(node.tensor_output_meta, dict): + shape = node.tensor_output_meta.get("shape") + + if isinstance(shape, list): + return tuple(shape) + + if isinstance(shape, tuple): + return shape + + if isinstance(node.tensor_output_meta, list) and node.tensor_output_meta: + first_meta = node.tensor_output_meta[0] + + if isinstance(first_meta, dict): + shape = first_meta.get("shape") + + if isinstance(shape, list): + return tuple(shape) + + if isinstance(shape, tuple): + return shape + + return () + + +def output_shape(node: IRModels.Node) -> tuple[int, ...]: + shape, _ = graph_input_shape(node) + + if shape: + return shape + + raise UnsupportedLoweringError(f"{node.name}: missing concrete output shape") + + +def concrete_dim(dim: Any) -> int | None: + if isinstance(dim, int) and dim >= 0: + return dim + + if isinstance(dim, str) and dim.isdigit(): + return int(dim) + + return None + + +def concrete_shape(shape: tuple[Any, ...]) -> tuple[int, ...] | None: + dims = tuple(concrete_dim(dim) for dim in shape) + + if any(dim is None for dim in dims): + return None + + return tuple(int(dim) for dim in dims) + + +def reduction_dropped_shape(shape: tuple[Any, ...], axis: int) -> tuple[Any, ...]: + dropped = tuple(dim for index, dim in enumerate(shape) if index != axis) + return dropped or (1,) + + +def shape_matches_tensor(node: IRModels.Node, expected_shape: tuple[int, ...]) -> bool: + actual_shape = meta_shape(node) + return tuple(actual_shape) == tuple(expected_shape) + + +def element_count(shape: tuple[Any, ...]) -> int | None: + if not shape: + return None + + count = 1 + + for dim in shape: + if not isinstance(dim, int) or dim < 0: + return None + + count *= dim + + return count + + +def shape_attr(node: IRModels.Node) -> tuple[int, ...]: + raw_shape = node.attrs.get("shape") + + if raw_shape is None: + return output_shape(node) + + resolved_output_shape = output_shape(node) + dims: list[int] = [] + + for index, dim in enumerate(raw_shape): + if isinstance(dim, int) and dim >= 0: + dims.append(dim) + elif isinstance(dim, str) and dim.isdigit(): + dims.append(int(dim)) + elif index < len(resolved_output_shape): + dims.append(resolved_output_shape[index]) + else: + raise UnsupportedLoweringError(f"{node.name}: cannot resolve shape dim {dim!r}") + + return tuple(dims) + + +def axis_attr(node: IRModels.Node, default: int | None = None) -> int | None: + axis = node.attrs.get("axis", node.attrs.get("dim", node.attrs.get("arg_1", default))) + + if isinstance(axis, list): + if len(axis) != 1: + raise UnsupportedLoweringError(f"{node.name}: Cactus reduction lowering only supports one axis") + return int(axis[0]) + + if axis is None: + return None + + return int(axis) + + +def scalar_attr(node: IRModels.Node, name: str) -> Any | None: + value = node.attrs.get(name) + + if isinstance(value, dict) and "node" in value: + return None + + if isinstance(value, list) and len(value) == 1: + return value[0] + + return value + + +def scalar_weight_bound_value(context: models.GenerationContext, node: IRModels.Node) -> float | None: + resolver = context.component.weight_resolver + + if resolver is None: + return None + + record = resolver.resolve(node.name) + + if record is None or record.output_name is None: + return None + + if record.shape and math.prod(int(dim) for dim in record.shape) != 1: + return None + + value = read_scalar_cactus_weight(resolver.weights_dir / record.output_name) + + if value is None: + return None + + return float(value) * float(record.scale_factor) + + +def read_scalar_cactus_weight(path: Path) -> float | None: + if not path.exists(): + return None + + with path.open("rb") as f: + header = f.read(84) + + if len(header) != 84 or header[:4] != b"CACT": + return None + + fields = struct.unpack("= 2: + return float(struct.unpack("= 4: + return float(struct.unpack(" int: + if alignment <= 0: + alignment = 32 + + remainder = offset % alignment + return offset if remainder == 0 else offset + alignment - remainder + + +def attr_value(node: IRModels.Node, name: str, default: Any) -> Any: + value = node.attrs.get(name) + return default if value is None else value + + +def epsilon_attr(node: IRModels.Node) -> float: + return float(node.attrs.get("epsilon", node.attrs.get("eps", 1e-5))) + + +def required_int_attr(node: IRModels.Node, name: str) -> int: + value = node.attrs.get(name) + + if value is None: + raise UnsupportedLoweringError(f"{node.name}: missing required attr {name}") + + return int(value) + + +def has_attrs(node: IRModels.Node, names: tuple[str, ...]) -> bool: + return all(node.attrs.get(name) is not None for name in names) + + +def first_int(value: Any, default: int) -> int: + if value is None: + return default + + if isinstance(value, (list, tuple)): + if not value: + return default + return int(value[0]) + + return int(value) + + +def tuple_int_values(value: Any, default: int) -> tuple[int, int]: + if value is None: + return (default, default) + + if isinstance(value, (list, tuple)): + if not value: + return (default, default) + + if len(value) == 1: + item = int(value[0]) + return (item, item) + + return (int(value[0]), int(value[1])) + + item = int(value) + return (item, item) + + +def tuple_ints(value: Any) -> tuple[int, ...]: + if not isinstance(value, (list, tuple)): + raise TypeError(f"Expected list/tuple of ints, got {type(value).__name__}") + + return tuple(int(item) for item in value) + + +def parent_rank(node: IRModels.Node) -> int: + if not node.parents: + raise UnsupportedLoweringError(f"{node.name}: cannot infer parent rank without parents") + + rank = len(meta_shape(node.parents[0])) + + if rank == 0: + raise UnsupportedLoweringError(f"{node.name}: cannot infer parent rank without parent shape metadata") + + return rank + + +def output_rank(node: IRModels.Node) -> int: + return len(meta_shape(node)) + + +def swap_permutation(rank: int, dim0: int, dim1: int) -> tuple[int, ...]: + dim0 = normalize_dim(dim0, rank) + dim1 = normalize_dim(dim1, rank) + permutation = list(range(rank)) + permutation[dim0], permutation[dim1] = permutation[dim1], permutation[dim0] + return tuple(permutation) + + +def normalize_dim(dim: int, rank: int) -> int: + if dim < 0: + dim += rank + + if dim < 0 or dim >= rank: + raise ValueError(f"Dimension {dim} is outside rank {rank}") + + return dim + + +def slice_length(node: IRModels.Node, start: int) -> int: + if "length" in node.attrs: + return int(node.attrs["length"]) + + end = node.attrs.get("end") + + if end is None: + return open_slice_length(node, start) + + end = int(end) + + if end >= constants.OPEN_SLICE_END: + return open_slice_length(node, start) + + return max(end - start, 0) + + +def open_slice_length(node: IRModels.Node, start: int) -> int: + if not node.parents: + return 0 + + axis = int(node.attrs.get("axis", node.attrs.get("dim", 0))) + parent_shape = meta_shape(node.parents[0]) + + if axis < 0: + axis += len(parent_shape) + + if axis < 0 or axis >= len(parent_shape): + return 0 + + axis_size = parent_shape[axis] + + if not isinstance(axis_size, int): + return 0 + + return max(axis_size - start, 0) + + +def binary_method_to_scalar_method(method: str) -> str: + return { + "add": "scalar_add", + "subtract": "scalar_subtract", + "multiply": "scalar_multiply", + "divide": "scalar_divide", + "not_equal": "scalar_not_equal", + }[method] + + +def ensure_tensor_sequence(value: Any) -> list[Any]: + if isinstance(value, list): + return value + + if isinstance(value, tuple): + return list(value) + + return [value] + + +def dump_result_manifest(result: models.GenerationResult, output_path: str | Path) -> Path: + path = Path(output_path) + path.parent.mkdir(parents=True, exist_ok=True) + path.write_text( + json.dumps( + { + "component_paths": {name: str(path_) for name, path_ in result.component_paths.items()}, + "component_manifest_paths": {name: str(path_) for name, path_ in result.component_manifest_paths.items()}, + "engine_manifest_path": str(result.engine_manifest_path) if result.engine_manifest_path is not None else None, + "runtime_plan_path": str(result.runtime_plan_path) if result.runtime_plan_path is not None else None, + "unsupported_nodes": list(result.unsupported_nodes), + "warnings": list(result.warnings), + "ok": result.ok, + }, + indent=4, + ), + encoding="utf-8", + ) + return path diff --git a/python/cactus/transpiler/Generator/lowering_weights.py b/python/cactus/transpiler/Generator/lowering_weights.py new file mode 100644 index 000000000..7668a7d8c --- /dev/null +++ b/python/cactus/transpiler/Generator/lowering_weights.py @@ -0,0 +1,293 @@ +from __future__ import annotations + +import json +import re +from typing import Any + +from . import constants +from . import models +from .errors import UnsupportedLoweringError +from .lowering_utils import * +from ..IR import models as IRModels + + +LFM_GROUPED_MOE_TARGET_RE = re.compile(r"^layers\.(\d+)\.feed_forward\.experts\.(gate_up_proj|down_proj)$") + + +def bind_weight_placeholder(context: models.GenerationContext, node: IRModels.Node, tensor: Any, logical_name: str | None = None) -> None: + resolver = context.component.weight_resolver + + if resolver is None: + context.component.warn(f"{node.name}: weight lowered as graph input because no weights_dir was provided") + context.component.add_runtime_input(tensor, logical_name) + return + + source_target = resolver.source_target_for(node.name) or node.target + record = resolver.resolve(node.name) + node_id = models.tensor_node_id(tensor) + + if node_id is None: + message = f"{node.name}: lowered weight tensor does not expose a Cactus node id" + + if context.config.strict and not context.config.allow_unsupported_ops: + raise UnsupportedLoweringError(message) + + context.component.warn(message) + return + + if record is None or record.output_name is None: + generated = generated_weight_placeholder_values(context, node, source_target) + + if generated is not None: + values, shape, dtype = generated + path = write_constant_tensor(context, node, values, shape, dtype) + context.component.add_weight_binding( + models.WeightBinding( + placeholder=node.name, + source_target=source_target, + node_id=node_id, + path=path, + output_name=path, + source_name=source_target, + value_id=node.name, + precision=precision_name(context.graph, dtype), + component=context.component.name, + binding_kind="generated_constant", + ) + ) + return + + message = f"{node.name}: could not resolve converted weight for source target {source_target}" + + if context.config.strict and not context.config.allow_unsupported_ops: + raise UnsupportedLoweringError(message) + + context.component.warn(message) + context.component.add_runtime_input(tensor, logical_name) + return + + context.component.add_weight_binding( + models.WeightBinding( + placeholder=node.name, + source_target=source_target, + node_id=node_id, + path=record.output_name, + output_name=record.output_name, + source_name=record.source_name or record.hf_name or source_target, + value_id=node.name, + precision=record.precision, + component=record.component, + scale_factor=record.scale_factor, + adapter_family=record.adapter_family, + transform=record.transform, + qdq_restore=record.qdq_restore, + ) + ) + + +def generated_weight_placeholder_values( + context: models.GenerationContext, + node: IRModels.Node, + source_target: str, +) -> tuple[list[float], tuple[int, ...], int] | None: + if not source_target.endswith(".pos_emb.inv_freq"): + return None + + shape = concrete_shape(meta_shape(node)) + + if shape is None or len(shape) != 1: + return None + + rope_theta = model_rope_theta(context) + + if rope_theta is None: + return None + + rotary_dim = int(shape[0]) * 2 + values = [1.0 / (float(rope_theta) ** (index / rotary_dim)) for index in range(0, rotary_dim, 2)] + return values, tuple(shape), cactus_precision(context.graph, tensor_dtype(node)) + + +def model_rope_theta(context: models.GenerationContext) -> float | None: + weights_dir = context.config.weights_dir + + if weights_dir is None: + return None + + config_path = weights_dir / "config.json" + + if not config_path.is_file(): + config_path = weights_dir / "hf_config.json" + + if not config_path.is_file(): + return None + + try: + config = json.loads(config_path.read_text(encoding="utf-8")) + except json.JSONDecodeError: + return None + + rope_parameters = config.get("rope_parameters") + + if isinstance(rope_parameters, Mapping): + value = rope_parameters.get("rope_theta") or rope_parameters.get("theta") + + if value is not None: + return float(value) + + value = config.get("rope_theta") or config.get("rotary_emb_base") + return float(value) if value is not None else None + + +def lower_lfm_grouped_moe_placeholder(context: models.GenerationContext, node: IRModels.Node) -> models.LfmMoeWeightBundle | None: + resolver = context.component.weight_resolver + + if resolver is None: + return None + + source_target = resolver.source_target_for(node.name) or node.target + parsed = parse_lfm_grouped_moe_target(source_target) + + if parsed is None: + return None + + layer_index, grouped_kind = parsed + placeholder_shape = concrete_shape(meta_shape(node)) + num_experts = placeholder_shape[0] if placeholder_shape is not None and placeholder_shape else None + + if not isinstance(num_experts, int): + num_experts = count_lfm_moe_experts(resolver, layer_index) + + if num_experts <= 0: + raise UnsupportedLoweringError(f"{node.name}: could not infer LFM MoE expert count for {source_target}") + + if grouped_kind == "gate_up_proj": + return models.LfmMoeWeightBundle( + w1_weights=tuple( + lfm_expert_weight_input(context, node, layer_index, expert_index, "w1") + for expert_index in range(num_experts) + ), + w3_weights=tuple( + lfm_expert_weight_input(context, node, layer_index, expert_index, "w3") + for expert_index in range(num_experts) + ), + ) + + return models.LfmMoeWeightBundle( + w2_weights=tuple( + lfm_expert_weight_input(context, node, layer_index, expert_index, "w2") + for expert_index in range(num_experts) + ) + ) + + +def parse_lfm_grouped_moe_target(source_target: str) -> tuple[int, str] | None: + target = source_target + + while target.startswith("model."): + target = target[len("model."):] + + match = LFM_GROUPED_MOE_TARGET_RE.fullmatch(target) + + if match is None: + return None + + return int(match.group(1)), match.group(2) + + +def count_lfm_moe_experts(resolver: models.WeightResolver, layer_index: int) -> int: + prefix = f"model.layers.{layer_index}.feed_forward.experts." + expert_indices: set[int] = set() + + for record in resolver.records: + for alias in record.aliases: + if not alias.startswith(prefix) or not alias.endswith(".w1.weight"): + continue + + suffix = alias[len(prefix):] + expert = suffix.split(".", 1)[0] + + if expert.isdigit(): + expert_indices.add(int(expert)) + + return max(expert_indices) + 1 if expert_indices else 0 + + +def lfm_expert_weight_input( + context: models.GenerationContext, + grouped_node: IRModels.Node, + layer_index: int, + expert_index: int, + role: str, +) -> Any: + source_target = f"model.layers.{layer_index}.feed_forward.experts.{expert_index}.{role}.weight" + record = resolve_required_source_weight(context, grouped_node, source_target) + tensor = context.graph.input(tuple(record.shape), dtype=weight_record_precision(context, record)) + bind_weight_record( + context, + grouped_node, + tensor, + source_target, + record, + value_id=f"{grouped_node.name}:{role}:{expert_index}", + ) + return tensor + + +def resolve_required_source_weight( + context: models.GenerationContext, + node: IRModels.Node, + source_target: str, +) -> models.WeightRecord: + resolver = context.component.weight_resolver + + if resolver is None: + raise UnsupportedLoweringError(f"{node.name}: no weight resolver available for {source_target}") + + for candidate in (source_target, source_target.removeprefix("model.")): + record = resolver.resolve_source_target(candidate) + + if record is not None and record.output_name is not None: + return record + + raise UnsupportedLoweringError(f"{node.name}: could not resolve converted weight for source target {source_target}") + + +def weight_record_precision(context: models.GenerationContext, record: models.WeightRecord) -> int: + if record.precision and hasattr(context.graph, record.precision): + return int(getattr(context.graph, record.precision)) + + return int(getattr(context.graph, constants.DEFAULT_INPUT_PRECISION)) + + +def bind_weight_record( + context: models.GenerationContext, + node: IRModels.Node, + tensor: Any, + source_target: str, + record: models.WeightRecord, + *, + value_id: str, +) -> None: + node_id = models.tensor_node_id(tensor) + + if node_id is None: + raise UnsupportedLoweringError(f"{node.name}: lowered weight tensor does not expose a Cactus node id") + + context.component.add_weight_binding( + models.WeightBinding( + placeholder=node.name, + source_target=source_target, + node_id=node_id, + path=str(record.output_name), + output_name=str(record.output_name), + source_name=record.source_name or record.hf_name or source_target, + value_id=value_id, + precision=record.precision, + component=record.component, + scale_factor=record.scale_factor, + adapter_family=record.adapter_family, + transform=record.transform, + qdq_restore=record.qdq_restore, + ) + ) diff --git a/python/cactus/transpiler/Generator/lowerings.py b/python/cactus/transpiler/Generator/lowerings.py new file mode 100644 index 000000000..68ed2c0e5 --- /dev/null +++ b/python/cactus/transpiler/Generator/lowerings.py @@ -0,0 +1,10 @@ +from __future__ import annotations + +from .lowering_basic_ops import * +from .lowering_cache import * +from .lowering_constant_ops import * +from .lowering_core import * +from .lowering_nn_ops import * +from .lowering_special_ops import * +from .lowering_utils import * +from .lowering_weights import * diff --git a/python/cactus/transpiler/Generator/models.py b/python/cactus/transpiler/Generator/models.py index de7134251..cbeba4c24 100644 --- a/python/cactus/transpiler/Generator/models.py +++ b/python/cactus/transpiler/Generator/models.py @@ -171,6 +171,13 @@ class WeightBinding: binding_kind: str = "mmap_weight" +@dataclass(slots=True, frozen=True) +class LfmMoeWeightBundle: + w1_weights: tuple[GraphTensor, ...] = () + w3_weights: tuple[GraphTensor, ...] = () + w2_weights: tuple[GraphTensor, ...] = () + + @dataclass(slots=True) class WeightResolver: weights_dir: Path @@ -185,6 +192,9 @@ def from_graph(cls, graph: IRModels.Graph, weights_dir: Path, manifest_path: Pat def resolve(self, placeholder_name: str) -> WeightRecord | None: return resolve_weight_record(self, placeholder_name) + def resolve_source_target(self, source_target: str) -> WeightRecord | None: + return resolve_source_weight_record(self, source_target) + def source_target_for(self, placeholder_name: str) -> str | None: return self.placeholder_targets.get(placeholder_name) @@ -264,7 +274,7 @@ def component_graph_from_ir(cls: type[ComponentGraph], name: str, ir_graph: IRMo def component_metadata_from_ir(ir_graph: IRModels.Graph) -> dict[str, str]: - metadata: dict[str, str] = {} + metadata: dict[str, str] = dict(ir_graph.metadata) for node in ir_graph.nodes: chunk_tokens = node.ir_metadata.get("prefill_chunk_tokens") @@ -521,6 +531,10 @@ def resolve_weight_record(resolver: WeightResolver, placeholder_name: str) -> We if source_target is None: return None + return resolve_source_weight_record(resolver, source_target) + + +def resolve_source_weight_record(resolver: WeightResolver, source_target: str) -> WeightRecord | None: for variant in weight_name_variants(source_target): record = resolver.records_by_name.get(variant) @@ -548,6 +562,16 @@ def weight_name_variants(name: str) -> tuple[str, ...]: ) ) + if name.endswith("proj_out.weight"): + variants.extend( + ( + "model.decoder.embed_tokens.weight", + "decoder.embed_tokens.weight", + "model.embed_tokens.weight", + "embed_tokens.weight", + ) + ) + if name.startswith("_orig_mod."): variants.append(name[len("_orig_mod."):]) diff --git a/python/cactus/transpiler/IR/extra_matcher_cache.py b/python/cactus/transpiler/IR/extra_matcher_cache.py new file mode 100644 index 000000000..e31d77b21 --- /dev/null +++ b/python/cactus/transpiler/IR/extra_matcher_cache.py @@ -0,0 +1,235 @@ +from __future__ import annotations + +from typing import Any + +from . import models, match_utils +from .extra_matcher_common import * +from ..Fusions import models as FModels + + +def match_cache_roll_append_structure(source: models.Node, graph: models.Graph, fusion: FModels.FusionGraph, bindings: dict[str, models.Node], spec: dict[str, Any]) -> bool: + cache_node = match_utils.get_first_input_by_role(fusion, bindings, spec["cache_role"]) + new_data_node = match_utils.get_first_input_by_role(fusion, bindings, spec["new_data_role"]) + + if cache_node is None or new_data_node is None: + return False + + cache_shape = match_utils.get_tensor_shape(cache_node) + new_data_shape = match_utils.get_tensor_shape(new_data_node) + output_shape = match_utils.get_tensor_shape(source) + + if not cache_shape or not new_data_shape or not output_shape: + return False + + if not match_utils.values_equal(cache_shape, output_shape): + return False + + if len(cache_shape) != len(new_data_shape): + return False + + window_size = get_known_int(cache_shape[-1]) + + if window_size is None or window_size <= 0: + return False + + if not cache_append_shapes_match(cache_shape, new_data_shape): + return False + + arange_node = bindings.get(spec["arange_node"]) + add_node = bindings.get(spec["add_node"]) + mod_node = bindings.get(spec["mod_node"]) + + if arange_node is None or add_node is None or mod_node is None: + return False + + if not match_utils.values_equal(get_attr(arange_node, "start", "arg_0"), 0): + return False + + if not match_utils.values_equal(get_attr(arange_node, "end", "arg_1"), window_size): + return False + + if not match_utils.values_equal(get_attr(add_node, "other", "value", "arg_1"), 1): + return False + + if not match_utils.values_equal(get_attr(mod_node, "other", "value", "arg_1"), window_size): + return False + + return ( + slice_node_matches(bindings.get(spec["scatter_slice_dim0_node"]), axis=0, start=0) + and slice_node_matches(bindings.get(spec["scatter_slice_dim1_node"]), axis=1, start=0) + and slice_node_matches(bindings.get(spec["value_slice_dim0_node"]), axis=0, start=0) + and slice_node_matches(bindings.get(spec["value_slice_dim1_node"]), axis=1, start=0) + and slice_node_matches(bindings.get(spec["value_slice_last_node"]), axis=2, start=-1) + and scatter_node_matches(bindings.get(spec["scatter_dim2_node"]), axis=2, start=-1) + and scatter_node_matches(bindings.get(spec["scatter_dim1_node"]), axis=1, start=0) + and scatter_node_matches(bindings.get(spec["scatter_dim0_node"]), axis=0, start=0) + ) + + +def match_empty_cache_initializer(source: models.Node, graph: models.Graph, fusion: FModels.FusionGraph, bindings: dict[str, models.Node], spec: dict[str, Any]) -> bool: + empty_node = match_utils.get_first_input_by_role(fusion, bindings, spec["empty_role"]) + new_value_node = match_utils.get_first_input_by_role(fusion, bindings, spec["new_value_role"]) + output_node = bindings.get(spec["output_node"]) + + if empty_node is None or new_value_node is None or output_node is None: + return False + + empty_shape = match_utils.get_tensor_shape(empty_node) + new_shape = match_utils.get_tensor_shape(new_value_node) + output_shape = match_utils.get_tensor_shape(output_node) + + if not empty_shape or not new_shape or not output_shape: + return False + + if element_count(empty_shape) != 0: + return False + + return match_utils.values_equal(new_shape, output_shape) + + +def match_conv_cache_initialize_structure(source: models.Node, graph: models.Graph, fusion: FModels.FusionGraph, bindings: dict[str, models.Node], spec: dict[str, Any]) -> bool: + rows_node = match_utils.get_first_input_by_role(fusion, bindings, spec["rows_role"]) + empty_node = match_utils.get_first_input_by_role(fusion, bindings, spec["empty_role"]) + output_node = bindings.get(spec["output_node"]) + + if rows_node is None or empty_node is None or output_node is None: + return False + + rows_shape = match_utils.get_tensor_shape(rows_node) + empty_shape = match_utils.get_tensor_shape(empty_node) + output_shape = match_utils.get_tensor_shape(output_node) + + if len(rows_shape) != 3 or len(empty_shape) != 3 or len(output_shape) != 3: + return False + + if not match_utils.values_equal(rows_shape, output_shape): + return False + + if not match_utils.values_equal(empty_shape, output_shape): + return False + + if empty_node.target != "aten.full_like.default": + return False + + if not match_utils.values_equal(get_attr(empty_node, "fill_value", "arg_1"), 0): + return False + + return get_known_int(rows_shape[1]) is not None and get_known_int(rows_shape[2]) is not None + + +def match_short_conv_prefill_structure(source: models.Node, graph: models.Graph, fusion: FModels.FusionGraph, bindings: dict[str, models.Node], spec: dict[str, Any]) -> bool: + conv_node = bindings.get(spec["conv_node"]) + slice_node = bindings.get(spec["slice_node"]) + x_node = match_utils.get_first_input_by_role(fusion, bindings, spec["x_role"]) + weight_node = match_utils.get_first_input_by_role(fusion, bindings, spec["weight_role"]) + + if conv_node is None or slice_node is None or x_node is None or weight_node is None: + return False + + x_shape = match_utils.get_tensor_shape(x_node) + weight_shape = match_utils.get_tensor_shape(weight_node) + conv_shape = match_utils.get_tensor_shape(conv_node) + output_shape = match_utils.get_tensor_shape(slice_node) + + if len(x_shape) != 3 or len(weight_shape) != 3 or len(conv_shape) != 3 or len(output_shape) != 3: + return False + + hidden_dim = get_known_int(x_shape[1]) + kernel_size = get_known_int(weight_shape[2]) + input_length = get_known_int(x_shape[2]) + + if hidden_dim is None or kernel_size is None or input_length is None: + return False + + if not match_utils.values_equal(weight_shape[0], hidden_dim): + return False + + if not match_utils.values_equal(weight_shape[1], 1): + return False + + if first_int_attr(conv_node, "stride", 1) != 1: + return False + + if first_int_attr(conv_node, "dilation", 1) != 1: + return False + + if first_int_attr(conv_node, "padding", 0) != kernel_size - 1: + return False + + if first_int_attr(conv_node, "groups", 1) != hidden_dim: + return False + + slice_axis = first_int_attr(slice_node, "axis", first_int_attr(slice_node, "dim", -1)) + if normalize_dim(slice_axis, 3) != 2: + return False + + if first_int_attr(slice_node, "start", 0) != 0: + return False + + return match_utils.values_equal(output_shape, x_shape) and get_shape_dim(conv_node, 2) == input_length + kernel_size - 1 + + +def match_short_conv_decode_structure(source: models.Node, graph: models.Graph, fusion: FModels.FusionGraph, bindings: dict[str, models.Node], spec: dict[str, Any]) -> bool: + cache_window_node = match_utils.get_first_input_by_role(fusion, bindings, spec["cache_window_role"]) + multiply_node = bindings.get(spec["multiply_node"]) + sum_node = bindings.get(spec["sum_node"]) + weight_node = match_utils.get_first_input_by_role(fusion, bindings, spec["weight_role"]) + + if cache_window_node is None or multiply_node is None or sum_node is None or weight_node is None: + return False + + if cache_window_node.target != "cactus.conv_cache_append" and not ( + cache_window_node.cache is not None and cache_window_node.cache.kind == FModels.CacheKind.CONV + ): + return False + + cache_window_shape = match_utils.get_tensor_shape(cache_window_node) + weight_shape = match_utils.get_tensor_shape(weight_node) + multiply_shape = match_utils.get_tensor_shape(multiply_node) + output_shape = match_utils.get_tensor_shape(sum_node) + + if len(cache_window_shape) != 3 or len(multiply_shape) != 3 or len(output_shape) != 2: + return False + + if not match_utils.values_equal(cache_window_shape, multiply_shape): + return False + + batch_size = get_known_int(cache_window_shape[0]) + hidden_dim = get_known_int(cache_window_shape[1]) + window_size = get_known_int(cache_window_shape[2]) + + if batch_size is None or hidden_dim is None or window_size is None: + return False + + if not match_utils.values_equal(output_shape, [batch_size, hidden_dim]): + return False + + if not short_conv_decode_weight_shape_matches(weight_shape, hidden_dim, window_size): + return False + + sum_axis = first_int_attr(sum_node, "axis", first_int_attr(sum_node, "dim", -1)) + if normalize_dim(sum_axis, 3) != 2: + return False + + keepdim = get_attr(sum_node, "keepdim") + return keepdim is match_utils.MISSING or keepdim is False + + +def match_cache_output_consumers(source: models.Node, graph: models.Graph, fusion: FModels.FusionGraph, bindings: dict[str, models.Node], spec: dict[str, Any]) -> bool: + node = bindings.get(spec["node"]) + + if node is None: + return False + + allow_output_child = bool(spec.get("allow_output_child", True)) + + for child in node.children: + if child.target in match_utils.METADATA_ONLY_TARGETS: + continue + + if allow_output_child and child.is_output: + continue + + return False + + return True diff --git a/python/cactus/transpiler/IR/extra_matcher_common.py b/python/cactus/transpiler/IR/extra_matcher_common.py new file mode 100644 index 000000000..0615129b0 --- /dev/null +++ b/python/cactus/transpiler/IR/extra_matcher_common.py @@ -0,0 +1,346 @@ +from __future__ import annotations + +from typing import Any, Callable + +from . import models, match_utils +from ..Fusions import models as FModels + + +ExtraMatcher = Callable[[models.Node, models.Graph, FModels.FusionGraph, dict[str, models.Node], dict[str, Any]], bool] + +WEIGHT_VALUE_KINDS = ( + FModels.ValueKind.PARAMETER, + FModels.ValueKind.BUFFER, + FModels.ValueKind.LIFTED_CONSTANT, +) + +COPY_TARGETS = { + "aten.clone.default", + "aten._to_copy.default", + "aten.contiguous.default", +} + +TRANSPOSE_TARGETS = { + "aten.t.default", + "aten.transpose.int", + "aten.permute.default", +} + +COS_TARGETS = {"aten.cos.default"} +SIN_TARGETS = {"aten.sin.default"} +SPLIT_TARGETS = { + "aten.chunk.default", + "aten.split.Tensor", + "aten.split_with_sizes.default", + "aten.unbind.int", +} + + +def get_bound_parent(bindings: dict[str, models.Node], node_name: str, parent_index: int) -> models.Node | None: + """ + Returns a real parent node from a bound synthetic node/input position. + + `bindings` maps synthetic fusion-node names to real IR nodes. This helper + first finds the real node bound to `node_name`, then asks match_utils for + the parent at `parent_index`. It returns None when the synthetic node was + not bound or the requested parent index does not exist. + """ + node = bindings.get(node_name) + + if node is None: + return None + + return match_utils.get_parent(node, parent_index) + + +def get_node_dim(node: models.Node, dim_index: int) -> Any: + """ + Returns one tensor dimension from a real IR node's output shape. + + This reads `tensor_output_meta["shape"]`, supports negative indices in the + Python style, and returns `match_utils.MISSING` instead of raising when the + shape is unavailable or the dimension index is invalid. Matchers use that + missing sentinel so they can either fail closed or honor `allow_missing`. + """ + shape = match_utils.get_tensor_shape(node) + + if not shape: + return match_utils.MISSING + + if dim_index < 0: + dim_index += len(shape) + + if dim_index < 0 or dim_index >= len(shape): + return match_utils.MISSING + + return shape[dim_index] + + +def walk_children(node: models.Node, max_depth: int) -> tuple[models.Node, ...]: + """ + Returns all reachable child nodes within a bounded forward traversal. + + Extra matchers use this when they need local forward evidence, such as + proving an LSTM gate tensor is split later. The traversal tracks object ids + to avoid revisiting shared children in a DAG and to remain safe if graph + rewrites produce unexpected cycles. + """ + found: list[models.Node] = [] + seen: set[int] = set() + + def visit(current: models.Node, depth: int) -> None: + """ + Recursively collects children while respecting the maximum depth. + + The outer `walk_children` owns the result list and visited set; this + nested helper only handles traversal mechanics. + """ + if depth > max_depth: + return + + for child in current.children: + if id(child) in seen: + continue + + seen.add(id(child)) + found.append(child) + visit(child, depth + 1) + + visit(node, 0) + return tuple(found) + + +def get_node_dtype(node: models.Node) -> Any: + """ + Reads a node's output dtype from tensor metadata. + + The exported JSON stores tensor metadata as dictionaries. If metadata is + missing or has some non-dict form, the function returns the shared MISSING + sentinel so callers can decide whether to fail or allow missing metadata. + """ + if not isinstance(node.tensor_output_meta, dict): + return match_utils.MISSING + + return node.tensor_output_meta.get("dtype", match_utils.MISSING) + + +def normalize_dim(dim: Any, rank: int) -> Any: + """ + Converts a possibly negative dimension index into a non-negative one. + + This mirrors Python indexing for known integer dimensions. Non-integer + values are returned unchanged so symbolic or malformed values do not crash + the matcher; they simply fail later if an exact comparison is required. + """ + if not isinstance(dim, int): + return dim + + return dim + rank if dim < 0 else dim + + +def get_attr(node: models.Node, *names: str) -> Any: + """ + Returns the first available attr from a node. + + Some attrs are normalized with friendly names such as `split_sizes`, while + older/generated JSON may only expose fallback names such as `arg_1`. This + lets extra matchers accept either form without duplicating logic. + """ + for name in names: + value = node.attrs.get(name, match_utils.MISSING) + + if value is not match_utils.MISSING: + return value + + return match_utils.MISSING + + +def get_known_int(value: Any) -> int | None: + """ + Converts plain integer-like values into int, otherwise returns None. + + Shape metadata can contain symbolic values, so matchers use this helper + when a semantic check is only safe for known concrete dimensions. + """ + return value if isinstance(value, int) and not isinstance(value, bool) else None + + +def get_shape_dim(node: models.Node | None, dim_index: int) -> int | None: + """ + Returns a concrete int dimension from a node shape when available. + + The helper supports negative indexing and intentionally returns None for + missing/symbolic dimensions so callers can decide whether to fail or skip. + """ + if node is None: + return None + + return get_known_int(get_node_dim(node, dim_index)) + + +def shapes_are_broadcastable(left_shape: list[Any], right_shape: list[Any]) -> bool: + """ + Implements standard trailing-dimension broadcasting for two shapes. + + It compares dimensions from right to left. Each pair is compatible if the + dims are equal or either side is 1. Extra leading dimensions on the longer + shape are allowed, matching normal tensor broadcasting rules. + """ + for left_dim, right_dim in zip(reversed(left_shape), reversed(right_shape)): + if match_utils.values_equal(left_dim, 1) or match_utils.values_equal(right_dim, 1): + continue + + if not match_utils.values_equal(left_dim, right_dim): + return False + + return True + + +def get_layer_key(node: models.Node) -> str | None: + """ + Extracts a coarse layer identifier from module stack, node name, or target. + + Cached attention checks need to know whether different matched nodes belong + to the same decoder layer. This tries the richer module stack first, then + falls back to name/target strings, and returns keys like `layers.0` when it + can identify a layer index. + """ + for value in (module_stack_text(node), node.name, node.target): + layer_key = first_layer_key(value) + + if layer_key is not None: + return layer_key + + return None + + +def module_stack_text(node: models.Node) -> str: + """ + Converts module_stack metadata into searchable text. + + Exported module stacks may be None or structured objects/lists. The matcher + only needs a best-effort string to search for layer-like path fragments, so + None becomes an empty string and everything else uses repr. + """ + return "" if node.module_stack is None else repr(node.module_stack) + + +def first_layer_key(value: str) -> str | None: + """ + Finds the first layer/block path fragment in a string. + + The helper normalizes `/` and `_` into `.` so names like + `model_layers_12_attn` can be treated similarly to + `model.layers.12.attn`. It then looks for common containers such as + `layers`, `layer`, `blocks`, or `h` followed by a numeric index. + """ + parts = value.replace("/", ".").replace("_", ".").split(".") + + for index, part in enumerate(parts[:-1]): + if part in {"layers", "layer", "h", "blocks", "block"} and parts[index + 1].isdigit(): + return f"{part}.{parts[index + 1]}" + + return None + + +def element_count(shape: list[Any]) -> int | None: + count = 1 + + for dim in shape: + known_dim = get_known_int(dim) + + if known_dim is None: + return None + + count *= known_dim + + return count + + +def first_int_attr(node: models.Node, name: str, default: int) -> int: + value = get_attr(node, name) + + if value is match_utils.MISSING: + return default + + if isinstance(value, (list, tuple)): + if not value: + return default + + return int(value[0]) + + return int(value) + + +def cache_append_shapes_match(cache_shape: list[Any], new_data_shape: list[Any]) -> bool: + for cache_dim, new_dim in zip(cache_shape[:-1], new_data_shape[:-1]): + if not match_utils.values_equal(cache_dim, new_dim): + return False + + return match_utils.values_equal(new_data_shape[-1], 1) + + +def short_conv_decode_weight_shape_matches(weight_shape: list[Any], hidden_dim: int, window_size: int) -> bool: + if len(weight_shape) == 2: + return match_utils.values_equal(weight_shape[0], hidden_dim) and match_utils.values_equal(weight_shape[1], window_size) + + if len(weight_shape) != 3: + return False + + return ( + match_utils.values_equal(weight_shape, [1, hidden_dim, window_size]) + or match_utils.values_equal(weight_shape, [hidden_dim, 1, window_size]) + ) + + +def slice_node_matches(node: models.Node | None, *, axis: int, start: int) -> bool: + if node is None: + return False + + return match_utils.values_equal(get_attr(node, "axis", "dim", "arg_1"), axis) and match_utils.values_equal(get_attr(node, "start", "arg_2"), start) + + +def scatter_node_matches(node: models.Node | None, *, axis: int, start: int) -> bool: + if node is None: + return False + + return match_utils.values_equal(get_attr(node, "dim", "axis", "arg_2"), axis) and match_utils.values_equal(get_attr(node, "start", "arg_3"), start) + + +def normalized_slice_axis(node: models.Node, rank: int) -> Any: + axis = node.attrs.get("axis", node.attrs.get("dim", match_utils.MISSING)) + + if axis is match_utils.MISSING: + return match_utils.MISSING + + return normalize_dim(axis, rank) + + +def slice_step_is_one(node: models.Node) -> bool: + step = node.attrs.get("step", 1) + return step is None or match_utils.values_equal(step, 1) + + +def attr_matches(node: models.Node, attr: str, expected: Any, *, default: Any = match_utils.MISSING) -> bool: + actual = node.attrs.get(attr, default) + + if actual is match_utils.MISSING: + return False + + return match_utils.values_equal(actual, expected) + + +def optional_attr_matches(node: models.Node, attr: str, expected: Any) -> bool: + actual = node.attrs.get(attr, match_utils.MISSING) + + if actual is match_utils.MISSING or actual is None: + return True + + return match_utils.values_equal(actual, expected) + + +def add_dims(left: Any, right: Any) -> Any: + if isinstance(left, int) and isinstance(right, int): + return left + right + + return left diff --git a/python/cactus/transpiler/IR/extra_matcher_generic.py b/python/cactus/transpiler/IR/extra_matcher_generic.py new file mode 100644 index 000000000..608200d82 --- /dev/null +++ b/python/cactus/transpiler/IR/extra_matcher_generic.py @@ -0,0 +1,208 @@ +from __future__ import annotations + +from typing import Any + +from . import models, match_utils +from .extra_matcher_common import * +from ..Fusions import models as FModels + + +def match_node_attr_equals(source: models.Node, graph: models.Graph, fusion: FModels.FusionGraph, bindings: dict[str, models.Node], spec: dict[str, Any]) -> bool: + """ + Checks that one matched node has a specific normalized attr value. + + The spec names a synthetic node, an attr key, and the expected value. The + function resolves the real node from bindings, reads its normalized attrs, + and compares with `values_equal`. Missing attrs fail unless + `allow_missing=True` is present in the spec. + """ + node = bindings.get(spec["node"]) + + if node is None: + return False + + actual = node.attrs.get(spec["attr"], match_utils.MISSING) + + if actual is match_utils.MISSING: + return bool(spec.get("allow_missing", False)) + + return match_utils.values_equal(actual, spec["value"]) + + +def match_node_attrs_equal(source: models.Node, graph: models.Graph, fusion: FModels.FusionGraph, bindings: dict[str, models.Node], spec: dict[str, Any]) -> bool: + """ + Checks that attrs on two matched nodes are equal to each other. + + This is useful for pattern checks like layernorm, where mean and variance + should reduce over the same dimension. The spec identifies left/right + synthetic nodes and attr names; both attrs must exist and compare equal + unless missing attrs are explicitly allowed. + """ + left_node = bindings.get(spec["left_node"]) + right_node = bindings.get(spec["right_node"]) + + if left_node is None or right_node is None: + return False + + left_value = left_node.attrs.get(spec["left_attr"], match_utils.MISSING) + right_value = right_node.attrs.get(spec["right_attr"], match_utils.MISSING) + + if left_value is match_utils.MISSING or right_value is match_utils.MISSING: + return bool(spec.get("allow_missing", False)) + + return match_utils.values_equal(left_value, right_value) + + +def match_input_value_kind(source: models.Node, graph: models.Graph, fusion: FModels.FusionGraph, bindings: dict[str, models.Node], spec: dict[str, Any]) -> bool: + """ + Checks the semantic kind of an external input to a fusion. + + The fusion graph declares inputs by role, such as `bias` or `weight`. This + matcher resolves the first real node for that role and verifies its + `value_kind` is in the allowed set, preventing activations from being + mistaken for parameters/buffers. + """ + input_node = match_utils.get_first_input_by_role(fusion, bindings, spec["role"]) + + if input_node is None: + return bool(spec.get("allow_missing", False)) + + return input_node.value_kind in spec["allowed_value_kinds"] + + +def match_parent_tensor_dim_equal(source: models.Node, graph: models.Graph, fusion: FModels.FusionGraph, bindings: dict[str, models.Node], spec: dict[str, Any]) -> bool: + """ + Compares one dimension from two parent tensors of matched nodes. + + The spec names the left/right synthetic nodes, which parent index to inspect + on each real node, and which shape dimension to compare. It is used for + checks such as attention query/key head-dim compatibility. Missing parents + or shapes follow the spec's `allow_missing` behavior. + """ + left_parent = get_bound_parent(bindings, spec["left_node"], spec["left_parent_index"]) + right_parent = get_bound_parent(bindings, spec["right_node"], spec["right_parent_index"]) + + if left_parent is None or right_parent is None: + return bool(spec.get("allow_missing", False)) + + left_dim = get_node_dim(left_parent, spec["left_dim"]) + right_dim = get_node_dim(right_parent, spec["right_dim"]) + + if left_dim is match_utils.MISSING or right_dim is match_utils.MISSING: + return bool(spec.get("allow_missing", False)) + + return match_utils.values_equal(left_dim, right_dim) + + +def match_input_broadcastable_to_node(source: models.Node, graph: models.Graph, fusion: FModels.FusionGraph, bindings: dict[str, models.Node], spec: dict[str, Any]) -> bool: + """ + Checks that a declared fusion input can broadcast to a target node shape. + + This is mainly for mask-like inputs. It resolves an input by role, resolves + the target synthetic node, reads both output shapes, and applies standard + broadcasting. Missing nodes or shapes are controlled by `allow_missing`. + """ + input_node = match_utils.get_first_input_by_role(fusion, bindings, spec["role"]) + target_node = bindings.get(spec["node"]) + + if input_node is None or target_node is None: + return bool(spec.get("allow_missing", False)) + + input_shape = match_utils.get_tensor_shape(input_node) + target_shape = match_utils.get_tensor_shape(target_node) + + if not input_shape or not target_shape: + return bool(spec.get("allow_missing", False)) + + return shapes_are_broadcastable(input_shape, target_shape) + + +def match_same_layer(source: models.Node, graph: models.Graph, fusion: FModels.FusionGraph, bindings: dict[str, models.Node], spec: dict[str, Any]) -> bool: + """ + Checks that multiple matched nodes appear to belong to the same layer. + + The spec provides synthetic node names. For each bound real node, we derive + a layer key from module stack/name/target text. The check passes only when + every derived key is identical, which helps avoid mixing cache/key/value + nodes from different decoder layers. + """ + layer_keys = [] + + for node_name in spec["nodes"]: + node = bindings.get(node_name) + + if node is None: + return False + + layer_key = get_layer_key(node) + + if layer_key is None: + return bool(spec.get("allow_missing", False)) + + layer_keys.append(layer_key) + + return len(set(layer_keys)) == 1 + + +def match_slice_halves(source: models.Node, graph: models.Graph, fusion: FModels.FusionGraph, bindings: dict[str, models.Node], spec: dict[str, Any]) -> bool: + """ + Checks that two slice nodes split the same tensor into adjacent equal halves. + + This protects decomposed GLU fusion. GLU is valid when the exported graph + computes `left * sigmoid(right)` where `left` and `right` are the first and + second halves of one source tensor along the same feature axis. + """ + left = bindings.get(spec["left_node"]) + right = bindings.get(spec["right_node"]) + + if left is None or right is None: + return False + + left_source = match_utils.get_parent(left, 0) + right_source = match_utils.get_parent(right, 0) + + if left_source is None or left_source is not right_source: + return False + + source_shape = match_utils.get_tensor_shape(left_source) + left_shape = match_utils.get_tensor_shape(left) + right_shape = match_utils.get_tensor_shape(right) + + if not source_shape or not left_shape or not right_shape: + return bool(spec.get("allow_missing", False)) + + left_axis = normalized_slice_axis(left, len(source_shape)) + right_axis = normalized_slice_axis(right, len(source_shape)) + + if left_axis is match_utils.MISSING or right_axis is match_utils.MISSING or left_axis != right_axis: + return False + + if not slice_step_is_one(left) or not slice_step_is_one(right): + return False + + left_size = match_utils.get_dim(left_shape, left_axis) + right_size = match_utils.get_dim(right_shape, right_axis) + source_size = match_utils.get_dim(source_shape, left_axis) + + if left_size is None or right_size is None: + return bool(spec.get("allow_missing", False)) + + if not match_utils.values_equal(left_size, right_size): + return False + + if not attr_matches(left, "start", 0, default=0): + return False + + if not attr_matches(right, "start", left_size): + return False + + if not optional_attr_matches(left, "end", left_size): + return False + + if not optional_attr_matches(right, "end", add_dims(left_size, right_size)): + return False + + if source_size is not None and not match_utils.values_equal(source_size, add_dims(left_size, right_size)): + return False + + return True diff --git a/python/cactus/transpiler/IR/extra_matcher_linear.py b/python/cactus/transpiler/IR/extra_matcher_linear.py new file mode 100644 index 000000000..49d3ee637 --- /dev/null +++ b/python/cactus/transpiler/IR/extra_matcher_linear.py @@ -0,0 +1,399 @@ +from __future__ import annotations + +from typing import Any + +from . import models, match_utils +from .extra_matcher_common import * +from ..Fusions import models as FModels + + +def classify_linear_rhs(node: models.Node) -> tuple[str, models.Node] | None: + """ + Classifies the right-hand side of a linear/matmul as direct or transposed. + + The node passed here is the actual RHS input to the matmul. The helper + delegates to the recursive implementation with an empty `seen` set so we + can safely walk through layout-only wrapper nodes without infinite loops. + """ + return classify_linear_rhs_(node, set()) + + +def classify_linear_rhs_(node: models.Node, seen: set[int]) -> tuple[str, models.Node] | None: + """ + Walks backward through simple RHS wrapper ops to find the source weight. + + It accepts either a direct parameter/buffer/lifted constant weight, or a + valid transpose/permute/t of such a weight. Copy-like wrappers such as + `_to_copy` and `contiguous` are skipped recursively. The return value is + `(layout_kind, source_weight)`, where layout_kind is `direct` or + `transposed`; None means the RHS is not safely recognizable as a weight. + """ + if id(node) in seen: + return None + + seen.add(id(node)) + + if node.target in COPY_TARGETS and len(node.parents) == 1: + return classify_linear_rhs_(node.parents[0], seen) + + if node.value_kind in WEIGHT_VALUE_KINDS: + return ("direct", node) + + if node.target in TRANSPOSE_TARGETS and len(node.parents) == 1 and node.parents[0].value_kind in WEIGHT_VALUE_KINDS: + if linear_transpose_is_valid(node, node.parents[0]): + return ("transposed", node.parents[0]) + + return None + + +def linear_transpose_is_valid(node: models.Node, source_weight: models.Node) -> bool: + """ + Checks whether a transpose-like op really swaps a 2D linear weight. + + Cactus needs to know whether the RHS has already been converted from + `[out_dim, in_dim]` to `[in_dim, out_dim]`. This accepts `t`, `transpose` + over dimensions 0/1, and `permute([1, 0])`, but rejects non-2D weights or + arbitrary permutations that would not be a standard linear transpose. + """ + source_shape = match_utils.get_tensor_shape(source_weight) + + if len(source_shape) != 2: + return False + + if node.target == "aten.t.default": + return True + + if node.target == "aten.transpose.int": + dim0 = node.attrs.get("dim0", match_utils.MISSING) + dim1 = node.attrs.get("dim1", match_utils.MISSING) + return {normalize_dim(dim0, 2), normalize_dim(dim1, 2)} == {0, 1} + + permutation = node.attrs.get("permutation", match_utils.MISSING) + return match_utils.values_equal(permutation, [1, 0]) + + +def linear_shapes_match(input_node: models.Node, rhs_node: models.Node, output_node: models.Node, source_weight: models.Node, layout_kind: str, spec: dict[str, Any]) -> bool: + """ + Verifies the dimensions around a linear/matmul are internally consistent. + + The RHS node is the tensor consumed by the real matmul, so its last two dims + should be `[input_dim, output_dim]`. If that RHS came from a transposed + original weight, the original source weight should be `[output_dim, + input_dim]`. Missing shapes are handled through the constraint spec's + `allow_missing` flag so callers can choose conservative or permissive mode. + """ + input_shape = match_utils.get_tensor_shape(input_node) + rhs_shape = match_utils.get_tensor_shape(rhs_node) + output_shape = match_utils.get_tensor_shape(output_node) + + if len(rhs_shape) < 2: + return bool(spec.get("allow_missing", False)) + + if not input_shape or not output_shape: + return bool(spec.get("allow_missing", False)) + + input_dim = input_shape[-1] + output_dim = output_shape[-1] + + if not match_utils.values_equal(rhs_shape[-2], input_dim): + return False + + if not match_utils.values_equal(rhs_shape[-1], output_dim): + return False + + source_shape = match_utils.get_tensor_shape(source_weight) + + if layout_kind == "transposed" and len(source_shape) >= 2: + if not match_utils.values_equal(source_shape[-2], output_dim): + return False + + if not match_utils.values_equal(source_shape[-1], input_dim): + return False + + return True + + +def compatible_shapes(left_shape: list[Any], right_shape: list[Any], spec: dict[str, Any]) -> bool: + """ + Checks whether two tensor shapes are equal or broadcast-compatible. + + This is used for things like RoPE cos/sin tables, where shapes may be + exactly identical or broadcastable across singleton dimensions. If either + shape is missing, the result is controlled by `allow_missing` in the spec. + """ + if not left_shape or not right_shape: + return bool(spec.get("allow_missing", False)) + + return ( + match_utils.values_equal(left_shape, right_shape) + or shapes_are_broadcastable(left_shape, right_shape) + or shapes_are_broadcastable(right_shape, left_shape) + ) + + +def compatible_dtypes(left_node: models.Node, right_node: models.Node, spec: dict[str, Any]) -> bool: + """ + Checks whether two nodes have matching output dtypes when required. + + Most table-like values in one fusion should have the same dtype. The check + can be disabled with `require_same_dtype=False`; otherwise missing dtype + metadata follows `allow_missing`, and available dtype strings are compared + with the shared `values_equal` normalization behavior. + """ + if not spec.get("require_same_dtype", True): + return True + + left_dtype = get_node_dtype(left_node) + right_dtype = get_node_dtype(right_node) + + if left_dtype is match_utils.MISSING or right_dtype is match_utils.MISSING: + return bool(spec.get("allow_missing", False)) + + return match_utils.values_equal(left_dtype, right_dtype) + + +def compatible_trig_sources(cos_node: models.Node, sin_node: models.Node, spec: dict[str, Any]) -> bool: + """ + Checks whether RoPE cos and sin tables appear to come from the same angles. + + The helper walks backward from the cos/sin input nodes until it finds + `aten.cos` and `aten.sin`, then compares the node feeding those trig ops. + If no trig source is visible, the matcher can either allow that partial + evidence or fail closed with `require_trig_sources=True`. + """ + cos_sources = collect_trig_angle_sources(cos_node, COS_TARGETS, int(spec.get("max_trig_depth", 8))) + sin_sources = collect_trig_angle_sources(sin_node, SIN_TARGETS, int(spec.get("max_trig_depth", 8))) + + if cos_sources or sin_sources: + return cos_sources == sin_sources + + return not spec.get("require_trig_sources", False) + + +def collect_trig_angle_sources(node: models.Node, trig_targets: set[str], max_depth: int) -> frozenset[str]: + """ + Finds the angle-producing parent nodes for visible trig ops upstream. + + Starting from `node`, this walks backward through parents up to `max_depth`. + When it reaches a target such as `aten.cos.default` or `aten.sin.default`, + it records that trig node's first parent as the angle source. The result is + a frozenset so cos and sin paths can be compared independent of ordering. + """ + sources: set[str] = set() + seen: set[int] = set() + + def visit(current: models.Node, depth: int) -> None: + """ + Recursively walks parents while avoiding cycles and depth blowups. + + The exported graph should be a DAG, but the `seen` guard makes this + helper robust if malformed graph data or future graph rewrites ever + introduce repeated references. + """ + if depth > max_depth or id(current) in seen: + return + + seen.add(id(current)) + + if current.target in trig_targets: + sources.add(current.parents[0].name if current.parents else current.name) + return + + for parent in current.parents: + visit(parent, depth + 1) + + visit(node, 0) + return frozenset(sources) + + +def rope_tables_match_input(x_node: models.Node, cos_shape: list[Any], sin_shape: list[Any], spec: dict[str, Any]) -> bool: + """ + Checks whether RoPE cos/sin table shapes can apply to the rotated input. + + RoPE tables often have fewer broadcast dimensions than the activation, and + the final table dimension can represent either the full hidden/head dim or + half of it depending on the rotate-half implementation. This delegates the + actual table-vs-input rule to `rope_table_shape_matches_x` for both tables. + """ + x_shape = match_utils.get_tensor_shape(x_node) + + if not x_shape: + return bool(spec.get("allow_missing", False)) + + return rope_table_shape_matches_x(cos_shape, x_shape) and rope_table_shape_matches_x(sin_shape, x_shape) + + +def rope_table_shape_matches_x(table_shape: list[Any], x_shape: list[Any]) -> bool: + """ + Checks one RoPE table shape against the activation being rotated. + + The table is accepted if it is broadcastable to the activation shape. If + not, we allow the common rotate-half form where the table's last dimension + is half of the activation's last dimension. Symbolic or missing dimensions + are intentionally conservative here. + """ + if not table_shape: + return False + + if shapes_are_broadcastable(table_shape, x_shape): + return True + + table_last = table_shape[-1] + x_last = x_shape[-1] + + if isinstance(table_last, int) and isinstance(x_last, int) and table_last > 0: + return x_last in {table_last, table_last * 2} + + return False + + +def gate_shape_matches(gate_node: models.Node, gate_count: int, spec: dict[str, Any]) -> bool: + """ + Checks whether an LSTM combined-gate tensor can hold `gate_count` gates. + + LSTM gate projections commonly produce a last dimension equal to + `4 * hidden_size`. We do not know hidden_size directly here, so this checks + that the final dimension is divisible by the expected gate count. Symbolic + dimensions can be allowed through `allow_symbolic`. + """ + gate_shape = match_utils.get_tensor_shape(gate_node) + + if not gate_shape: + return bool(spec.get("allow_missing", False)) + + hidden_dim = gate_shape[-1] + + if not isinstance(hidden_dim, int): + return bool(spec.get("allow_symbolic", True)) + + return hidden_dim % gate_count == 0 + + +def has_explicit_gate_split(gate_node: models.Node, gate_count: int, max_depth: int) -> bool: + """ + Looks for a visible split/chunk/unbind that separates combined LSTM gates. + + Shape divisibility alone is not enough to prove LSTM semantics. This walks + forward from the combined gate node and requires a split-like op whose attr + or fanout indicates the expected number of gates. If no such split is + visible, the LSTM fusion fails closed. + """ + for node in walk_children(gate_node, max_depth): + if node.target not in SPLIT_TARGETS: + continue + + if split_node_has_gate_count(node, gate_count): + return True + + return False + + +def split_node_has_gate_count(node: models.Node, gate_count: int) -> bool: + """ + Determines whether a split-like node creates the expected number of gates. + + Different export paths encode splits differently: `unbind` is inferred from + child count, `chunk` stores the chunk count as positional attr `arg_1`, and + `split_with_sizes` stores the explicit split sizes. If those attrs are not + available, child count is used as the fallback evidence. + """ + if node.target == "aten.unbind.int": + return len(node.children) == gate_count + + if node.target == "aten.chunk.default": + return match_utils.values_equal(node.attrs.get("arg_1", match_utils.MISSING), gate_count) + + split_sizes = node.attrs.get("arg_1", match_utils.MISSING) + + if isinstance(split_sizes, (list, tuple)): + return len(split_sizes) == gate_count + + return len(node.children) == gate_count + + +def match_linear_weight_layout(source: models.Node, graph: models.Graph, fusion: FModels.FusionGraph, bindings: dict[str, models.Node], spec: dict[str, Any]) -> bool: + """ + Verifies that a linear/matmul RHS is a compatible model weight layout. + + The matcher resolves the fusion's input, weight, and output nodes by role. + It then classifies the weight path as direct or transposed and checks that + dimensions line up with the matmul output. This is what prevents a random + activation matrix from being treated as a Cactus linear weight. + """ + input_node = match_utils.get_first_input_by_role(fusion, bindings, spec["input_role"]) + weight_node = match_utils.get_first_input_by_role(fusion, bindings, spec["weight_role"]) + output_node = bindings.get(spec["output_node"]) + + if input_node is None or weight_node is None or output_node is None: + return bool(spec.get("allow_missing", False)) + + layout = classify_linear_rhs(weight_node) + + if layout is None: + return False + + layout_kind, source_weight = layout + allowed_layouts = spec.get("allowed_layouts", ("direct", "transposed")) + + if layout_kind not in allowed_layouts: + return False + + return linear_shapes_match(input_node, weight_node, output_node, source_weight, layout_kind, spec) + + +def match_rope_tables_compatible(source: models.Node, graph: models.Graph, fusion: FModels.FusionGraph, bindings: dict[str, models.Node], spec: dict[str, Any]) -> bool: + """ + Verifies that RoPE cos/sin inputs look like compatible rotation tables. + + It gets the matched x/cos/sin values by role, checks cos/sin shape and dtype + compatibility, optionally proves they came from matching angle sources, and + then checks that both tables can broadcast or half-dim match the RoPE input. + """ + x_node = match_utils.get_first_input_by_role(fusion, bindings, spec["x_role"]) + cos_node = match_utils.get_first_input_by_role(fusion, bindings, spec["cos_role"]) + sin_node = match_utils.get_first_input_by_role(fusion, bindings, spec["sin_role"]) + + if cos_node is None or sin_node is None: + return bool(spec.get("allow_missing", False)) + + cos_shape = match_utils.get_tensor_shape(cos_node) + sin_shape = match_utils.get_tensor_shape(sin_node) + + if not compatible_shapes(cos_shape, sin_shape, spec): + return False + + if not compatible_dtypes(cos_node, sin_node, spec): + return False + + if not compatible_trig_sources(cos_node, sin_node, spec): + return False + + if x_node is not None and not rope_tables_match_input(x_node, cos_shape, sin_shape, spec): + return False + + return True + + +def match_lstm_gate_structure(source: models.Node, graph: models.Graph, fusion: FModels.FusionGraph, bindings: dict[str, models.Node], spec: dict[str, Any]) -> bool: + """ + Verifies that an LSTM gate node has real four-gate structure evidence. + + The matcher first checks that the combined gate tensor's final dimension is + divisible by the expected gate count. By default it then requires an + explicit downstream split/chunk/unbind into that many gates, so LSTM fusion + fails closed when the graph only vaguely resembles an LSTM. + """ + gate_node = bindings.get(spec["gate_node"]) + + if gate_node is None: + return bool(spec.get("allow_missing", False)) + + gate_count = int(spec.get("gate_count", 4)) + + if not gate_shape_matches(gate_node, gate_count, spec): + return False + + if spec.get("require_explicit_gate_split", True): + return has_explicit_gate_split(gate_node, gate_count, int(spec.get("max_depth", 4))) + + return True diff --git a/python/cactus/transpiler/IR/extra_matcher_moe.py b/python/cactus/transpiler/IR/extra_matcher_moe.py new file mode 100644 index 000000000..6018f6cb3 --- /dev/null +++ b/python/cactus/transpiler/IR/extra_matcher_moe.py @@ -0,0 +1,434 @@ +from __future__ import annotations + +from typing import Any + +from . import models, match_utils +from .extra_matcher_common import * +from ..Fusions import models as FModels + + +def get_repeated_subgraph(fusion: FModels.FusionGraph, name: str) -> FModels.RepeatedSubgraph | None: + """ + Fetches a named repeated subgraph declaration from a fusion graph. + + MoE fusions store expert branches as repeated subgraphs. This helper finds + the declaration by name so the MoE-specific matchers can reuse its pattern + and constraints when collecting real expert-branch bindings. + """ + for subgraph in fusion.repeated_subgraphs: + if subgraph.name == name: + return subgraph + + return None + + +def collect_repeated_subgraph_bindings(graph: models.Graph, subgraph: FModels.RepeatedSubgraph) -> tuple[dict[str, models.Node], ...]: + """ + Finds every real subgraph binding that matches a repeated pattern. + + This scans each node as a possible root for `subgraph.graph`, runs the same + graph binder used by normal fusion matching, then validates the full set of + fusion matchers against that candidate. The result is not just a count: it + returns each binding dictionary so later matchers can inspect which real + nodes were bound to expert_gate, expert_weighted, routing_weight, etc. + """ + from . import match + + matches: list[dict[str, models.Node]] = [] + + for candidate in graph.nodes: + candidate_bindings = match_utils.bind_fusion_graph(candidate, subgraph.graph, match.match_nodes) + + if candidate_bindings is None: + continue + + if not match.match_fusion_bindings(candidate, graph, subgraph.graph, candidate_bindings): + continue + + matches.append(candidate_bindings) + + return tuple(matches) + + +def relevant_moe_branch_bindings(graph: models.Graph, fusion: FModels.FusionGraph, bindings: dict[str, models.Node], spec: dict[str, Any]) -> tuple[dict[str, models.Node], ...]: + """ + Filters repeated expert branches down to ones driven by MoE routing. + + The first step collects every expert-like branch in the whole graph. Then + this keeps only branches whose declared `routing_weight` input is reachable + from a top-k node or that top-k node's probability input. This is the key + safety check that prevents us from fusing arbitrary MLP-looking branches as + MoE experts when they are not actually selected by the router. + """ + subgraph = get_repeated_subgraph(fusion, spec["repeated_subgraph"]) + + if subgraph is None: + return () + + topk_node = bindings.get(spec.get("topk_node", "topk")) + combine_node = bindings.get(spec.get("combine_node", "moe_combine")) + topk_nodes = (topk_node,) if topk_node is not None else find_reachable_topk_nodes(graph, combine_node, spec) + + if not topk_nodes: + return () + + max_depth = int(spec.get("max_depth", 16)) + branch_bindings = collect_repeated_subgraph_bindings(graph, subgraph) + relevant = [] + + for branch_binding in branch_bindings: + routing_weight = match_utils.get_first_input_by_role(subgraph.graph, branch_binding, spec["routing_weight_role"]) + + if routing_weight is None: + continue + + routing_sources = tuple(source for topk in topk_nodes for source in (topk, first_parent(topk)) if source is not None) + + if not has_any_path(routing_sources, routing_weight, max_depth): + continue + + relevant.append(branch_binding) + + return tuple(relevant) + + +def required_moe_branch_count(subgraph: FModels.RepeatedSubgraph, topk_node: models.Node | None, spec: dict[str, Any]) -> int: + """ + Determines how many routed expert branches must be proven. + + The spec may override this directly with `min_count`. Otherwise, if a top-k + node exposes its `k` attr, that is the strongest graph-local value because + it tells us how many experts each token selects. If neither exists, the + repeated subgraph's own `min_count` is used. + """ + if "min_count" in spec: + return int(spec["min_count"]) + + if topk_node is not None and spec.get("min_count_from_topk", True): + k = topk_node.attrs.get("k", match_utils.MISSING) + + if isinstance(k, int): + return k + + return subgraph.min_count + + +def get_num_experts_from_inputs(fusion: FModels.FusionGraph, bindings: dict[str, models.Node], spec: dict[str, Any]) -> int | None: + """ + Infers the MoE expert count from weight/bias shapes. + + For LFM grouped MoE, router bias is `[num_experts]`, gate/up expert weight + is `[num_experts, 2 * intermediate, hidden]`, and down expert weight is + `[num_experts, hidden, intermediate]`. Matching these first dimensions is + a strong safety check that the grouped matmuls are really expert weights. + """ + candidates: list[int] = [] + + for role in (spec["gate_up_weight_role"], spec["down_weight_role"], spec.get("router_bias_role")): + if role is None: + continue + + input_node = match_utils.get_first_input_by_role(fusion, bindings, role) + shape = match_utils.get_tensor_shape(input_node) + + if not shape: + continue + + dim = get_known_int(shape[0]) + + if dim is not None: + candidates.append(dim) + + if not candidates: + return None + + if len(set(candidates)) != 1: + return None + + return candidates[0] + + +def split_sizes_are_valid(split_node: models.Node, gate_up_node: models.Node, down_weight: models.Node | None, spec: dict[str, Any]) -> bool: + """ + Checks the grouped gate/up projection split. + + LFM projects gate and up together, then splits the final dimension into two + equal chunks. The second chunk should match the down-projection expert + input dimension when that shape is available. + """ + split_sizes = get_attr(split_node, "split_sizes", "arg_1") + + if not isinstance(split_sizes, (list, tuple)) or len(split_sizes) != 2: + return bool(spec.get("allow_missing", False)) + + left_size = get_known_int(split_sizes[0]) + right_size = get_known_int(split_sizes[1]) + + if left_size is None or right_size is None: + return bool(spec.get("allow_symbolic", True)) + + if left_size != right_size: + return False + + gate_up_output_dim = get_shape_dim(gate_up_node, -1) + + if gate_up_output_dim is not None and gate_up_output_dim != left_size + right_size: + return False + + down_input_dim = get_shape_dim(down_weight, -1) + + if down_input_dim is not None and down_input_dim != right_size: + return False + + return True + + +def grouped_moe_rows_match_topk(hidden_node: models.Node | None, gate_up_node: models.Node, topk_node: models.Node, spec: dict[str, Any]) -> bool: + """ + Checks that routed grouped-mm rows equal token count times top-k. + + This links the routing decision to the grouped expert matmul size. Missing + or symbolic dimensions can be allowed by `allow_missing` because shape + metadata may be incomplete in some exported graphs. + """ + token_count = get_shape_dim(hidden_node, 0) + grouped_rows = get_shape_dim(gate_up_node, 0) + topk = get_known_int(get_attr(topk_node, "k", "arg_1")) + + if token_count is None or grouped_rows is None or topk is None: + return bool(spec.get("allow_missing", False)) + + return grouped_rows == token_count * topk + + +def grouped_moe_weight_permute_is_valid(node: models.Node | None) -> bool: + """ + Checks the expert weight permutation used before grouped_mm. + + The LFM export stores expert weights with expert dimension first and then + permutes the last two dims to feed grouped_mm. That appears as + `permute([0, 2, 1])`. + """ + if node is None: + return False + + return match_utils.values_equal(get_attr(node, "permutation", "arg_1"), [0, 2, 1]) + + +def find_reachable_topk_nodes(graph: models.Graph, combine_node: models.Node | None, spec: dict[str, Any]) -> tuple[models.Node, ...]: + """ + Finds top-k router nodes that can flow into the MoE combine output. + + Some MoE fusion graphs do not bind top-k as a direct synthetic node because + the exported routing path can contain indexing/scatter/gather wrappers. + This helper searches the real graph for `aten.topk.default` nodes that have + a forward path to the matched combine node within the configured depth. + """ + if combine_node is None: + return () + + max_depth = int(spec.get("max_depth", 16)) + return tuple(node for node in graph.nodes if node.target == "aten.topk.default" and has_path(node, combine_node, max_depth)) + + +def first_parent(node: models.Node) -> models.Node | None: + """ + Returns the first parent of a node, if it has one. + + For top-k, parent 0 is usually the routing probabilities tensor. The MoE + routing matcher treats both the top-k node and its first parent as valid + routing sources because exported graphs may use either values or indices + through different wrapper chains. + """ + return node.parents[0] if node.parents else None + + +def branch_output_reaches_combine(branch_binding: dict[str, models.Node], subgraph: FModels.RepeatedSubgraph, combine_node: models.Node, spec: dict[str, Any]) -> bool: + """ + Checks whether one matched expert branch contributes to the combine node. + + The branch's exposed weighted output should eventually flow into the outer + MoE combine/sum/scatter-add node. This proves that a routed expert branch is + actually part of the candidate MoE output instead of being an unrelated + branch elsewhere in the graph. + """ + weighted_node_name = spec.get("weighted_node", subgraph.graph.root) + weighted_node = branch_binding.get(weighted_node_name) + + if weighted_node is None: + return False + + return has_path(weighted_node, combine_node, int(spec.get("max_depth", 16))) + + +def has_any_path(sources: tuple[models.Node | None, ...], dest: models.Node, max_depth: int) -> bool: + """ + Returns True if any non-None source has a forward path to `dest`. + + This is a small convenience wrapper around `has_path` for cases where a + matcher has multiple plausible routing sources, such as top-k itself and + the routing-probability parent of top-k. + """ + for source in sources: + if source is not None and has_path(source, dest, max_depth): + return True + + return False + + +def has_path(source: models.Node, dest: models.Node, max_depth: int) -> bool: + """ + Performs a bounded forward reachability check in the IR DAG. + + The function starts at `source`, walks through `children`, and succeeds if + it reaches `dest` within `max_depth`. The visited set prevents repeated + work on shared DAG nodes and protects against accidental cycles. + """ + if source is dest: + return True + + seen: set[int] = set() + + def visit(node: models.Node, depth: int) -> bool: + if depth > max_depth or id(node) in seen: + return False + + seen.add(id(node)) + + for child in node.children: + if child is dest: + return True + + if visit(child, depth + 1): + return True + + return False + + return visit(source, 0) + + +def match_moe_expert_branch_routing(source: models.Node, graph: models.Graph, fusion: FModels.FusionGraph, bindings: dict[str, models.Node], spec: dict[str, Any]) -> bool: + """ + Verifies that enough matched expert branches are driven by router top-k. + + This locates the repeated expert-branch pattern, finds or discovers a top-k + node connected to the MoE combine output, then collects expert branch + bindings whose routing-weight input is reachable from that routing source. + The match succeeds only if the number of routed branches meets the expected + count from `topk.k`, an explicit spec override, or the subgraph minimum. + """ + subgraph = get_repeated_subgraph(fusion, spec["repeated_subgraph"]) + combine_node = bindings.get(spec.get("combine_node", "moe_combine")) + topk_node = bindings.get(spec.get("topk_node", "topk")) + + if subgraph is None: + return bool(spec.get("allow_missing", False)) + + if topk_node is None: + topk_nodes = find_reachable_topk_nodes(graph, combine_node, spec) + topk_node = topk_nodes[0] if topk_nodes else None + + if topk_node is None: + return bool(spec.get("allow_missing", False)) + + branch_bindings = relevant_moe_branch_bindings(graph, fusion, bindings, spec) + min_count = required_moe_branch_count(subgraph, topk_node, spec) + + return len(branch_bindings) >= min_count + + +def match_moe_routing_weights_combine(source: models.Node, graph: models.Graph, fusion: FModels.FusionGraph, bindings: dict[str, models.Node], spec: dict[str, Any]) -> bool: + """ + Verifies that routed expert outputs are actually combined into MoE output. + + This reuses the routed expert-branch filtering, then checks that each + branch's weighted expert output has a forward path to the matched combine + node. That proves both sides of MoE semantics: routing weights are used on + expert outputs, and those weighted outputs feed the final combine op. + """ + subgraph = get_repeated_subgraph(fusion, spec["repeated_subgraph"]) + topk_node = bindings.get(spec.get("topk_node", "topk")) + combine_node = bindings.get(spec["combine_node"]) + + if subgraph is None or combine_node is None: + return bool(spec.get("allow_missing", False)) + + if topk_node is None: + topk_nodes = find_reachable_topk_nodes(graph, combine_node, spec) + topk_node = topk_nodes[0] if topk_nodes else None + + if topk_node is None: + return bool(spec.get("allow_missing", False)) + + branch_bindings = relevant_moe_branch_bindings(graph, fusion, bindings, spec) + combined_branch_count = sum( + 1 for branch_binding in branch_bindings + if branch_output_reaches_combine(branch_binding, subgraph, combine_node, spec) + ) + + return combined_branch_count >= required_moe_branch_count(subgraph, topk_node, spec) + + +def match_grouped_moe_structure(source: models.Node, graph: models.Graph, fusion: FModels.FusionGraph, bindings: dict[str, models.Node], spec: dict[str, Any]) -> bool: + """ + Verifies the grouped-mm MoE structure exported by LFM-style MoE models. + + This matcher handles the compact representation where expert branches are + not expanded into separate per-expert matmuls. The graph edges prove the + local dataflow; this semantic check verifies top-k routing, expert-count + consistency, shared grouped-mm offsets, valid gate/up splitting, and final + top-k combine shape. + """ + topk_node = bindings.get(spec["topk_node"]) + topk_indices_node = bindings.get(spec["topk_indices_node"]) + gate_up_grouped_node = bindings.get(spec["gate_up_grouped_node"]) + down_grouped_node = bindings.get(spec["down_grouped_node"]) + split_node = bindings.get(spec["split_node"]) + offsets_node = bindings.get(spec["offsets_node"]) + combine_node = bindings.get(spec["combine_node"]) + hidden_node = bindings.get("moe_hidden_view") + gate_up_weight_transpose = bindings.get("moe_gate_up_weight_transpose") + down_weight_transpose = bindings.get("moe_down_weight_transpose") + down_weight = match_utils.get_first_input_by_role(fusion, bindings, spec["down_weight_role"]) + + if any(node is None for node in (topk_node, topk_indices_node, gate_up_grouped_node, down_grouped_node, split_node, offsets_node, combine_node)): + return bool(spec.get("allow_missing", False)) + + if get_known_int(get_attr(topk_node, "k", "arg_1")) is None: + return False + + if not match_utils.values_equal(get_attr(topk_indices_node, "index", "arg_1"), 1): + return False + + if not grouped_moe_weight_permute_is_valid(gate_up_weight_transpose): + return False + + if not grouped_moe_weight_permute_is_valid(down_weight_transpose): + return False + + gate_up_offsets = match_utils.get_parent(gate_up_grouped_node, 2) + down_offsets = match_utils.get_parent(down_grouped_node, 2) + + if gate_up_offsets is not offsets_node or down_offsets is not offsets_node: + return False + + if not match_utils.values_equal(get_attr(offsets_node, "dim", "arg_1"), 0): + return False + + if not split_sizes_are_valid(split_node, gate_up_grouped_node, down_weight, spec): + return False + + num_experts = get_num_experts_from_inputs(fusion, bindings, spec) + histc_node = bindings.get("moe_histc") + histc_bins = get_known_int(get_attr(histc_node, "bins", "arg_1")) if histc_node is not None else None + + if num_experts is not None and histc_bins is not None and num_experts != histc_bins: + return False + + if not grouped_moe_rows_match_topk(hidden_node, gate_up_grouped_node, topk_node, spec): + return False + + if not match_utils.values_equal(get_attr(combine_node, "dim", "arg_1"), 1): + return False + + return True diff --git a/python/cactus/transpiler/IR/extra_matchers.py b/python/cactus/transpiler/IR/extra_matchers.py index 2762088ef..c3038cf60 100644 --- a/python/cactus/transpiler/IR/extra_matchers.py +++ b/python/cactus/transpiler/IR/extra_matchers.py @@ -1,1643 +1,81 @@ -from typing import Any, Callable - -from . import models, match_utils -from ..Fusions import models as FModels - -#TODO: Clean up this file -ExtraMatcher = Callable[[models.Node, models.Graph, FModels.FusionGraph, dict[str, models.Node], dict[str, Any]], bool] - - -WEIGHT_VALUE_KINDS = ( - FModels.ValueKind.PARAMETER, - FModels.ValueKind.BUFFER, - FModels.ValueKind.LIFTED_CONSTANT, -) - -COPY_TARGETS = { - "aten.clone.default", - "aten._to_copy.default", - "aten.contiguous.default", -} - -TRANSPOSE_TARGETS = { - "aten.t.default", - "aten.transpose.int", - "aten.permute.default", -} - -COS_TARGETS = {"aten.cos.default"} -SIN_TARGETS = {"aten.sin.default"} -SPLIT_TARGETS = { - "aten.chunk.default", - "aten.split.Tensor", - "aten.split_with_sizes.default", - "aten.unbind.int", -} - - -############################################ Extra Matcher Utils!!!!! ############################################ -def get_bound_parent(bindings: dict[str, models.Node], node_name: str, parent_index: int) -> models.Node | None: - """ - Returns a real parent node from a bound synthetic node/input position. - - `bindings` maps synthetic fusion-node names to real IR nodes. This helper - first finds the real node bound to `node_name`, then asks match_utils for - the parent at `parent_index`. It returns None when the synthetic node was - not bound or the requested parent index does not exist. - """ - node = bindings.get(node_name) - - if node is None: - return None - - return match_utils.get_parent(node, parent_index) - - -def get_node_dim(node: models.Node, dim_index: int) -> Any: - """ - Returns one tensor dimension from a real IR node's output shape. - - This reads `tensor_output_meta["shape"]`, supports negative indices in the - Python style, and returns `match_utils.MISSING` instead of raising when the - shape is unavailable or the dimension index is invalid. Matchers use that - missing sentinel so they can either fail closed or honor `allow_missing`. - """ - shape = match_utils.get_tensor_shape(node) - - if not shape: - return match_utils.MISSING - - if dim_index < 0: - dim_index += len(shape) - - if dim_index < 0 or dim_index >= len(shape): - return match_utils.MISSING - - return shape[dim_index] - - -def classify_linear_rhs(node: models.Node) -> tuple[str, models.Node] | None: - """ - Classifies the right-hand side of a linear/matmul as direct or transposed. - - The node passed here is the actual RHS input to the matmul. The helper - delegates to the recursive implementation with an empty `seen` set so we - can safely walk through layout-only wrapper nodes without infinite loops. - """ - return classify_linear_rhs_(node, set()) - - -def classify_linear_rhs_(node: models.Node, seen: set[int]) -> tuple[str, models.Node] | None: - """ - Walks backward through simple RHS wrapper ops to find the source weight. - - It accepts either a direct parameter/buffer/lifted constant weight, or a - valid transpose/permute/t of such a weight. Copy-like wrappers such as - `_to_copy` and `contiguous` are skipped recursively. The return value is - `(layout_kind, source_weight)`, where layout_kind is `direct` or - `transposed`; None means the RHS is not safely recognizable as a weight. - """ - if id(node) in seen: - return None - - seen.add(id(node)) - - if node.target in COPY_TARGETS and len(node.parents) == 1: - return classify_linear_rhs_(node.parents[0], seen) - - if node.value_kind in WEIGHT_VALUE_KINDS: - return ("direct", node) - - if node.target in TRANSPOSE_TARGETS and len(node.parents) == 1 and node.parents[0].value_kind in WEIGHT_VALUE_KINDS: - if linear_transpose_is_valid(node, node.parents[0]): - return ("transposed", node.parents[0]) - - return None - - -def linear_transpose_is_valid(node: models.Node, source_weight: models.Node) -> bool: - """ - Checks whether a transpose-like op really swaps a 2D linear weight. - - Cactus needs to know whether the RHS has already been converted from - `[out_dim, in_dim]` to `[in_dim, out_dim]`. This accepts `t`, `transpose` - over dimensions 0/1, and `permute([1, 0])`, but rejects non-2D weights or - arbitrary permutations that would not be a standard linear transpose. - """ - source_shape = match_utils.get_tensor_shape(source_weight) - - if len(source_shape) != 2: - return False - - if node.target == "aten.t.default": - return True - - if node.target == "aten.transpose.int": - dim0 = node.attrs.get("dim0", match_utils.MISSING) - dim1 = node.attrs.get("dim1", match_utils.MISSING) - return {normalize_dim(dim0, 2), normalize_dim(dim1, 2)} == {0, 1} - - permutation = node.attrs.get("permutation", match_utils.MISSING) - return match_utils.values_equal(permutation, [1, 0]) - - -def linear_shapes_match(input_node: models.Node, rhs_node: models.Node, output_node: models.Node, source_weight: models.Node, layout_kind: str, spec: dict[str, Any]) -> bool: - """ - Verifies the dimensions around a linear/matmul are internally consistent. - - The RHS node is the tensor consumed by the real matmul, so its last two dims - should be `[input_dim, output_dim]`. If that RHS came from a transposed - original weight, the original source weight should be `[output_dim, - input_dim]`. Missing shapes are handled through the constraint spec's - `allow_missing` flag so callers can choose conservative or permissive mode. - """ - input_shape = match_utils.get_tensor_shape(input_node) - rhs_shape = match_utils.get_tensor_shape(rhs_node) - output_shape = match_utils.get_tensor_shape(output_node) - - if len(rhs_shape) < 2: - return bool(spec.get("allow_missing", False)) - - if not input_shape or not output_shape: - return bool(spec.get("allow_missing", False)) - - input_dim = input_shape[-1] - output_dim = output_shape[-1] - - if not match_utils.values_equal(rhs_shape[-2], input_dim): - return False - - if not match_utils.values_equal(rhs_shape[-1], output_dim): - return False - - source_shape = match_utils.get_tensor_shape(source_weight) - - if layout_kind == "transposed" and len(source_shape) >= 2: - if not match_utils.values_equal(source_shape[-2], output_dim): - return False - - if not match_utils.values_equal(source_shape[-1], input_dim): - return False - - return True - - -def compatible_shapes(left_shape: list[Any], right_shape: list[Any], spec: dict[str, Any]) -> bool: - """ - Checks whether two tensor shapes are equal or broadcast-compatible. - - This is used for things like RoPE cos/sin tables, where shapes may be - exactly identical or broadcastable across singleton dimensions. If either - shape is missing, the result is controlled by `allow_missing` in the spec. - """ - if not left_shape or not right_shape: - return bool(spec.get("allow_missing", False)) - - return ( - match_utils.values_equal(left_shape, right_shape) - or shapes_are_broadcastable(left_shape, right_shape) - or shapes_are_broadcastable(right_shape, left_shape) - ) - - -def compatible_dtypes(left_node: models.Node, right_node: models.Node, spec: dict[str, Any]) -> bool: - """ - Checks whether two nodes have matching output dtypes when required. - - Most table-like values in one fusion should have the same dtype. The check - can be disabled with `require_same_dtype=False`; otherwise missing dtype - metadata follows `allow_missing`, and available dtype strings are compared - with the shared `values_equal` normalization behavior. - """ - if not spec.get("require_same_dtype", True): - return True - - left_dtype = get_node_dtype(left_node) - right_dtype = get_node_dtype(right_node) - - if left_dtype is match_utils.MISSING or right_dtype is match_utils.MISSING: - return bool(spec.get("allow_missing", False)) - - return match_utils.values_equal(left_dtype, right_dtype) - - -def compatible_trig_sources(cos_node: models.Node, sin_node: models.Node, spec: dict[str, Any]) -> bool: - """ - Checks whether RoPE cos and sin tables appear to come from the same angles. - - The helper walks backward from the cos/sin input nodes until it finds - `aten.cos` and `aten.sin`, then compares the node feeding those trig ops. - If no trig source is visible, the matcher can either allow that partial - evidence or fail closed with `require_trig_sources=True`. - """ - cos_sources = collect_trig_angle_sources(cos_node, COS_TARGETS, int(spec.get("max_trig_depth", 8))) - sin_sources = collect_trig_angle_sources(sin_node, SIN_TARGETS, int(spec.get("max_trig_depth", 8))) - - if cos_sources or sin_sources: - return cos_sources == sin_sources - - return not spec.get("require_trig_sources", False) - - -def collect_trig_angle_sources(node: models.Node, trig_targets: set[str], max_depth: int) -> frozenset[str]: - """ - Finds the angle-producing parent nodes for visible trig ops upstream. - - Starting from `node`, this walks backward through parents up to `max_depth`. - When it reaches a target such as `aten.cos.default` or `aten.sin.default`, - it records that trig node's first parent as the angle source. The result is - a frozenset so cos and sin paths can be compared independent of ordering. - """ - sources: set[str] = set() - seen: set[int] = set() - - def visit(current: models.Node, depth: int) -> None: - """ - Recursively walks parents while avoiding cycles and depth blowups. - - The exported graph should be a DAG, but the `seen` guard makes this - helper robust if malformed graph data or future graph rewrites ever - introduce repeated references. - """ - if depth > max_depth or id(current) in seen: - return - - seen.add(id(current)) - - if current.target in trig_targets: - sources.add(current.parents[0].name if current.parents else current.name) - return - - for parent in current.parents: - visit(parent, depth + 1) - - visit(node, 0) - return frozenset(sources) - - -def rope_tables_match_input(x_node: models.Node, cos_shape: list[Any], sin_shape: list[Any], spec: dict[str, Any]) -> bool: - """ - Checks whether RoPE cos/sin table shapes can apply to the rotated input. - - RoPE tables often have fewer broadcast dimensions than the activation, and - the final table dimension can represent either the full hidden/head dim or - half of it depending on the rotate-half implementation. This delegates the - actual table-vs-input rule to `rope_table_shape_matches_x` for both tables. - """ - x_shape = match_utils.get_tensor_shape(x_node) - - if not x_shape: - return bool(spec.get("allow_missing", False)) - - return rope_table_shape_matches_x(cos_shape, x_shape) and rope_table_shape_matches_x(sin_shape, x_shape) - - -def rope_table_shape_matches_x(table_shape: list[Any], x_shape: list[Any]) -> bool: - """ - Checks one RoPE table shape against the activation being rotated. - - The table is accepted if it is broadcastable to the activation shape. If - not, we allow the common rotate-half form where the table's last dimension - is half of the activation's last dimension. Symbolic or missing dimensions - are intentionally conservative here. - """ - if not table_shape: - return False - - if shapes_are_broadcastable(table_shape, x_shape): - return True - - table_last = table_shape[-1] - x_last = x_shape[-1] - - if isinstance(table_last, int) and isinstance(x_last, int) and table_last > 0: - return x_last in {table_last, table_last * 2} - - return False - - -def gate_shape_matches(gate_node: models.Node, gate_count: int, spec: dict[str, Any]) -> bool: - """ - Checks whether an LSTM combined-gate tensor can hold `gate_count` gates. - - LSTM gate projections commonly produce a last dimension equal to - `4 * hidden_size`. We do not know hidden_size directly here, so this checks - that the final dimension is divisible by the expected gate count. Symbolic - dimensions can be allowed through `allow_symbolic`. - """ - gate_shape = match_utils.get_tensor_shape(gate_node) - - if not gate_shape: - return bool(spec.get("allow_missing", False)) - - hidden_dim = gate_shape[-1] - - if not isinstance(hidden_dim, int): - return bool(spec.get("allow_symbolic", True)) - - return hidden_dim % gate_count == 0 - - -def has_explicit_gate_split(gate_node: models.Node, gate_count: int, max_depth: int) -> bool: - """ - Looks for a visible split/chunk/unbind that separates combined LSTM gates. - - Shape divisibility alone is not enough to prove LSTM semantics. This walks - forward from the combined gate node and requires a split-like op whose attr - or fanout indicates the expected number of gates. If no such split is - visible, the LSTM fusion fails closed. - """ - for node in walk_children(gate_node, max_depth): - if node.target not in SPLIT_TARGETS: - continue - - if split_node_has_gate_count(node, gate_count): - return True - - return False - - -def split_node_has_gate_count(node: models.Node, gate_count: int) -> bool: - """ - Determines whether a split-like node creates the expected number of gates. - - Different export paths encode splits differently: `unbind` is inferred from - child count, `chunk` stores the chunk count as positional attr `arg_1`, and - `split_with_sizes` stores the explicit split sizes. If those attrs are not - available, child count is used as the fallback evidence. - """ - if node.target == "aten.unbind.int": - return len(node.children) == gate_count - - if node.target == "aten.chunk.default": - return match_utils.values_equal(node.attrs.get("arg_1", match_utils.MISSING), gate_count) - - split_sizes = node.attrs.get("arg_1", match_utils.MISSING) - - if isinstance(split_sizes, (list, tuple)): - return len(split_sizes) == gate_count - - return len(node.children) == gate_count - - -def walk_children(node: models.Node, max_depth: int) -> tuple[models.Node, ...]: - """ - Returns all reachable child nodes within a bounded forward traversal. - - Extra matchers use this when they need local forward evidence, such as - proving an LSTM gate tensor is split later. The traversal tracks object ids - to avoid revisiting shared children in a DAG and to remain safe if graph - rewrites produce unexpected cycles. - """ - found: list[models.Node] = [] - seen: set[int] = set() - - def visit(current: models.Node, depth: int) -> None: - """ - Recursively collects children while respecting the maximum depth. - - The outer `walk_children` owns the result list and visited set; this - nested helper only handles traversal mechanics. - """ - if depth > max_depth: - return - - for child in current.children: - if id(child) in seen: - continue - - seen.add(id(child)) - found.append(child) - visit(child, depth + 1) - - visit(node, 0) - return tuple(found) - - -def get_node_dtype(node: models.Node) -> Any: - """ - Reads a node's output dtype from tensor metadata. - - The exported JSON stores tensor metadata as dictionaries. If metadata is - missing or has some non-dict form, the function returns the shared MISSING - sentinel so callers can decide whether to fail or allow missing metadata. - """ - if not isinstance(node.tensor_output_meta, dict): - return match_utils.MISSING - - return node.tensor_output_meta.get("dtype", match_utils.MISSING) - - -def normalize_dim(dim: Any, rank: int) -> Any: - """ - Converts a possibly negative dimension index into a non-negative one. - - This mirrors Python indexing for known integer dimensions. Non-integer - values are returned unchanged so symbolic or malformed values do not crash - the matcher; they simply fail later if an exact comparison is required. - """ - if not isinstance(dim, int): - return dim - - return dim + rank if dim < 0 else dim - - -def get_repeated_subgraph(fusion: FModels.FusionGraph, name: str) -> FModels.RepeatedSubgraph | None: - """ - Fetches a named repeated subgraph declaration from a fusion graph. - - MoE fusions store expert branches as repeated subgraphs. This helper finds - the declaration by name so the MoE-specific matchers can reuse its pattern - and constraints when collecting real expert-branch bindings. - """ - for subgraph in fusion.repeated_subgraphs: - if subgraph.name == name: - return subgraph - - return None - - -def collect_repeated_subgraph_bindings(graph: models.Graph, subgraph: FModels.RepeatedSubgraph) -> tuple[dict[str, models.Node], ...]: - """ - Finds every real subgraph binding that matches a repeated pattern. - - This scans each node as a possible root for `subgraph.graph`, runs the same - graph binder used by normal fusion matching, then validates the full set of - fusion matchers against that candidate. The result is not just a count: it - returns each binding dictionary so later matchers can inspect which real - nodes were bound to expert_gate, expert_weighted, routing_weight, etc. - """ - from . import match - - matches: list[dict[str, models.Node]] = [] - - for candidate in graph.nodes: - candidate_bindings = match_utils.bind_fusion_graph(candidate, subgraph.graph, match.match_nodes) - - if candidate_bindings is None: - continue - - if not match.match_fusion_bindings(candidate, graph, subgraph.graph, candidate_bindings): - continue - - matches.append(candidate_bindings) - - return tuple(matches) - - -def relevant_moe_branch_bindings(graph: models.Graph, fusion: FModels.FusionGraph, bindings: dict[str, models.Node], spec: dict[str, Any]) -> tuple[dict[str, models.Node], ...]: - """ - Filters repeated expert branches down to ones driven by MoE routing. - - The first step collects every expert-like branch in the whole graph. Then - this keeps only branches whose declared `routing_weight` input is reachable - from a top-k node or that top-k node's probability input. This is the key - safety check that prevents us from fusing arbitrary MLP-looking branches as - MoE experts when they are not actually selected by the router. - """ - subgraph = get_repeated_subgraph(fusion, spec["repeated_subgraph"]) - - if subgraph is None: - return () - - topk_node = bindings.get(spec.get("topk_node", "topk")) - combine_node = bindings.get(spec.get("combine_node", "moe_combine")) - topk_nodes = (topk_node,) if topk_node is not None else find_reachable_topk_nodes(graph, combine_node, spec) - - if not topk_nodes: - return () - - max_depth = int(spec.get("max_depth", 16)) - branch_bindings = collect_repeated_subgraph_bindings(graph, subgraph) - relevant = [] - - for branch_binding in branch_bindings: - routing_weight = match_utils.get_first_input_by_role(subgraph.graph, branch_binding, spec["routing_weight_role"]) - - if routing_weight is None: - continue - - routing_sources = tuple(source for topk in topk_nodes for source in (topk, first_parent(topk)) if source is not None) - - if not has_any_path(routing_sources, routing_weight, max_depth): - continue - - relevant.append(branch_binding) - - return tuple(relevant) - - -def required_moe_branch_count(subgraph: FModels.RepeatedSubgraph, topk_node: models.Node | None, spec: dict[str, Any]) -> int: - """ - Determines how many routed expert branches must be proven. - - The spec may override this directly with `min_count`. Otherwise, if a top-k - node exposes its `k` attr, that is the strongest graph-local value because - it tells us how many experts each token selects. If neither exists, the - repeated subgraph's own `min_count` is used. - """ - if "min_count" in spec: - return int(spec["min_count"]) - - if topk_node is not None and spec.get("min_count_from_topk", True): - k = topk_node.attrs.get("k", match_utils.MISSING) - - if isinstance(k, int): - return k - - return subgraph.min_count - - -def get_attr(node: models.Node, *names: str) -> Any: - """ - Returns the first available attr from a node. - - Some attrs are normalized with friendly names such as `split_sizes`, while - older/generated JSON may only expose fallback names such as `arg_1`. This - lets extra matchers accept either form without duplicating logic. - """ - for name in names: - value = node.attrs.get(name, match_utils.MISSING) - - if value is not match_utils.MISSING: - return value - - return match_utils.MISSING - - -def get_known_int(value: Any) -> int | None: - """ - Converts plain integer-like values into int, otherwise returns None. - - Shape metadata can contain symbolic values, so matchers use this helper - when a semantic check is only safe for known concrete dimensions. - """ - return value if isinstance(value, int) and not isinstance(value, bool) else None - - -def get_shape_dim(node: models.Node | None, dim_index: int) -> int | None: - """ - Returns a concrete int dimension from a node shape when available. - - The helper supports negative indexing and intentionally returns None for - missing/symbolic dimensions so callers can decide whether to fail or skip. - """ - if node is None: - return None - - return get_known_int(get_node_dim(node, dim_index)) - - -def get_num_experts_from_inputs(fusion: FModels.FusionGraph, bindings: dict[str, models.Node], spec: dict[str, Any]) -> int | None: - """ - Infers the MoE expert count from weight/bias shapes. - - For LFM grouped MoE, router bias is `[num_experts]`, gate/up expert weight - is `[num_experts, 2 * intermediate, hidden]`, and down expert weight is - `[num_experts, hidden, intermediate]`. Matching these first dimensions is - a strong safety check that the grouped matmuls are really expert weights. - """ - candidates: list[int] = [] - - for role in (spec["gate_up_weight_role"], spec["down_weight_role"], spec.get("router_bias_role")): - if role is None: - continue - - input_node = match_utils.get_first_input_by_role(fusion, bindings, role) - shape = match_utils.get_tensor_shape(input_node) - - if not shape: - continue - - dim = get_known_int(shape[0]) - - if dim is not None: - candidates.append(dim) - - if not candidates: - return None - - if len(set(candidates)) != 1: - return None - - return candidates[0] - - -def split_sizes_are_valid(split_node: models.Node, gate_up_node: models.Node, down_weight: models.Node | None, spec: dict[str, Any]) -> bool: - """ - Checks the grouped gate/up projection split. - - LFM projects gate and up together, then splits the final dimension into two - equal chunks. The second chunk should match the down-projection expert - input dimension when that shape is available. - """ - split_sizes = get_attr(split_node, "split_sizes", "arg_1") - - if not isinstance(split_sizes, (list, tuple)) or len(split_sizes) != 2: - return bool(spec.get("allow_missing", False)) - - left_size = get_known_int(split_sizes[0]) - right_size = get_known_int(split_sizes[1]) - - if left_size is None or right_size is None: - return bool(spec.get("allow_symbolic", True)) - - if left_size != right_size: - return False - - gate_up_output_dim = get_shape_dim(gate_up_node, -1) - - if gate_up_output_dim is not None and gate_up_output_dim != left_size + right_size: - return False - - down_input_dim = get_shape_dim(down_weight, -1) - - if down_input_dim is not None and down_input_dim != right_size: - return False - - return True - - -def grouped_moe_rows_match_topk(hidden_node: models.Node | None, gate_up_node: models.Node, topk_node: models.Node, spec: dict[str, Any]) -> bool: - """ - Checks that routed grouped-mm rows equal token count times top-k. - - This links the routing decision to the grouped expert matmul size. Missing - or symbolic dimensions can be allowed by `allow_missing` because shape - metadata may be incomplete in some exported graphs. - """ - token_count = get_shape_dim(hidden_node, 0) - grouped_rows = get_shape_dim(gate_up_node, 0) - topk = get_known_int(get_attr(topk_node, "k", "arg_1")) - - if token_count is None or grouped_rows is None or topk is None: - return bool(spec.get("allow_missing", False)) - - return grouped_rows == token_count * topk - - -def grouped_moe_weight_permute_is_valid(node: models.Node | None) -> bool: - """ - Checks the expert weight permutation used before grouped_mm. - - The LFM export stores expert weights with expert dimension first and then - permutes the last two dims to feed grouped_mm. That appears as - `permute([0, 2, 1])`. - """ - if node is None: - return False - - return match_utils.values_equal(get_attr(node, "permutation", "arg_1"), [0, 2, 1]) - - -def find_reachable_topk_nodes(graph: models.Graph, combine_node: models.Node | None, spec: dict[str, Any]) -> tuple[models.Node, ...]: - """ - Finds top-k router nodes that can flow into the MoE combine output. - - Some MoE fusion graphs do not bind top-k as a direct synthetic node because - the exported routing path can contain indexing/scatter/gather wrappers. - This helper searches the real graph for `aten.topk.default` nodes that have - a forward path to the matched combine node within the configured depth. - """ - if combine_node is None: - return () - - max_depth = int(spec.get("max_depth", 16)) - return tuple(node for node in graph.nodes if node.target == "aten.topk.default" and has_path(node, combine_node, max_depth)) - - -def first_parent(node: models.Node) -> models.Node | None: - """ - Returns the first parent of a node, if it has one. - - For top-k, parent 0 is usually the routing probabilities tensor. The MoE - routing matcher treats both the top-k node and its first parent as valid - routing sources because exported graphs may use either values or indices - through different wrapper chains. - """ - return node.parents[0] if node.parents else None - - -def branch_output_reaches_combine(branch_binding: dict[str, models.Node], subgraph: FModels.RepeatedSubgraph, combine_node: models.Node, spec: dict[str, Any]) -> bool: - """ - Checks whether one matched expert branch contributes to the combine node. - - The branch's exposed weighted output should eventually flow into the outer - MoE combine/sum/scatter-add node. This proves that a routed expert branch is - actually part of the candidate MoE output instead of being an unrelated - branch elsewhere in the graph. - """ - weighted_node_name = spec.get("weighted_node", subgraph.graph.root) - weighted_node = branch_binding.get(weighted_node_name) - - if weighted_node is None: - return False - - return has_path(weighted_node, combine_node, int(spec.get("max_depth", 16))) - - -def has_any_path(sources: tuple[models.Node | None, ...], dest: models.Node, max_depth: int) -> bool: - """ - Returns True if any non-None source has a forward path to `dest`. - - This is a small convenience wrapper around `has_path` for cases where a - matcher has multiple plausible routing sources, such as top-k itself and - the routing-probability parent of top-k. - """ - for source in sources: - if source is not None and has_path(source, dest, max_depth): - return True - - return False - - -def has_path(source: models.Node, dest: models.Node, max_depth: int) -> bool: - """ - Performs a bounded forward reachability check in the IR DAG. - - The function starts at `source`, walks through `children`, and succeeds if - it reaches `dest` within `max_depth`. The visited set prevents repeated - work on shared DAG nodes and protects against accidental cycles. - """ - if source is dest: - return True - - seen: set[int] = set() - - def visit(node: models.Node, depth: int) -> bool: - if depth > max_depth or id(node) in seen: - return False - - seen.add(id(node)) - - for child in node.children: - if child is dest: - return True - - if visit(child, depth + 1): - return True - - return False - - return visit(source, 0) - - -def shapes_are_broadcastable(left_shape: list[Any], right_shape: list[Any]) -> bool: - """ - Implements standard trailing-dimension broadcasting for two shapes. - - It compares dimensions from right to left. Each pair is compatible if the - dims are equal or either side is 1. Extra leading dimensions on the longer - shape are allowed, matching normal tensor broadcasting rules. - """ - for left_dim, right_dim in zip(reversed(left_shape), reversed(right_shape)): - if match_utils.values_equal(left_dim, 1) or match_utils.values_equal(right_dim, 1): - continue - - if not match_utils.values_equal(left_dim, right_dim): - return False - - return True - - -def get_layer_key(node: models.Node) -> str | None: - """ - Extracts a coarse layer identifier from module stack, node name, or target. - - Cached attention checks need to know whether different matched nodes belong - to the same decoder layer. This tries the richer module stack first, then - falls back to name/target strings, and returns keys like `layers.0` when it - can identify a layer index. - """ - for value in (module_stack_text(node), node.name, node.target): - layer_key = first_layer_key(value) - - if layer_key is not None: - return layer_key - - return None - - -def module_stack_text(node: models.Node) -> str: - """ - Converts module_stack metadata into searchable text. - - Exported module stacks may be None or structured objects/lists. The matcher - only needs a best-effort string to search for layer-like path fragments, so - None becomes an empty string and everything else uses repr. - """ - return "" if node.module_stack is None else repr(node.module_stack) - - -def first_layer_key(value: str) -> str | None: - """ - Finds the first layer/block path fragment in a string. - - The helper normalizes `/` and `_` into `.` so names like - `model_layers_12_attn` can be treated similarly to - `model.layers.12.attn`. It then looks for common containers such as - `layers`, `layer`, `blocks`, or `h` followed by a numeric index. - """ - parts = value.replace("/", ".").replace("_", ".").split(".") - - for index, part in enumerate(parts[:-1]): - if part in {"layers", "layer", "h", "blocks", "block"} and parts[index + 1].isdigit(): - return f"{part}.{parts[index + 1]}" - - return None - - - -############################################ Extra Matchers!!!!! ############################################ -def match_extra_constraints(source: models.Node, graph: models.Graph, fusion: FModels.FusionGraph, bindings: dict[str, models.Node]) -> bool: - """ - Dispatches every structured FusionGraph constraint to its registered matcher. - - `fusion.constraints` is a dict keyed by constraint name. This function looks - up each name in `EXTRA_MATCHERS`, normalizes each raw spec into one or more - spec dictionaries, and requires every matcher invocation to return True. - Unknown constraint names or malformed specs fail closed. - """ - for constraint_name, raw_spec in fusion.constraints.items(): - matcher = EXTRA_MATCHERS.get(constraint_name) - - if matcher is None: - return False - - specs = normalize_constraint_specs(raw_spec) - - if specs is None: - return False - - for spec in specs: - if not matcher(source, graph, fusion, bindings, spec): - return False - - return True - - -def normalize_constraint_specs(raw_spec: Any) -> tuple[dict[str, Any], ...] | None: - """ - Normalizes one constraint value into a tuple of spec dictionaries. - - A constraint key can map to a single dict or a list/tuple of dicts. The - latter lets one matcher type be applied multiple times without duplicating - registry keys. Non-dict specs return None so the caller can fail safely. - """ - if isinstance(raw_spec, dict): - return (raw_spec,) - - if isinstance(raw_spec, (list, tuple)) and all(isinstance(spec, dict) for spec in raw_spec): - return tuple(raw_spec) - - return None - - -def match_linear_weight_layout(source: models.Node, graph: models.Graph, fusion: FModels.FusionGraph, bindings: dict[str, models.Node], spec: dict[str, Any]) -> bool: - """ - Verifies that a linear/matmul RHS is a compatible model weight layout. - - The matcher resolves the fusion's input, weight, and output nodes by role. - It then classifies the weight path as direct or transposed and checks that - dimensions line up with the matmul output. This is what prevents a random - activation matrix from being treated as a Cactus linear weight. - """ - input_node = match_utils.get_first_input_by_role(fusion, bindings, spec["input_role"]) - weight_node = match_utils.get_first_input_by_role(fusion, bindings, spec["weight_role"]) - output_node = bindings.get(spec["output_node"]) - - if input_node is None or weight_node is None or output_node is None: - return bool(spec.get("allow_missing", False)) - - layout = classify_linear_rhs(weight_node) - - if layout is None: - return False - - layout_kind, source_weight = layout - allowed_layouts = spec.get("allowed_layouts", ("direct", "transposed")) - - if layout_kind not in allowed_layouts: - return False - - return linear_shapes_match(input_node, weight_node, output_node, source_weight, layout_kind, spec) - - -def match_rope_tables_compatible(source: models.Node, graph: models.Graph, fusion: FModels.FusionGraph, bindings: dict[str, models.Node], spec: dict[str, Any]) -> bool: - """ - Verifies that RoPE cos/sin inputs look like compatible rotation tables. - - It gets the matched x/cos/sin values by role, checks cos/sin shape and dtype - compatibility, optionally proves they came from matching angle sources, and - then checks that both tables can broadcast or half-dim match the RoPE input. - """ - x_node = match_utils.get_first_input_by_role(fusion, bindings, spec["x_role"]) - cos_node = match_utils.get_first_input_by_role(fusion, bindings, spec["cos_role"]) - sin_node = match_utils.get_first_input_by_role(fusion, bindings, spec["sin_role"]) - - if cos_node is None or sin_node is None: - return bool(spec.get("allow_missing", False)) - - cos_shape = match_utils.get_tensor_shape(cos_node) - sin_shape = match_utils.get_tensor_shape(sin_node) - - if not compatible_shapes(cos_shape, sin_shape, spec): - return False - - if not compatible_dtypes(cos_node, sin_node, spec): - return False - - if not compatible_trig_sources(cos_node, sin_node, spec): - return False - - if x_node is not None and not rope_tables_match_input(x_node, cos_shape, sin_shape, spec): - return False - - return True - - -def match_lstm_gate_structure(source: models.Node, graph: models.Graph, fusion: FModels.FusionGraph, bindings: dict[str, models.Node], spec: dict[str, Any]) -> bool: - """ - Verifies that an LSTM gate node has real four-gate structure evidence. - - The matcher first checks that the combined gate tensor's final dimension is - divisible by the expected gate count. By default it then requires an - explicit downstream split/chunk/unbind into that many gates, so LSTM fusion - fails closed when the graph only vaguely resembles an LSTM. - """ - gate_node = bindings.get(spec["gate_node"]) - - if gate_node is None: - return bool(spec.get("allow_missing", False)) - - gate_count = int(spec.get("gate_count", 4)) - - if not gate_shape_matches(gate_node, gate_count, spec): - return False - - if spec.get("require_explicit_gate_split", True): - return has_explicit_gate_split(gate_node, gate_count, int(spec.get("max_depth", 4))) - - return True - - -def match_moe_expert_branch_routing(source: models.Node, graph: models.Graph, fusion: FModels.FusionGraph, bindings: dict[str, models.Node], spec: dict[str, Any]) -> bool: - """ - Verifies that enough matched expert branches are driven by router top-k. - - This locates the repeated expert-branch pattern, finds or discovers a top-k - node connected to the MoE combine output, then collects expert branch - bindings whose routing-weight input is reachable from that routing source. - The match succeeds only if the number of routed branches meets the expected - count from `topk.k`, an explicit spec override, or the subgraph minimum. - """ - subgraph = get_repeated_subgraph(fusion, spec["repeated_subgraph"]) - combine_node = bindings.get(spec.get("combine_node", "moe_combine")) - topk_node = bindings.get(spec.get("topk_node", "topk")) - - if subgraph is None: - return bool(spec.get("allow_missing", False)) - - if topk_node is None: - topk_nodes = find_reachable_topk_nodes(graph, combine_node, spec) - topk_node = topk_nodes[0] if topk_nodes else None - - if topk_node is None: - return bool(spec.get("allow_missing", False)) - - branch_bindings = relevant_moe_branch_bindings(graph, fusion, bindings, spec) - min_count = required_moe_branch_count(subgraph, topk_node, spec) - - return len(branch_bindings) >= min_count - - -def match_moe_routing_weights_combine(source: models.Node, graph: models.Graph, fusion: FModels.FusionGraph, bindings: dict[str, models.Node], spec: dict[str, Any]) -> bool: - """ - Verifies that routed expert outputs are actually combined into MoE output. - - This reuses the routed expert-branch filtering, then checks that each - branch's weighted expert output has a forward path to the matched combine - node. That proves both sides of MoE semantics: routing weights are used on - expert outputs, and those weighted outputs feed the final combine op. - """ - subgraph = get_repeated_subgraph(fusion, spec["repeated_subgraph"]) - topk_node = bindings.get(spec.get("topk_node", "topk")) - combine_node = bindings.get(spec["combine_node"]) - - if subgraph is None or combine_node is None: - return bool(spec.get("allow_missing", False)) - - if topk_node is None: - topk_nodes = find_reachable_topk_nodes(graph, combine_node, spec) - topk_node = topk_nodes[0] if topk_nodes else None - - if topk_node is None: - return bool(spec.get("allow_missing", False)) - - branch_bindings = relevant_moe_branch_bindings(graph, fusion, bindings, spec) - combined_branch_count = sum( - 1 for branch_binding in branch_bindings - if branch_output_reaches_combine(branch_binding, subgraph, combine_node, spec) - ) - - return combined_branch_count >= required_moe_branch_count(subgraph, topk_node, spec) +from __future__ import annotations + +from typing import Any + +from . import models +from .extra_matcher_cache import ( + match_cache_output_consumers, + match_cache_roll_append_structure, + match_conv_cache_initialize_structure, + match_empty_cache_initializer, + match_short_conv_decode_structure, + match_short_conv_prefill_structure, +) +from .extra_matcher_common import ExtraMatcher +from .extra_matcher_generic import ( + match_input_broadcastable_to_node, + match_input_value_kind, + match_node_attr_equals, + match_node_attrs_equal, + match_parent_tensor_dim_equal, + match_same_layer, + match_slice_halves, +) +from .extra_matcher_linear import ( + match_linear_weight_layout, + match_lstm_gate_structure, + match_rope_tables_compatible, +) +from .extra_matcher_moe import ( + match_grouped_moe_structure, + match_moe_expert_branch_routing, + match_moe_routing_weights_combine, +) +from ..Fusions import models as FModels -def match_grouped_moe_structure(source: models.Node, graph: models.Graph, fusion: FModels.FusionGraph, bindings: dict[str, models.Node], spec: dict[str, Any]) -> bool: +def match_extra_constraints(source: models.Node, graph: models.Graph, fusion: FModels.FusionGraph, bindings: dict[str, models.Node]) -> bool: """ - Verifies the grouped-mm MoE structure exported by LFM-style MoE models. + Dispatches every structured FusionGraph constraint to its registered matcher. - This matcher handles the compact representation where expert branches are - not expanded into separate per-expert matmuls. The graph edges prove the - local dataflow; this semantic check verifies top-k routing, expert-count - consistency, shared grouped-mm offsets, valid gate/up splitting, and final - top-k combine shape. + `fusion.constraints` is a dict keyed by constraint name. This function looks + up each name in `EXTRA_MATCHERS`, normalizes each raw spec into one or more + spec dictionaries, and requires every matcher invocation to return True. + Unknown constraint names or malformed specs fail closed. """ - topk_node = bindings.get(spec["topk_node"]) - topk_indices_node = bindings.get(spec["topk_indices_node"]) - gate_up_grouped_node = bindings.get(spec["gate_up_grouped_node"]) - down_grouped_node = bindings.get(spec["down_grouped_node"]) - split_node = bindings.get(spec["split_node"]) - offsets_node = bindings.get(spec["offsets_node"]) - combine_node = bindings.get(spec["combine_node"]) - hidden_node = bindings.get("moe_hidden_view") - gate_up_weight_transpose = bindings.get("moe_gate_up_weight_transpose") - down_weight_transpose = bindings.get("moe_down_weight_transpose") - down_weight = match_utils.get_first_input_by_role(fusion, bindings, spec["down_weight_role"]) - - if any(node is None for node in (topk_node, topk_indices_node, gate_up_grouped_node, down_grouped_node, split_node, offsets_node, combine_node)): - return bool(spec.get("allow_missing", False)) - - if get_known_int(get_attr(topk_node, "k", "arg_1")) is None: - return False - - if not match_utils.values_equal(get_attr(topk_indices_node, "index", "arg_1"), 1): - return False - - if not grouped_moe_weight_permute_is_valid(gate_up_weight_transpose): - return False - - if not grouped_moe_weight_permute_is_valid(down_weight_transpose): - return False - - gate_up_offsets = match_utils.get_parent(gate_up_grouped_node, 2) - down_offsets = match_utils.get_parent(down_grouped_node, 2) - - if gate_up_offsets is not offsets_node or down_offsets is not offsets_node: - return False - - if not match_utils.values_equal(get_attr(offsets_node, "dim", "arg_1"), 0): - return False - - if not split_sizes_are_valid(split_node, gate_up_grouped_node, down_weight, spec): - return False - - num_experts = get_num_experts_from_inputs(fusion, bindings, spec) - histc_node = bindings.get("moe_histc") - histc_bins = get_known_int(get_attr(histc_node, "bins", "arg_1")) if histc_node is not None else None - - if num_experts is not None and histc_bins is not None and num_experts != histc_bins: - return False - - if not grouped_moe_rows_match_topk(hidden_node, gate_up_grouped_node, topk_node, spec): - return False - - if not match_utils.values_equal(get_attr(combine_node, "dim", "arg_1"), 1): - return False - - return True - - -def match_cache_roll_append_structure(source: models.Node, graph: models.Graph, fusion: FModels.FusionGraph, bindings: dict[str, models.Node], spec: dict[str, Any]) -> bool: - cache_node = match_utils.get_first_input_by_role(fusion, bindings, spec["cache_role"]) - new_data_node = match_utils.get_first_input_by_role(fusion, bindings, spec["new_data_role"]) - - if cache_node is None or new_data_node is None: - return False - - cache_shape = match_utils.get_tensor_shape(cache_node) - new_data_shape = match_utils.get_tensor_shape(new_data_node) - output_shape = match_utils.get_tensor_shape(source) - - if not cache_shape or not new_data_shape or not output_shape: - return False - - if not match_utils.values_equal(cache_shape, output_shape): - return False - - if len(cache_shape) != len(new_data_shape): - return False - - window_size = get_known_int(cache_shape[-1]) - - if window_size is None or window_size <= 0: - return False - - if not cache_append_shapes_match(cache_shape, new_data_shape): - return False - - arange_node = bindings.get(spec["arange_node"]) - add_node = bindings.get(spec["add_node"]) - mod_node = bindings.get(spec["mod_node"]) - - if arange_node is None or add_node is None or mod_node is None: - return False - - if not match_utils.values_equal(get_attr(arange_node, "start", "arg_0"), 0): - return False - - if not match_utils.values_equal(get_attr(arange_node, "end", "arg_1"), window_size): - return False - - if not match_utils.values_equal(get_attr(add_node, "other", "value", "arg_1"), 1): - return False - - if not match_utils.values_equal(get_attr(mod_node, "other", "value", "arg_1"), window_size): - return False - - return ( - slice_node_matches(bindings.get(spec["scatter_slice_dim0_node"]), axis=0, start=0) - and slice_node_matches(bindings.get(spec["scatter_slice_dim1_node"]), axis=1, start=0) - and slice_node_matches(bindings.get(spec["value_slice_dim0_node"]), axis=0, start=0) - and slice_node_matches(bindings.get(spec["value_slice_dim1_node"]), axis=1, start=0) - and slice_node_matches(bindings.get(spec["value_slice_last_node"]), axis=2, start=-1) - and scatter_node_matches(bindings.get(spec["scatter_dim2_node"]), axis=2, start=-1) - and scatter_node_matches(bindings.get(spec["scatter_dim1_node"]), axis=1, start=0) - and scatter_node_matches(bindings.get(spec["scatter_dim0_node"]), axis=0, start=0) - ) - - -def match_empty_cache_initializer(source: models.Node, graph: models.Graph, fusion: FModels.FusionGraph, bindings: dict[str, models.Node], spec: dict[str, Any]) -> bool: - empty_node = match_utils.get_first_input_by_role(fusion, bindings, spec["empty_role"]) - new_value_node = match_utils.get_first_input_by_role(fusion, bindings, spec["new_value_role"]) - output_node = bindings.get(spec["output_node"]) - - if empty_node is None or new_value_node is None or output_node is None: - return False - - empty_shape = match_utils.get_tensor_shape(empty_node) - new_shape = match_utils.get_tensor_shape(new_value_node) - output_shape = match_utils.get_tensor_shape(output_node) - - if not empty_shape or not new_shape or not output_shape: - return False - - if element_count(empty_shape) != 0: - return False - - return match_utils.values_equal(new_shape, output_shape) - - -def match_conv_cache_initialize_structure(source: models.Node, graph: models.Graph, fusion: FModels.FusionGraph, bindings: dict[str, models.Node], spec: dict[str, Any]) -> bool: - rows_node = match_utils.get_first_input_by_role(fusion, bindings, spec["rows_role"]) - empty_node = match_utils.get_first_input_by_role(fusion, bindings, spec["empty_role"]) - output_node = bindings.get(spec["output_node"]) - - if rows_node is None or empty_node is None or output_node is None: - return False - - rows_shape = match_utils.get_tensor_shape(rows_node) - empty_shape = match_utils.get_tensor_shape(empty_node) - output_shape = match_utils.get_tensor_shape(output_node) - - if len(rows_shape) != 3 or len(empty_shape) != 3 or len(output_shape) != 3: - return False - - if not match_utils.values_equal(rows_shape, output_shape): - return False - - if not match_utils.values_equal(empty_shape, output_shape): - return False - - if empty_node.target != "aten.full_like.default": - return False - - if not match_utils.values_equal(get_attr(empty_node, "fill_value", "arg_1"), 0): - return False - - return get_known_int(rows_shape[1]) is not None and get_known_int(rows_shape[2]) is not None - - -def match_short_conv_prefill_structure(source: models.Node, graph: models.Graph, fusion: FModels.FusionGraph, bindings: dict[str, models.Node], spec: dict[str, Any]) -> bool: - conv_node = bindings.get(spec["conv_node"]) - slice_node = bindings.get(spec["slice_node"]) - x_node = match_utils.get_first_input_by_role(fusion, bindings, spec["x_role"]) - weight_node = match_utils.get_first_input_by_role(fusion, bindings, spec["weight_role"]) - - if conv_node is None or slice_node is None or x_node is None or weight_node is None: - return False - - x_shape = match_utils.get_tensor_shape(x_node) - weight_shape = match_utils.get_tensor_shape(weight_node) - conv_shape = match_utils.get_tensor_shape(conv_node) - output_shape = match_utils.get_tensor_shape(slice_node) - - if len(x_shape) != 3 or len(weight_shape) != 3 or len(conv_shape) != 3 or len(output_shape) != 3: - return False - - hidden_dim = get_known_int(x_shape[1]) - kernel_size = get_known_int(weight_shape[2]) - input_length = get_known_int(x_shape[2]) - - if hidden_dim is None or kernel_size is None or input_length is None: - return False - - if not match_utils.values_equal(weight_shape[0], hidden_dim): - return False - - if not match_utils.values_equal(weight_shape[1], 1): - return False - - if first_int_attr(conv_node, "stride", 1) != 1: - return False - - if first_int_attr(conv_node, "dilation", 1) != 1: - return False - - if first_int_attr(conv_node, "padding", 0) != kernel_size - 1: - return False - - if first_int_attr(conv_node, "groups", 1) != hidden_dim: - return False - - slice_axis = first_int_attr(slice_node, "axis", first_int_attr(slice_node, "dim", -1)) - if normalize_dim(slice_axis, 3) != 2: - return False - - if first_int_attr(slice_node, "start", 0) != 0: - return False - - return match_utils.values_equal(output_shape, x_shape) and get_shape_dim(conv_node, 2) == input_length + kernel_size - 1 - - -def match_short_conv_decode_structure(source: models.Node, graph: models.Graph, fusion: FModels.FusionGraph, bindings: dict[str, models.Node], spec: dict[str, Any]) -> bool: - cache_window_node = match_utils.get_first_input_by_role(fusion, bindings, spec["cache_window_role"]) - multiply_node = bindings.get(spec["multiply_node"]) - sum_node = bindings.get(spec["sum_node"]) - weight_node = match_utils.get_first_input_by_role(fusion, bindings, spec["weight_role"]) - - if cache_window_node is None or multiply_node is None or sum_node is None or weight_node is None: - return False - - if cache_window_node.target != "cactus.conv_cache_append" and not ( - cache_window_node.cache is not None and cache_window_node.cache.kind == FModels.CacheKind.CONV - ): - return False - - cache_window_shape = match_utils.get_tensor_shape(cache_window_node) - weight_shape = match_utils.get_tensor_shape(weight_node) - multiply_shape = match_utils.get_tensor_shape(multiply_node) - output_shape = match_utils.get_tensor_shape(sum_node) - - if len(cache_window_shape) != 3 or len(multiply_shape) != 3 or len(output_shape) != 2: - return False - - if not match_utils.values_equal(cache_window_shape, multiply_shape): - return False - - batch_size = get_known_int(cache_window_shape[0]) - hidden_dim = get_known_int(cache_window_shape[1]) - window_size = get_known_int(cache_window_shape[2]) - - if batch_size is None or hidden_dim is None or window_size is None: - return False - - if not match_utils.values_equal(output_shape, [batch_size, hidden_dim]): - return False - - if not short_conv_decode_weight_shape_matches(weight_shape, hidden_dim, window_size): - return False - - sum_axis = first_int_attr(sum_node, "axis", first_int_attr(sum_node, "dim", -1)) - if normalize_dim(sum_axis, 3) != 2: - return False - - keepdim = get_attr(sum_node, "keepdim") - return keepdim is match_utils.MISSING or keepdim is False - - -def match_cache_output_consumers(source: models.Node, graph: models.Graph, fusion: FModels.FusionGraph, bindings: dict[str, models.Node], spec: dict[str, Any]) -> bool: - node = bindings.get(spec["node"]) - - if node is None: - return False - - allow_output_child = bool(spec.get("allow_output_child", True)) - - for child in node.children: - if child.target in match_utils.METADATA_ONLY_TARGETS: - continue - - if allow_output_child and child.is_output: - continue - - return False - - return True - - -def element_count(shape: list[Any]) -> int | None: - count = 1 - - for dim in shape: - known_dim = get_known_int(dim) - - if known_dim is None: - return None - - count *= known_dim - - return count - - -def first_int_attr(node: models.Node, name: str, default: int) -> int: - value = get_attr(node, name) - - if value is match_utils.MISSING: - return default - - if isinstance(value, (list, tuple)): - if not value: - return default - - return int(value[0]) - - return int(value) - + for constraint_name, raw_spec in fusion.constraints.items(): + matcher = EXTRA_MATCHERS.get(constraint_name) -def cache_append_shapes_match(cache_shape: list[Any], new_data_shape: list[Any]) -> bool: - for cache_dim, new_dim in zip(cache_shape[:-1], new_data_shape[:-1]): - if not match_utils.values_equal(cache_dim, new_dim): + if matcher is None: return False - return match_utils.values_equal(new_data_shape[-1], 1) - - -def short_conv_decode_weight_shape_matches(weight_shape: list[Any], hidden_dim: int, window_size: int) -> bool: - if len(weight_shape) == 2: - return match_utils.values_equal(weight_shape[0], hidden_dim) and match_utils.values_equal(weight_shape[1], window_size) - - if len(weight_shape) != 3: - return False - - return ( - match_utils.values_equal(weight_shape, [1, hidden_dim, window_size]) - or match_utils.values_equal(weight_shape, [hidden_dim, 1, window_size]) - ) - - -def slice_node_matches(node: models.Node | None, *, axis: int, start: int) -> bool: - if node is None: - return False - - return match_utils.values_equal(get_attr(node, "axis", "dim", "arg_1"), axis) and match_utils.values_equal(get_attr(node, "start", "arg_2"), start) - - -def scatter_node_matches(node: models.Node | None, *, axis: int, start: int) -> bool: - if node is None: - return False - - return match_utils.values_equal(get_attr(node, "dim", "axis", "arg_2"), axis) and match_utils.values_equal(get_attr(node, "start", "arg_3"), start) - - -def match_node_attr_equals(source: models.Node, graph: models.Graph, fusion: FModels.FusionGraph, bindings: dict[str, models.Node], spec: dict[str, Any]) -> bool: - """ - Checks that one matched node has a specific normalized attr value. - - The spec names a synthetic node, an attr key, and the expected value. The - function resolves the real node from bindings, reads its normalized attrs, - and compares with `values_equal`. Missing attrs fail unless - `allow_missing=True` is present in the spec. - """ - node = bindings.get(spec["node"]) - - if node is None: - return False - - actual = node.attrs.get(spec["attr"], match_utils.MISSING) - - if actual is match_utils.MISSING: - return bool(spec.get("allow_missing", False)) - - return match_utils.values_equal(actual, spec["value"]) - - -def match_node_attrs_equal(source: models.Node, graph: models.Graph, fusion: FModels.FusionGraph, bindings: dict[str, models.Node], spec: dict[str, Any]) -> bool: - """ - Checks that attrs on two matched nodes are equal to each other. - - This is useful for pattern checks like layernorm, where mean and variance - should reduce over the same dimension. The spec identifies left/right - synthetic nodes and attr names; both attrs must exist and compare equal - unless missing attrs are explicitly allowed. - """ - left_node = bindings.get(spec["left_node"]) - right_node = bindings.get(spec["right_node"]) - - if left_node is None or right_node is None: - return False - - left_value = left_node.attrs.get(spec["left_attr"], match_utils.MISSING) - right_value = right_node.attrs.get(spec["right_attr"], match_utils.MISSING) - - if left_value is match_utils.MISSING or right_value is match_utils.MISSING: - return bool(spec.get("allow_missing", False)) - - return match_utils.values_equal(left_value, right_value) - - -def match_input_value_kind(source: models.Node, graph: models.Graph, fusion: FModels.FusionGraph, bindings: dict[str, models.Node], spec: dict[str, Any]) -> bool: - """ - Checks the semantic kind of an external input to a fusion. - - The fusion graph declares inputs by role, such as `bias` or `weight`. This - matcher resolves the first real node for that role and verifies its - `value_kind` is in the allowed set, preventing activations from being - mistaken for parameters/buffers. - """ - input_node = match_utils.get_first_input_by_role(fusion, bindings, spec["role"]) - - if input_node is None: - return bool(spec.get("allow_missing", False)) - - return input_node.value_kind in spec["allowed_value_kinds"] - - -def match_parent_tensor_dim_equal(source: models.Node, graph: models.Graph, fusion: FModels.FusionGraph, bindings: dict[str, models.Node], spec: dict[str, Any]) -> bool: - """ - Compares one dimension from two parent tensors of matched nodes. - - The spec names the left/right synthetic nodes, which parent index to inspect - on each real node, and which shape dimension to compare. It is used for - checks such as attention query/key head-dim compatibility. Missing parents - or shapes follow the spec's `allow_missing` behavior. - """ - left_parent = get_bound_parent(bindings, spec["left_node"], spec["left_parent_index"]) - right_parent = get_bound_parent(bindings, spec["right_node"], spec["right_parent_index"]) - - if left_parent is None or right_parent is None: - return bool(spec.get("allow_missing", False)) - - left_dim = get_node_dim(left_parent, spec["left_dim"]) - right_dim = get_node_dim(right_parent, spec["right_dim"]) - - if left_dim is match_utils.MISSING or right_dim is match_utils.MISSING: - return bool(spec.get("allow_missing", False)) - - return match_utils.values_equal(left_dim, right_dim) - - -def match_input_broadcastable_to_node(source: models.Node, graph: models.Graph, fusion: FModels.FusionGraph, bindings: dict[str, models.Node], spec: dict[str, Any]) -> bool: - """ - Checks that a declared fusion input can broadcast to a target node shape. - - This is mainly for mask-like inputs. It resolves an input by role, resolves - the target synthetic node, reads both output shapes, and applies standard - broadcasting. Missing nodes or shapes are controlled by `allow_missing`. - """ - input_node = match_utils.get_first_input_by_role(fusion, bindings, spec["role"]) - target_node = bindings.get(spec["node"]) - - if input_node is None or target_node is None: - return bool(spec.get("allow_missing", False)) - - input_shape = match_utils.get_tensor_shape(input_node) - target_shape = match_utils.get_tensor_shape(target_node) - - if not input_shape or not target_shape: - return bool(spec.get("allow_missing", False)) - - return shapes_are_broadcastable(input_shape, target_shape) - - -def match_same_layer(source: models.Node, graph: models.Graph, fusion: FModels.FusionGraph, bindings: dict[str, models.Node], spec: dict[str, Any]) -> bool: - """ - Checks that multiple matched nodes appear to belong to the same layer. - - The spec provides synthetic node names. For each bound real node, we derive - a layer key from module stack/name/target text. The check passes only when - every derived key is identical, which helps avoid mixing cache/key/value - nodes from different decoder layers. - """ - layer_keys = [] - - for node_name in spec["nodes"]: - node = bindings.get(node_name) + specs = normalize_constraint_specs(raw_spec) - if node is None: + if specs is None: return False - layer_key = get_layer_key(node) - - if layer_key is None: - return bool(spec.get("allow_missing", False)) - - layer_keys.append(layer_key) + for spec in specs: + if not matcher(source, graph, fusion, bindings, spec): + return False - return len(set(layer_keys)) == 1 + return True -def match_slice_halves(source: models.Node, graph: models.Graph, fusion: FModels.FusionGraph, bindings: dict[str, models.Node], spec: dict[str, Any]) -> bool: +def normalize_constraint_specs(raw_spec: Any) -> tuple[dict[str, Any], ...] | None: """ - Checks that two slice nodes split the same tensor into adjacent equal halves. + Normalizes one constraint value into a tuple of spec dictionaries. - This protects decomposed GLU fusion. GLU is valid when the exported graph - computes `left * sigmoid(right)` where `left` and `right` are the first and - second halves of one source tensor along the same feature axis. + A constraint key can map to a single dict or a list/tuple of dicts. The + latter lets one matcher type be applied multiple times without duplicating + registry keys. Non-dict specs return None so the caller can fail safely. """ - left = bindings.get(spec["left_node"]) - right = bindings.get(spec["right_node"]) - - if left is None or right is None: - return False - - left_source = match_utils.get_parent(left, 0) - right_source = match_utils.get_parent(right, 0) - - if left_source is None or left_source is not right_source: - return False - - source_shape = match_utils.get_tensor_shape(left_source) - left_shape = match_utils.get_tensor_shape(left) - right_shape = match_utils.get_tensor_shape(right) - - if not source_shape or not left_shape or not right_shape: - return bool(spec.get("allow_missing", False)) - - left_axis = normalized_slice_axis(left, len(source_shape)) - right_axis = normalized_slice_axis(right, len(source_shape)) - - if left_axis is match_utils.MISSING or right_axis is match_utils.MISSING or left_axis != right_axis: - return False - - if not slice_step_is_one(left) or not slice_step_is_one(right): - return False - - left_size = match_utils.get_dim(left_shape, left_axis) - right_size = match_utils.get_dim(right_shape, right_axis) - source_size = match_utils.get_dim(source_shape, left_axis) - - if left_size is None or right_size is None: - return bool(spec.get("allow_missing", False)) - - if not match_utils.values_equal(left_size, right_size): - return False - - if not attr_matches(left, "start", 0, default=0): - return False - - if not attr_matches(right, "start", left_size): - return False - - if not optional_attr_matches(left, "end", left_size): - return False - - if not optional_attr_matches(right, "end", add_dims(left_size, right_size)): - return False - - if source_size is not None and not match_utils.values_equal(source_size, add_dims(left_size, right_size)): - return False - - return True - - -def normalized_slice_axis(node: models.Node, rank: int) -> Any: - axis = node.attrs.get("axis", node.attrs.get("dim", match_utils.MISSING)) - - if axis is match_utils.MISSING: - return match_utils.MISSING - - return normalize_dim(axis, rank) - - -def slice_step_is_one(node: models.Node) -> bool: - step = node.attrs.get("step", 1) - return step is None or match_utils.values_equal(step, 1) - - -def attr_matches(node: models.Node, attr: str, expected: Any, *, default: Any = match_utils.MISSING) -> bool: - actual = node.attrs.get(attr, default) - - if actual is match_utils.MISSING: - return False - - return match_utils.values_equal(actual, expected) - - -def optional_attr_matches(node: models.Node, attr: str, expected: Any) -> bool: - actual = node.attrs.get(attr, match_utils.MISSING) - - if actual is match_utils.MISSING or actual is None: - return True - - return match_utils.values_equal(actual, expected) - + if isinstance(raw_spec, dict): + return (raw_spec,) -def add_dims(left: Any, right: Any) -> Any: - if isinstance(left, int) and isinstance(right, int): - return left + right + if isinstance(raw_spec, (list, tuple)) and all(isinstance(spec, dict) for spec in raw_spec): + return tuple(raw_spec) - return left + return None EXTRA_MATCHERS: dict[str, ExtraMatcher] = { @@ -1661,3 +99,4 @@ def add_dims(left: Any, right: Any) -> Any: "same_layer": match_same_layer, "slice_halves": match_slice_halves, } + diff --git a/python/cactus/transpiler/IR/models.py b/python/cactus/transpiler/IR/models.py index 6a3e22923..633612f57 100644 --- a/python/cactus/transpiler/IR/models.py +++ b/python/cactus/transpiler/IR/models.py @@ -72,6 +72,7 @@ class Graph: input_specs: tuple[CModels.GraphSpecRecord, ...] = field(default_factory=tuple) output_specs: tuple[CModels.GraphSpecRecord, ...] = field(default_factory=tuple) cache_annotations: tuple[CacheAnnotation, ...] = field(default_factory=tuple) + metadata: dict[str, str] = field(default_factory=dict) fusions: list["FusionResult"] = field(default_factory=list) @classmethod @@ -272,6 +273,7 @@ def generate_graph(cls: type[Graph], layer_map: CModels.LayerMap) -> Graph: input_specs=tuple(layer_map.input_specs), output_specs=tuple(layer_map.output_specs), cache_annotations=cache_annotations_from_nodes(nodes), + metadata={}, ) @@ -568,6 +570,7 @@ def rebuild_graph(nodes: tuple[Node, ...], original_graph: Graph, fusion_results input_specs=original_graph.input_specs, output_specs=original_graph.output_specs, cache_annotations=cache_annotations_from_nodes(nodes), + metadata=dict(original_graph.metadata), fusions=list(fusion_results), ) diff --git a/python/cactus/transpiler/IR/simplify_ir.py b/python/cactus/transpiler/IR/simplify_ir.py index 34c82aa3a..e59a0fc4d 100644 --- a/python/cactus/transpiler/IR/simplify_ir.py +++ b/python/cactus/transpiler/IR/simplify_ir.py @@ -7,8 +7,6 @@ from ..Fusions import models as FModels -#TODO: Clean up code - def simplify( layer_map: CModels.LayerMap, *, diff --git a/python/cactus/transpiler/ModelProfiles/profiles.py b/python/cactus/transpiler/ModelProfiles/profiles.py index aea2fd50e..5a4585430 100644 --- a/python/cactus/transpiler/ModelProfiles/profiles.py +++ b/python/cactus/transpiler/ModelProfiles/profiles.py @@ -27,6 +27,8 @@ "embedding", "vision", "audio", + "attention", + "normalization", "gemma4_token_merge", "gemma4_rmsnorm", "gemma4_attention", @@ -49,6 +51,10 @@ fusion_fields=( "generic", "audio", + "attention", + "normalization", + "mlp", + "cache", "whisper_audio_encoder", "whisper_attention", "whisper_mlp", @@ -70,6 +76,9 @@ fusion_fields=( "generic", "audio", + "attention", + "normalization", + "mlp", "parakeet_fastconformer", "parakeet_tdt", "conv", @@ -92,6 +101,10 @@ "generic", "embedding", "vision", + "attention", + "normalization", + "mlp", + "cache", "lfm_vlm_token_merge", "lfm_attention", "lfm_mlp", @@ -113,6 +126,10 @@ fusion_fields=( "generic", "embedding", + "attention", + "normalization", + "mlp", + "cache", "qwen2_5_rmsnorm", "qwen2_5_rope", "qwen2_5_attention", @@ -135,6 +152,10 @@ fusion_fields=( "generic", "embedding", + "attention", + "normalization", + "mlp", + "cache", "lfm_attention", "lfm_mlp", "lfm_moe", diff --git a/python/cactus/transpiler/main.py b/python/cactus/transpiler/main.py index 887587bda..490201d3d 100644 --- a/python/cactus/transpiler/main.py +++ b/python/cactus/transpiler/main.py @@ -1,6 +1,5 @@ import sys import os -import shutil from pathlib import Path sys.path.append(str(Path(__file__).resolve().parents[1])) @@ -8,30 +7,10 @@ from transpiler.Converter import constants from transpiler.Converter.convert import convert from transpiler.Converter.models import LayerMap -from transpiler.Generator.generate import generate_bundle +from transpiler.Generator.generate import generate_bundle, materialize_runtime_bundle_files from transpiler.ModelProfiles.profiles import GEMMA4_E2B_PROFILE -def materialize_runtime_bundle_files(source_dir: Path, bundle_dir: Path) -> None: - for source in source_dir.iterdir(): - if not source.is_file(): - continue - - target = bundle_dir / source.name - if target.exists() or target.is_symlink(): - try: - if target.samefile(source): - continue - except OSError: - pass - target.unlink() - - try: - os.symlink(source, target) - except OSError: - shutil.copy2(source, target) - - def main() -> None: model_id = "google/gemma-4-E2B" input_modalities = ("text", "vision", "audio") From fb6176ccfc88d260dfa2594baf6f78e9eedf93ee Mon Sep 17 00:00:00 2001 From: Parkirat Date: Mon, 3 Aug 2026 00:08:35 -0400 Subject: [PATCH 72/84] Performance upgrades --- cactus-engine/src/complete.cpp | 42 +- cactus-engine/src/engine.h | 16 +- cactus-engine/src/engine_image.cpp | 10 +- cactus-engine/src/init.cpp | 9 +- cactus-engine/src/model.cpp | 128 ++++- cactus-engine/src/telemetry_impl.cpp | 5 + cactus-engine/src/tokenizer.cpp | 140 ++++- cactus-engine/src/utils.h | 10 +- cactus-engine/tests/transcribe.cpp | 2 +- cactus-graph/src/builder.cpp | 21 +- cactus-graph/src/core.cpp | 1 + cactus-graph/src/io.cpp | 32 +- cactus-graph/src/ops_nn.cpp | 74 ++- cactus-kernels/src/matmul.cpp | 16 +- cactus-kernels/src/threading.h | 8 + python/cactus/cli/__init__.py | 10 +- python/cactus/cli/convert.py | 20 +- python/cactus/cli/model.py | 68 ++- python/cactus/cli/transpiler.py | 248 +++++++++ .../convert/cactus_adapters/tokenizer.py | 4 + python/cactus/convert/cli.py | 4 +- python/cactus/convert/export/files.py | 45 ++ .../cactus/convert/model_adapters/adapters.py | 9 + python/cactus/convert/tests/test_policy.py | 29 + .../transpiler/Converter/cache_utils.py | 4 + .../cactus/transpiler/Converter/constants.py | 3 +- python/cactus/transpiler/Converter/convert.py | 5 +- python/cactus/transpiler/Converter/models.py | 2 + .../cactus/transpiler/Converter/overrides.py | 118 ++++ python/cactus/transpiler/Fusions/edges.py | 3 + .../transpiler/Fusions/fusion_attention.py | 16 +- .../transpiler/Fusions/fusion_direct.py | 1 + .../transpiler/Fusions/fusion_neural.py | 77 +++ .../transpiler/Fusions/fusion_special.py | 18 +- python/cactus/transpiler/Fusions/fusions.py | 51 +- python/cactus/transpiler/Fusions/nodes.py | 15 +- .../Generator/component_split_boundaries.py | 171 +++++- .../Generator/component_split_builder.py | 360 +++++++++++- .../Generator/component_split_plans.py | 145 ++++- .../Generator/component_split_types.py | 17 + .../cactus/transpiler/Generator/constants.py | 19 + .../cactus/transpiler/Generator/generate.py | 102 +++- .../Generator/lowering_basic_ops.py | 107 +++- .../transpiler/Generator/lowering_cache.py | 59 +- .../transpiler/Generator/lowering_core.py | 6 +- .../transpiler/Generator/lowering_nn_ops.py | 2 + .../Generator/lowering_special_ops.py | 197 ++++++- .../transpiler/Generator/lowering_utils.py | 82 +++ .../transpiler/Generator/lowering_weights.py | 38 +- python/cactus/transpiler/Generator/models.py | 5 + python/cactus/transpiler/IR/models.py | 40 +- .../cactus/transpiler/IR/special_fusions.py | 513 +++++++++++++++++- .../transpiler/ModelProfiles/profiles.py | 7 +- 53 files changed, 2982 insertions(+), 152 deletions(-) create mode 100644 python/cactus/cli/transpiler.py diff --git a/cactus-engine/src/complete.cpp b/cactus-engine/src/complete.cpp index 7baa740b3..1ece66cf4 100644 --- a/cactus-engine/src/complete.cpp +++ b/cactus-engine/src/complete.cpp @@ -311,6 +311,13 @@ std::vector> build_stop_sequences( stop_token_sequences.push_back(tokenizer->encode(stop_seq)); } + if (!tokenizer->has_chat_template()) { + stop_token_sequences.push_back(tokenizer->encode("\nUser:")); + stop_token_sequences.push_back(tokenizer->encode("\nAssistant:")); + stop_token_sequences.push_back(tokenizer->encode("\nUser")); + stop_token_sequences.push_back(tokenizer->encode("\nAssistant")); + } + if (model_type == Config::ModelType::GEMMA4) { stop_token_sequences.push_back(tokenizer->encode("")); if (has_tools) { @@ -374,6 +381,7 @@ struct PreparedPrompt { std::vector tokens; size_t context_token_count = 0; std::vector> images; + bool force_session_reset = false; std::vector> audio_features; size_t audio_num_frames = 0; @@ -527,8 +535,6 @@ PreparedPrompt prepare_prompt( auto samples_16k = resample_to_16k_fp32(waveform_fp32, 16000); if (!samples_16k.empty()) { auto audio_prep = cactus::audio::preprocess_audio_for_gemma4(samples_16k, handle->model->get_config()); - size_t graph_soft_tokens = handle->model->audio_soft_token_count_for_frames(audio_prep.num_frames); - if (graph_soft_tokens > 0) audio_prep.num_soft_tokens = std::min(audio_prep.num_soft_tokens, graph_soft_tokens); prompt.audio_features.push_back(std::move(audio_prep.features)); size_t u = user_indices.size() - 1; prompt.messages[user_indices[u]].audio_soft_token_count = audio_prep.num_soft_tokens; @@ -543,8 +549,6 @@ PreparedPrompt prepare_prompt( auto samples_16k = resample_to_16k_fp32(wav.samples, wav.sample_rate); if (samples_16k.empty()) continue; auto audio_prep = cactus::audio::preprocess_audio_for_gemma4(samples_16k, handle->model->get_config()); - size_t graph_soft_tokens = handle->model->audio_soft_token_count_for_frames(audio_prep.num_frames); - if (graph_soft_tokens > 0) audio_prep.num_soft_tokens = std::min(audio_prep.num_soft_tokens, graph_soft_tokens); prompt.audio_features.push_back(std::move(audio_prep.features)); prompt.messages[user_indices[u]].audio_soft_token_count = audio_prep.num_soft_tokens; counts[u] = audio_prep.num_soft_tokens; @@ -572,6 +576,23 @@ PreparedPrompt prepare_prompt( setup_tool_constraints(handle, prompt.tools, prompt.options.force_tools, prompt.options.temperature); } + if (tokenizer->uses_default_chat_style(prompt.messages, formatted_tools, prompt.options.enable_thinking_if_supported)) { + auto last_user = std::find_if(prompt.messages.rbegin(), prompt.messages.rend(), [](const ChatMessage& msg) { + return msg.role == "user"; + }); + if (last_user != prompt.messages.rend()) { + std::vector latest_turn; + for (const auto& msg : prompt.messages) { + if (msg.role == "system" || msg.role == "developer") { + latest_turn.push_back(msg); + } + } + latest_turn.push_back(*last_user); + prompt.messages = std::move(latest_turn); + prompt.force_session_reset = true; + } + } + prompt.rendered = tokenizer->format_chat_prompt( prompt.messages, add_generation_prompt, @@ -766,6 +787,15 @@ int cactus_complete( handle_error_response(err, response_buffer, buffer_size); return -1; } + const auto runtime_model_type = handle->model->get_config().model_type; + if (runtime_model_type == Config::ModelType::WHISPER || + runtime_model_type == Config::ModelType::PARAKEET || + runtime_model_type == Config::ModelType::PARAKEET_TDT) { + const std::string err = "Model is speech-to-text; use `cactus transcribe` instead of `cactus run`"; + CACTUS_LOG_ERROR("complete", err); + handle_error_response(err, response_buffer, buffer_size); + return -1; + } auto* tokenizer = handle->model->get_tokenizer(); auto prompt = prepare_prompt(handle, messages_json, options_json, tools_json, true, true, pcm_buffer, pcm_buffer_size); @@ -892,6 +922,10 @@ int cactus_complete( uint32_t next_token; size_t prompt_tokens; + if (prompt.force_session_reset && !handle->processed_tokens.empty()) { + reset_cache(handle); + } + if (has_audio && !handle->processed_tokens.empty()) { auto& cache = handle->processed_tokens; size_t common = 0; diff --git a/cactus-engine/src/engine.h b/cactus-engine/src/engine.h index 53ec2b0b9..7520818c7 100644 --- a/cactus-engine/src/engine.h +++ b/cactus-engine/src/engine.h @@ -115,6 +115,10 @@ struct Config { uint32_t tile_size = 512; float max_pixels_tolerance = 2.0f; bool do_image_splitting = true; + bool do_resize = true; + bool do_rescale = true; + bool do_normalize = true; + bool do_convert_rgb = true; bool encoder_act_gelu = false; bool decoder_act_gelu = false; uint32_t num_encoder_layers = 0; @@ -331,6 +335,7 @@ class Tokenizer { virtual std::vector apply_chat_template(const std::vector& messages, bool add_generation_prompt = true) const; virtual std::string format_chat_prompt(const std::vector& messages, bool add_generation_prompt = true, const std::string& tools_json = "", bool enable_thinking_if_supported = false) const; + virtual bool uses_default_chat_style(const std::vector& messages, const std::string& tools_json = "", bool enable_thinking_if_supported = false) const; virtual uint32_t get_vocab_size() const = 0; virtual uint32_t get_unk_token() const = 0; @@ -355,7 +360,7 @@ class Tokenizer { void set_lfm2_vision_config(const Config& cfg) { lfm2_vision_config_ = cfg; has_lfm2_vision_config_ = true; } protected: - enum class ModelType { UNKNOWN, GEMMA4, GEMMA, QWEN, LFM2, NEEDLE }; + enum class ModelType { UNKNOWN, GEMMA4, GEMMA, QWEN, LFM2, NEEDLE, WHISPER, PARAKEET_TDT }; ModelType model_type_ = ModelType::UNKNOWN; enum class ModelVariant { DEFAULT, VLM, EXTRACT, RAG}; ModelVariant model_variant_ = ModelVariant::DEFAULT; @@ -378,6 +383,7 @@ class Tokenizer { void detect_model_type(const std::string& config_path); void load_chat_template(const std::string& template_file); + std::string format_default_chat_style(const std::vector& messages, bool add_generation_prompt) const; std::string format_gemma4_style(const std::vector& messages, bool add_generation_prompt, const std::string& tools_json, bool enable_thinking_if_supported = false) const; std::string format_qwen_style(const std::vector& messages, bool add_generation_prompt, const std::string& tools_json, bool enable_thinking_if_supported = false) const; std::string format_lfm2_style(const std::vector& messages, bool add_generation_prompt, const std::string& tools_json, bool enable_thinking_if_supported = false) const; @@ -811,6 +817,7 @@ class Model { void write_media_embeds_row(Component& comp, int embeds_idx, const uint8_t* feature_row, size_t feature_row_bytes, Precision feature_precision); void reset_encoder_cross_kv_route_state(); + void release_encoder_cross_kv_prepare_graphs(); bool finish_encoder_cross_kv_prepare(); bool finish_encoder_cross_kv_prepare_after_source(); bool prepare_encoder_cross_kv_from_text(const std::vector& tokens); @@ -840,15 +847,16 @@ class Model { size_t processed, size_t start_position); void run_full_context_text(); uint32_t argmax_component_logits(Component& comp, size_t logit_row = std::numeric_limits::max(), - float* out_uncertainty = nullptr); - uint32_t argmax_logits_at(const BufferDesc& desc, void* ptr, size_t row_off, float* out_uncertainty); + float* out_uncertainty = nullptr, float repetition_penalty = 1.0f); + uint32_t argmax_logits_at(const BufferDesc& desc, void* ptr, size_t row_off, float* out_uncertainty, + float repetition_penalty = 1.0f); std::vector argmax_component_logits_batch(Component& comp, size_t batch); void write_int_input(Component& comp, const std::string& name, int64_t value); void write_int_input_at(Component& comp, const std::string& name, size_t index, int64_t value); void write_bytes_input(Component& comp, const std::string& name, const void* data, size_t byte_size); int input_index(const Component& comp, const std::string& name) const; int output_index(const Component& comp, const std::string& name) const; - uint32_t argmax_last_logits(float* out_uncertainty = nullptr); + uint32_t argmax_last_logits(float* out_uncertainty = nullptr, float repetition_penalty = 1.0f); bool load_handoff_probe(); void maybe_capture_handoff_probe_hidden(const Component& comp, const std::string& output_name = "probe_hidden"); void run_vision_encoder(const std::string& image_path); diff --git a/cactus-engine/src/engine_image.cpp b/cactus-engine/src/engine_image.cpp index 3fdb92e66..b24c1ab0a 100644 --- a/cactus-engine/src/engine_image.cpp +++ b/cactus-engine/src/engine_image.cpp @@ -912,11 +912,11 @@ static Siglip2Preprocessor::Config make_lfm2_vl_siglip_config(const Config& conf sp_cfg.rescale_factor = config.rescale_factor > 0.0f ? config.rescale_factor : (1.0f / 255.0f); sp_cfg.image_mean[0] = sp_cfg.image_mean[1] = sp_cfg.image_mean[2] = config.image_mean; sp_cfg.image_std[0] = sp_cfg.image_std[1] = sp_cfg.image_std[2] = config.image_std; - sp_cfg.do_image_splitting = true; - sp_cfg.do_resize = true; - sp_cfg.do_rescale = true; - sp_cfg.do_normalize = true; - sp_cfg.do_convert_rgb = true; + sp_cfg.do_image_splitting = config.do_image_splitting; + sp_cfg.do_resize = config.do_resize; + sp_cfg.do_rescale = config.do_rescale; + sp_cfg.do_normalize = config.do_normalize; + sp_cfg.do_convert_rgb = config.do_convert_rgb; return sp_cfg; } diff --git a/cactus-engine/src/init.cpp b/cactus-engine/src/init.cpp index 7672d88f2..e588e305c 100644 --- a/cactus-engine/src/init.cpp +++ b/cactus-engine/src/init.cpp @@ -1,6 +1,7 @@ #include "../cactus_engine.h" #include "utils.h" #include "telemetry.h" +#include "../../cactus-kernels/src/threading.h" #include #include #include @@ -19,7 +20,8 @@ static constexpr size_t RAG_MIN_CHUNK_TOKENS = 24; static constexpr size_t RAG_CHUNK_OVERLAP = 32; static void apply_no_cloud_telemetry_env() { - if (cactus::ffi::env_flag_enabled("CACTUS_NO_CLOUD_TELE")) { + if (cactus::ffi::env_flag_enabled("CACTUS_NO_CLOUD_TELE") || + cactus::ffi::env_flag_enabled("CACTUS_DISABLE_CLOUD_HANDOFF")) { cactus::telemetry::setCloudDisabled(true); } } @@ -437,7 +439,10 @@ cactus_model_t cactus_init(const char* model_path, const char* corpus_dir, bool } void cactus_destroy(cactus_model_t model) { - if (model) delete static_cast(model); + if (!model) return; + CactusThreading::get_thread_pool().wait_all(); + delete static_cast(model); + CactusThreading::get_thread_pool().wait_all(); } void cactus_reset(cactus_model_t model) { diff --git a/cactus-engine/src/model.cpp b/cactus-engine/src/model.cpp index 1e97380c3..a7601c9ce 100644 --- a/cactus-engine/src/model.cpp +++ b/cactus-engine/src/model.cpp @@ -75,6 +75,40 @@ void write_scalar_value(Precision precision, uint8_t* data, size_t index, float } } +bool is_media_generation_token(const Config& config, uint32_t token_id) { + if (config.image_token_id != 0 && token_id == config.image_token_id) return true; + if (config.audio_token_id != 0 && token_id == config.audio_token_id) return true; + + if (config.model_type == Config::ModelType::GEMMA4) { + switch (token_id) { + case 255999: // <|image> + case 258880: // <|image|> + case 258882: // + case 256000: // <|audio> + case 258881: // <|audio|> + case 258883: // + return true; + default: + break; + } + } + + if (config.model_type == Config::ModelType::LFM2 && config.use_image_tokens) { + switch (token_id) { + case 396: // + case 497: // <|img_thumbnail|> + case 498: // <|image_start|> + case 499: // <|image_end|> + case 500: // <|image_split|> + return true; + default: + break; + } + } + + return false; +} + bool copy_component_tensor(CactusGraph& source_graph, const BufferDesc& src_desc, size_t src_node, @@ -1076,25 +1110,28 @@ void Model::copy_component_outputs_to_chunk_inputs_range(const Component& source bool Model::cache_states_compatible(const Component& source, const Component& target) const { if (source.cache_states.empty() || source.cache_states.size() != target.cache_states.size()) return false; + if (!source.graph || !target.graph) return false; for (size_t i = 0; i < source.cache_states.size(); ++i) { const auto& src = source.cache_states[i]; const auto& dst = target.cache_states[i]; if (src.layer_key != dst.layer_key) return false; if (src.key_node_id < 0 || src.value_node_id < 0 || dst.key_node_id < 0 || dst.value_node_id < 0) return false; + for (auto [src_node, dst_node] : {std::pair{src.key_node_id, dst.key_node_id}, std::pair{src.value_node_id, dst.value_node_id}}) { + const size_t src_id = static_cast(src_node); + const size_t dst_id = static_cast(dst_node); + if (source.graph->get_node_op_type(src_id) != target.graph->get_node_op_type(dst_id)) return false; + } } return true; } void Model::move_cache_states(Component& source, Component& target, size_t logical_current) { - if (source.cache_states.empty() || source.cache_states.size() != target.cache_states.size()) { + if (!cache_states_compatible(source, target)) { throw std::runtime_error("prefill and step cache states are not compatible"); } for (size_t i = 0; i < source.cache_states.size(); ++i) { const auto& src = source.cache_states[i]; const auto& dst = target.cache_states[i]; - if (src.layer_key != dst.layer_key) { - throw std::runtime_error("prefill and step cache layer mismatch: " + src.layer_key + " != " + dst.layer_key); - } int moved_src = -1, moved_dst = -1; for (auto [src_node, dst_node] : {std::pair{src.key_node_id, dst.key_node_id}, std::pair{src.value_node_id, dst.value_node_id}}) { if (src_node < 0 || dst_node < 0) continue; @@ -1293,11 +1330,11 @@ Model::ChunkedPrefillResult Model::run_chunked_prefill(const std::vectorgraph->get_node_op_type(id) == OpType::KV_CACHE_STATE && decoder_prefill_->graph->get_node_window_size(id) > 0; }); - const size_t tail_tokens = tokens.size() - whole_chunks_end; - const size_t padding_cutoff = std::max(1, effective_chunk / 16); const bool has_conv_state = any_cache_node([&](size_t id) { return decoder_prefill_->graph->get_node_op_type(id) == OpType::CONV_CACHE_STATE; }); + const size_t tail_tokens = tokens.size() - whole_chunks_end; + const size_t padding_cutoff = std::max(1, effective_chunk / 16); const bool pad_tail = !has_conv_state && !has_recurrent_state && !has_sliding_window_cache @@ -1862,22 +1899,36 @@ void Model::run_audio_encoder(const std::vector& audio_features) { const size_t node_id = static_cast(audio_encoder_->output_node_ids[i]); const auto& desc = audio_encoder_->graph->get_output_buffer(node_id); void* ptr = audio_encoder_->graph->get_output(node_id); + const size_t output_rows = desc.shape.size() >= 2 + ? desc.shape[desc.shape.size() - 2] + : desc.shape.empty() ? 0 : desc.shape[0]; + size_t rows_to_copy = output_rows; + if (family_ == "gemma4" && output_rows > 0) { + rows_to_copy = std::min(output_rows, (frames_in_chunk + 3) / 4); + } + if (rows_to_copy == 0) continue; + const size_t row_bytes = output_rows > 0 ? desc.byte_size / output_rows : desc.byte_size; + const size_t append_bytes = rows_to_copy * row_bytes; auto& slot = media_features_[name]; const size_t prev_bytes = slot.size(); - slot.resize(prev_bytes + desc.byte_size); - std::memcpy(slot.data() + prev_bytes, ptr, desc.byte_size); + slot.resize(prev_bytes + append_bytes); + std::memcpy(slot.data() + prev_bytes, ptr, append_bytes); auto shape_it = media_feature_shapes_.find(name); if (shape_it == media_feature_shapes_.end() || shape_it->second.empty()) { media_feature_shapes_[name] = desc.shape; + auto& shape = media_feature_shapes_[name]; + if (shape.size() >= 2) shape[shape.size() - 2] = rows_to_copy; + else if (!shape.empty()) shape[0] = rows_to_copy; } else if (desc.shape.size() >= 2 && shape_it->second.size() == desc.shape.size()) { - shape_it->second[shape_it->second.size() - 2] += desc.shape[desc.shape.size() - 2]; + shape_it->second[shape_it->second.size() - 2] += rows_to_copy; } media_feature_precisions_[name] = desc.precision; } } } -uint32_t Model::argmax_logits_at(const BufferDesc& desc, void* ptr, size_t row_off, float* out_uncertainty) { +uint32_t Model::argmax_logits_at(const BufferDesc& desc, void* ptr, size_t row_off, float* out_uncertainty, + float repetition_penalty) { size_t vocab = desc.shape.empty() ? 0 : desc.shape.back(); uint32_t best = 0; float best_v = -std::numeric_limits::infinity(); @@ -1888,10 +1939,15 @@ uint32_t Model::argmax_logits_at(const BufferDesc& desc, void* ptr, size_t row_o if (tool_it != tool_bias.end()) value += tool_it->second; auto vocab_it = vocab_bias_.find(static_cast(token_id)); if (vocab_it != vocab_bias_.end()) value += vocab_it->second; + if (repetition_penalty > 1.0f + && std::find(token_history_.begin(), token_history_.end(), static_cast(token_id)) != token_history_.end()) { + value = value >= 0.0f ? value / repetition_penalty : value * repetition_penalty; + } return value; }; auto observe_logit = [&](size_t i, float v) { if (static_cast(i) == suppressed_token_id_) return; + if (is_media_generation_token(config_, static_cast(i))) return; v = score_with_bias(i, v); if (v > best_v) { second_v = best_v; @@ -1922,7 +1978,8 @@ uint32_t Model::argmax_logits_at(const BufferDesc& desc, void* ptr, size_t row_o return best; } -uint32_t Model::argmax_component_logits(Component& comp, size_t logit_row, float* out_uncertainty) { +uint32_t Model::argmax_component_logits(Component& comp, size_t logit_row, float* out_uncertainty, + float repetition_penalty) { size_t out_node = static_cast(comp.output_node_ids.empty() ? 0 : comp.output_node_ids[0]); const auto& desc = comp.graph->get_output_buffer(out_node); void* ptr = comp.graph->get_output(out_node); @@ -1934,7 +1991,7 @@ uint32_t Model::argmax_component_logits(Component& comp, size_t logit_row, float } else if (decode_route_ == DecodeRoute::FULL_CONTEXT_TEXT) { row = std::min(last_logit_position_, seq > 0 ? seq - 1 : 0); } - return argmax_logits_at(desc, ptr, row * vocab, out_uncertainty); + return argmax_logits_at(desc, ptr, row * vocab, out_uncertainty, repetition_penalty); } std::vector Model::argmax_component_logits_batch(Component& comp, size_t batch) { @@ -1954,8 +2011,8 @@ std::vector Model::argmax_component_logits_batch(Component& comp, size return out; } -uint32_t Model::argmax_last_logits(float* out_uncertainty) { - return argmax_component_logits(*decoder_, std::numeric_limits::max(), out_uncertainty); +uint32_t Model::argmax_last_logits(float* out_uncertainty, float repetition_penalty) { + return argmax_component_logits(*decoder_, std::numeric_limits::max(), out_uncertainty, repetition_penalty); } bool Model::prefill_and_sample_first_token(const std::vector& tokens, uint32_t& out_token, @@ -2008,9 +2065,7 @@ bool Model::prefill_and_sample_first_token(const std::vector& tokens, ChunkedPrefillResult chunked; if (decoder_prefill_) { chunked = run_chunked_prefill(tokens, cache_total_seq_len_, get_prefill_chunk_size(), true); - if (chunked.logical_tokens == tokens.size() - && tokens.size() > 0 - && (chunked.padding_tokens > 0 || family_ == "gemma4")) { + if (chunked.logical_tokens == tokens.size() && chunked.padding_tokens > 0 && tokens.size() > 0) { // Cache already moved into the step component; drop the last row and re-run it for step-decoder logits. set_cache_current_len(*decoder_, tokens.size() - 1); cache_total_seq_len_ = tokens.size() - 1; @@ -2233,6 +2288,21 @@ void Model::reset_encoder_cross_kv_route_state() { encoder_cross_kv_source_len_ = 0; } +void Model::release_encoder_cross_kv_prepare_graphs() { + if (source_encoder_ && source_encoder_ != decoder_) { + if (source_encoder_->graph) { + source_encoder_->graph->release_runtime_buffers(); + source_encoder_->graph->release_all_weight_pages(); + } + } + if (decoder_cross_kv_ && decoder_cross_kv_ != decoder_) { + if (decoder_cross_kv_->graph) { + decoder_cross_kv_->graph->release_runtime_buffers(); + decoder_cross_kv_->graph->release_all_weight_pages(); + } + } +} + bool Model::finish_encoder_cross_kv_prepare() { if (!source_encoder_ || !decoder_cross_kv_ || !decoder_) return false; source_encoder_->graph->execute(); @@ -2243,6 +2313,7 @@ bool Model::finish_encoder_cross_kv_prepare() { copy_component_outputs_to_inputs(*decoder_cross_kv_, *decoder_); } encoder_cross_kv_ready_ = true; + release_encoder_cross_kv_prepare_graphs(); return true; } @@ -2253,12 +2324,14 @@ bool Model::finish_encoder_cross_kv_prepare_after_source() { copy_component_outputs_to_inputs(*decoder_cross_kv_, *decoder_); } encoder_cross_kv_ready_ = true; + release_encoder_cross_kv_prepare_graphs(); return true; } bool Model::prepare_encoder_cross_kv_from_text(const std::vector& tokens) { if (!source_encoder_ || !decoder_cross_kv_ || !decoder_ || tokens.empty()) return false; if (encoder_cross_kv_source_kind_ != "text_tokens") return false; + if (!load_component_graph(*source_encoder_) || !load_component_graph(*decoder_cross_kv_)) return false; reset_encoder_cross_kv_route_state(); for (auto& buf : source_encoder_->input_buffers) { @@ -2297,6 +2370,7 @@ bool Model::prepare_encoder_cross_kv_from_text(const std::vector& toke bool Model::prepare_encoder_cross_kv_from_audio(const std::vector& audio_features) { if (!source_encoder_ || !decoder_cross_kv_ || !decoder_) return false; + if (!load_component_graph(*source_encoder_) || !load_component_graph(*decoder_cross_kv_)) return false; int feature_idx = input_index(*source_encoder_, "input_features"); if (feature_idx < 0) { feature_idx = input_index(*source_encoder_, "audio_features"); @@ -2893,7 +2967,8 @@ bool Model::run_chunk_prefill_path(const std::vector& tokens, const auto& feat_shape = media_feature_shapes_["image_features"]; const size_t hidden = emb_shape.empty() ? 0 : emb_shape.back(); const size_t emb_seq = emb_shape.size() >= 2 ? emb_shape[emb_shape.size() - 2] : 0; - const size_t feat_rows = feat_shape.empty() ? 0 : feat_shape[0]; + const size_t feat_rows = feat_shape.size() >= 2 ? feat_shape[feat_shape.size() - 2] + : feat_shape.empty() ? 0 : feat_shape[0]; const Precision emb_prec = store_prec["inputs_embeds"]; const Precision feat_prec = media_feature_precisions_["image_features"]; const size_t emb_elem = PrecisionTraits::size_of(emb_prec); @@ -2986,7 +3061,10 @@ bool Model::run_chunk_prefill_path(const std::vector& tokens, if (count > 0) valid_seq = count; } valid_seq = std::min(valid_seq, full_seq); - const size_t whole_chunks_end = (valid_seq / chunk_seq) * chunk_seq; + size_t whole_chunks_end = (valid_seq / chunk_seq) * chunk_seq; + if (whole_chunks_end > 0 && decoder_ != nullptr && !cache_states_compatible(*decoder_prefill_chunk_, *decoder_)) { + whole_chunks_end = 0; + } for (size_t chunk_start = 0; chunk_start < whole_chunks_end; chunk_start += chunk_seq) { for (const auto& kv : store_bytes) { const std::string& name = kv.first; @@ -3167,7 +3245,7 @@ void Model::prefill_with_media(const std::vector& tokens, uint32_t Model::decode(const std::vector& tokens, float /*temperature*/, float /*top_p*/, size_t /*top_k*/, const std::string& /*profile_file*/, float* out_entropy, - float /*min_p*/, float /*repetition_penalty*/) { + float /*min_p*/, float repetition_penalty) { if (tokens.empty()) return 0; if (decode_route_ == DecodeRoute::ENCODER_CROSS_KV_STEP) { if (!encoder_cross_kv_ready_ && encoder_cross_kv_source_kind_ == "text_tokens") { @@ -3191,7 +3269,7 @@ uint32_t Model::decode(const std::vector& tokens, float /*temperature* } cache_total_seq_len_ += tokens.size(); if (out_entropy) *out_entropy = 0.0f; - uint32_t result = argmax_last_logits(); + uint32_t result = argmax_last_logits(nullptr, repetition_penalty); record_sampled_token(result); return result; } @@ -3200,7 +3278,7 @@ uint32_t Model::decode(const std::vector& tokens, float /*temperature* run_full_context_text(); cache_total_seq_len_ = context_tokens_.size(); cache_token_ids_ = context_tokens_; - uint32_t result = argmax_last_logits(out_entropy); + uint32_t result = argmax_last_logits(out_entropy, repetition_penalty); record_sampled_token(result); return result; } @@ -3211,7 +3289,7 @@ uint32_t Model::decode(const std::vector& tokens, float /*temperature* cache_total_seq_len_ += tokens.size(); cache_token_ids_.insert(cache_token_ids_.end(), tokens.begin(), tokens.end()); maybe_roll_compact(); - uint32_t result = argmax_last_logits(out_entropy); + uint32_t result = argmax_last_logits(out_entropy, repetition_penalty); record_sampled_token(result); return result; } @@ -4203,6 +4281,10 @@ bool Config::from_json(const std::string& config_path) { else if (key == "tile_size") tile_size = static_cast(std::stoul(value)); else if (key == "max_pixels_tolerance") max_pixels_tolerance = std::stof(value); else if (key == "do_image_splitting") do_image_splitting = (value == "true" || value == "1"); + else if (key == "do_resize") do_resize = (value == "true" || value == "1"); + else if (key == "do_rescale") do_rescale = (value == "true" || value == "1"); + else if (key == "do_normalize") do_normalize = (value == "true" || value == "1"); + else if (key == "do_convert_rgb") do_convert_rgb = (value == "true" || value == "1"); else if (key == "precision") { if (value == "INT8") precision = Precision::INT8; else if (value == "FP16") precision = Precision::FP16; diff --git a/cactus-engine/src/telemetry_impl.cpp b/cactus-engine/src/telemetry_impl.cpp index 150440b53..0f5d85b1f 100644 --- a/cactus-engine/src/telemetry_impl.cpp +++ b/cactus-engine/src/telemetry_impl.cpp @@ -1302,6 +1302,11 @@ void init(const char* project_id_param, const char* project_scope_param, const c cloud_disabled_from_env = true; } } + if (const char* env = std::getenv("CACTUS_DISABLE_CLOUD_HANDOFF")) { + if (env[0] != '\0' && !(env[0] == '0' && env[1] == '\0')) { + cloud_disabled_from_env = true; + } + } const char* env_url = std::getenv("CACTUS_SUPABASE_URL"); const char* env_key = std::getenv("CACTUS_SUPABASE_KEY"); const char* env_project = std::getenv("CACTUS_PROJECT_ID"); diff --git a/cactus-engine/src/tokenizer.cpp b/cactus-engine/src/tokenizer.cpp index d792578db..cebdc6226 100644 --- a/cactus-engine/src/tokenizer.cpp +++ b/cactus-engine/src/tokenizer.cpp @@ -6,6 +6,7 @@ #include #include #include +#include #include namespace cactus { @@ -361,7 +362,7 @@ void Tokenizer::load_chat_template(const std::string& template_file) { } void Tokenizer::detect_model_type(const std::string& config_path) { - model_type_ = ModelType::GEMMA4; + model_type_ = ModelType::UNKNOWN; std::ifstream file(config_path); if (!file.is_open()) return; @@ -382,6 +383,11 @@ void Tokenizer::detect_model_type(const std::string& config_path) { model_type_ = ModelType::GEMMA4; } else if (lower_line.find("gemma") != std::string::npos) { model_type_ = ModelType::GEMMA; + } else if (lower_line.find("whisper") != std::string::npos) { + model_type_ = ModelType::WHISPER; + } else if (lower_line.find("parakeet_tdt") != std::string::npos || + lower_line.find("parakeet-tdt") != std::string::npos) { + model_type_ = ModelType::PARAKEET_TDT; } } else if (lower_line.find("model_variant") != std::string::npos) { if (lower_line.find("vlm") != std::string::npos) { model_variant_ = ModelVariant::VLM; } @@ -416,15 +422,32 @@ std::string Tokenizer::get_default_stop_sequence() const { if (model_type_ == ModelType::GEMMA) { return ""; } - return ""; + if (model_type_ == ModelType::GEMMA4) { + return ""; + } + return ""; } std::vector Tokenizer::apply_chat_template(const std::vector& messages, bool add_generation_prompt) const { return encode(format_chat_prompt(messages, add_generation_prompt)); } +bool Tokenizer::uses_default_chat_style(const std::vector& messages, const std::string& tools_json, + bool enable_thinking_if_supported) const { + return !has_chat_template_ && tools_json.empty() && !enable_thinking_if_supported && + std::all_of(messages.begin(), messages.end(), [](const ChatMessage& msg) { + return msg.images.empty() && msg.audio.empty() && msg.tool_calls.empty() && + msg.audio_soft_token_count == 0 && msg.role != "tool"; + }) && + (model_type_ == ModelType::UNKNOWN || model_type_ == ModelType::GEMMA4); +} + std::string Tokenizer::format_chat_prompt(const std::vector& messages, bool add_generation_prompt, const std::string& tools_json, bool enable_thinking_if_supported) const { + if (uses_default_chat_style(messages, tools_json, enable_thinking_if_supported)) { + return format_default_chat_style(messages, add_generation_prompt); + } + if (model_type_ == ModelType::QWEN) { return format_qwen_style(messages, add_generation_prompt, tools_json, enable_thinking_if_supported); } @@ -440,6 +463,30 @@ std::string Tokenizer::format_chat_prompt(const std::vector& messag return format_gemma4_style(messages, add_generation_prompt, tools_json, enable_thinking_if_supported); } +std::string Tokenizer::format_default_chat_style(const std::vector& messages, bool add_generation_prompt) const { + std::string result; + if (model_type_ == ModelType::GEMMA || model_type_ == ModelType::GEMMA4) { + result = ""; + } else if (model_type_ == ModelType::LFM2) { + result = "<|startoftext|>"; + } + + for (const auto& msg : messages) { + if (msg.role == "system" || msg.role == "developer") { + result += "System: " + msg.content + "\n"; + } else if (msg.role == "assistant" || msg.role == "model") { + result += "Assistant: " + msg.content + "\n"; + } else { + result += "User: " + msg.content + "\n"; + } + } + + if (add_generation_prompt) { + result += "Assistant:"; + } + return result; +} + std::string Tokenizer::format_gemma_style(const std::vector& messages, bool add_generation_prompt, const std::string& tools_json) const { std::string result = ""; @@ -716,6 +763,93 @@ std::string Tokenizer::format_gemma4_style(const std::vector& messa if (tw == 0) tw = side; return static_cast((th / p / k) * (tw / p / k)); }; + + if (tools_json.empty() && !enable_thinking_if_supported) { + bool has_media = false; + for (const auto& msg : messages) { + if (!msg.images.empty() || msg.audio_soft_token_count > 0) { + has_media = true; + break; + } + } + if (has_media && std::getenv("CACTUS_GEMMA4_RAW_MEDIA_PROMPT") != nullptr) { + result = ""; + for (size_t i = first_msg; i < messages.size(); i++) { + const auto& msg = messages[i]; + bool needs_separator = false; + auto append_separator = [&]() { + if (needs_separator) result += " "; + }; + for (const auto& image_path : msg.images) { + size_t n = compute_soft_tokens(image_path); + if (n > 0) { + append_separator(); + result += "<|image>"; + for (size_t j = 0; j < n; j++) result += "<|image|>"; + result += ""; + needs_separator = true; + } + } + if (msg.audio_soft_token_count > 0) { + append_separator(); + result += "<|audio>"; + for (size_t j = 0; j < msg.audio_soft_token_count; j++) result += "<|audio|>"; + result += ""; + needs_separator = true; + } + if (!msg.content.empty()) { + append_separator(); + result += msg.content; + } + if (i + 1 < messages.size()) result += "\n"; + } + return result; + } + + result = ""; + for (size_t i = first_msg; i < messages.size(); i++) { + const auto& msg = messages[i]; + if (msg.role == "system" || msg.role == "developer") { + result += "System: " + msg.content + "\n"; + continue; + } + if (msg.role == "assistant" || msg.role == "model") { + result += "Assistant: " + msg.content + "\n"; + continue; + } + + result += "User: "; + bool needs_separator = false; + auto append_separator = [&]() { + if (needs_separator) result += " "; + }; + for (const auto& image_path : msg.images) { + size_t n = compute_soft_tokens(image_path); + if (n > 0) { + append_separator(); + result += "<|image>"; + for (size_t j = 0; j < n; j++) result += "<|image|>"; + result += ""; + needs_separator = true; + } + } + if (msg.audio_soft_token_count > 0) { + append_separator(); + result += "<|audio>"; + for (size_t j = 0; j < msg.audio_soft_token_count; j++) result += "<|audio|>"; + result += ""; + needs_separator = true; + } + if (!msg.content.empty()) { + append_separator(); + result += msg.content; + } + result += "\n"; + } + if (add_generation_prompt) result += "Assistant:"; + return result; + } + bool in_model_turn = false; std::vector pending_call_names; auto close_model_turn = [&]() { @@ -754,13 +888,13 @@ std::string Tokenizer::format_gemma4_style(const std::vector& messa result += ""; } } - result += msg.content; if (msg.audio_soft_token_count > 0) { result += "<|audio>"; for (size_t j = 0; j < msg.audio_soft_token_count; j++) result += "<|audio|>"; result += ""; } + result += msg.content; result += "\n"; } } diff --git a/cactus-engine/src/utils.h b/cactus-engine/src/utils.h index 95e2368da..ce280dd29 100644 --- a/cactus-engine/src/utils.h +++ b/cactus-engine/src/utils.h @@ -432,15 +432,7 @@ inline AudioPreprocessResult preprocess_audio_for_gemma4( result.num_frames = mel.size() / mel_bins; result.features = transpose_mel_to_frame_major(mel, mel_bins, result.num_frames); - size_t max_soft = model_config.audio_soft_tokens; - size_t max_frames_per_chunk = max_soft > 0 ? (max_soft * 2 - 1) * 2 - 1 : 0; - if (max_soft > 0 && result.num_frames > max_frames_per_chunk) { - size_t num_chunks = (result.num_frames + max_frames_per_chunk - 1) / max_frames_per_chunk; - result.num_soft_tokens = num_chunks * max_soft; - } else { - size_t after_stage1 = (result.num_frames + 1) / 2; - result.num_soft_tokens = (after_stage1 + 1) / 2; - } + result.num_soft_tokens = (result.num_frames + 3) / 4; return result; } diff --git a/cactus-engine/tests/transcribe.cpp b/cactus-engine/tests/transcribe.cpp index 84d22b0e9..e00195c9b 100644 --- a/cactus-engine/tests/transcribe.cpp +++ b/cactus-engine/tests/transcribe.cpp @@ -92,7 +92,7 @@ static std::string whisper_prompt(const std::string& model_path, const std::stri std::string p = model_path; std::transform(p.begin(), p.end(), p.begin(), [](unsigned char c) { return (char)std::tolower(c); }); if (p.find("whisper") != std::string::npos) - return "<|startoftranscript|><|" + language + "|><|transcribe|>"; + return "<|startoftranscript|><|" + language + "|><|transcribe|><|notimestamps|>"; return ""; } diff --git a/cactus-graph/src/builder.cpp b/cactus-graph/src/builder.cpp index 10c4c94ec..764fd6bfe 100644 --- a/cactus-graph/src/builder.cpp +++ b/cactus-graph/src/builder.cpp @@ -13,6 +13,25 @@ bool use_fp16_kv_cache_for_builder() { return value != nullptr && std::strcmp(value, "1") == 0; } +bool op_honors_explicit_output_precision(OpType op_type) { + switch (op_type) { + case OpType::PRECISION_CAST: + case OpType::EMBEDDING: + case OpType::TOPK: + case OpType::SCATTER_TOPK: + case OpType::SAMPLE: + case OpType::KV_CACHE_APPEND: + case OpType::CONV_CACHE_STATE: + case OpType::RECURRENT_CACHE_STATE: + case OpType::IMAGE_PREPROCESS: + case OpType::MEL_FILTER_BANK: + case OpType::SPECTROGRAM: + return true; + default: + return false; + } +} + } // namespace size_t CactusGraph::attach_conv_bias(size_t node, size_t bias, size_t expected_size, const char* op_name) { @@ -1429,7 +1448,7 @@ size_t CactusGraph::add_node(OpType op_type, const std::vector& inputs, } Precision result_precision = params.output_precision; - if (op_type == OpType::PRECISION_CAST || op_type == OpType::EMBEDDING) { + if (op_honors_explicit_output_precision(op_type)) { result_precision = params.output_precision; } else if (!inputs.empty()) { result_precision = nodes_[node_index_map_[inputs[0]]]->output_buffer.precision; diff --git a/cactus-graph/src/core.cpp b/cactus-graph/src/core.cpp index 7dc709ec6..5a6b77966 100644 --- a/cactus-graph/src/core.cpp +++ b/cactus-graph/src/core.cpp @@ -234,6 +234,7 @@ void BufferDesc::set_external(void* ptr) { delete[] pooled_data; } pooled_data = nullptr; + pooled_byte_size = 0; } GraphNode::GraphNode(size_t node_id, OpType type) : id(node_id), op_type(type) {} diff --git a/cactus-graph/src/io.cpp b/cactus-graph/src/io.cpp index 7a7ff543f..61f267dd1 100644 --- a/cactus-graph/src/io.cpp +++ b/cactus-graph/src/io.cpp @@ -53,6 +53,36 @@ namespace { return filename; } + bool shape_matches_mmap_weight(const std::vector& logical_shape, + const std::vector& file_shape, + Precision precision, + const GraphFile::MappedFile& mapped_file) { + if (logical_shape == file_shape) { + return true; + } + + if (!PrecisionTraits::is_cq(precision) || mapped_file.group_size() == 0 || logical_shape.size() != file_shape.size()) { + return false; + } + + const size_t padded_k = mapped_file.group_size() * mapped_file.num_groups(); + size_t padded_dims = 0; + for (size_t i = 0; i < logical_shape.size(); ++i) { + if (logical_shape[i] == file_shape[i]) { + continue; + } + + if (file_shape[i] == padded_k && logical_shape[i] < file_shape[i]) { + ++padded_dims; + continue; + } + + return false; + } + + return padded_dims == 1; + } + inline void write_u32(std::ostream& out, uint32_t v) { out.write(reinterpret_cast(&v), sizeof(v)); } @@ -545,7 +575,7 @@ void CactusGraph::bind_mmap_weights(size_t node_id, const std::string& filename) const auto& shape = mapped_file->shape(); Precision precision = mapped_file->precision(); auto& buffer = node.output_buffer; - if (buffer.shape != shape) { + if (!shape_matches_mmap_weight(buffer.shape, shape, precision, *mapped_file)) { throw std::runtime_error("mmap weight shape mismatch for node " + std::to_string(node_id)); } if (buffer.precision != precision) { diff --git a/cactus-graph/src/ops_nn.cpp b/cactus-graph/src/ops_nn.cpp index 8ab813c05..3caffa739 100644 --- a/cactus-graph/src/ops_nn.cpp +++ b/cactus-graph/src/ops_nn.cpp @@ -88,6 +88,7 @@ namespace { thread_local std::vector moe_expert_offsets_buf; thread_local std::vector moe_expert_tokens_buf; thread_local std::vector moe_routing_denom_buf; + constexpr float MOE_FP16_MAX = 65504.0f; void ensure_moe_buffers(size_t max_tokens, size_t hidden_dim, size_t intermediate_dim, size_t num_experts, size_t top_k) { @@ -125,6 +126,30 @@ namespace { throw std::runtime_error("moe_layer only supports FP16 or TQ expert weights"); } + + inline __fp16 moe_finite_f16(float value) { + if (!std::isfinite(value)) { + value = std::copysign(MOE_FP16_MAX, value); + } + value = std::clamp(value, -MOE_FP16_MAX, MOE_FP16_MAX); + return static_cast<__fp16>(value); + } + + void sanitize_moe_fp16(__fp16* data, size_t count) { + for (size_t i = 0; i < count; ++i) { + float value = static_cast(data[i]); + if (!std::isfinite(value) || std::abs(value) > MOE_FP16_MAX) { + data[i] = moe_finite_f16(value); + } + } + } + + void multiply_moe_fp16_sanitized(__fp16* lhs, const __fp16* rhs, size_t count) { + for (size_t i = 0; i < count; ++i) { + float value = static_cast(lhs[i]) * static_cast(rhs[i]); + lhs[i] = moe_finite_f16(value); + } + } } void compute_moe_layer_node(GraphNode& node, const std::vector>& nodes, const std::unordered_map& node_index_map) { @@ -148,6 +173,9 @@ void compute_moe_layer_node(GraphNode& node, const std::vector