MOE

cactus-engine/src/model.cpp

@@ -1187,7 +1187,6 @@ void Model::move_cache_states(Component& source, Component& target, size_t logic
                 logical_current != std::numeric_limits<size_t>::max()) {
                 auto* meta = static_cast<uint64_t*>(target.graph->get_output(static_cast<size_t>(dst_node)));
                 if (meta && logical_current < meta[0]) {
-                    meta[4] = std::min<uint64_t>(meta[4], logical_current);
                     meta[0] = logical_current;
                 }
             }
@@ -1202,7 +1201,6 @@ void Model::set_cache_current_len(Component& comp, size_t len) {
             if (comp.graph->get_node_op_type(static_cast<size_t>(node_id)) != OpType::KV_CACHE_STATE) continue;
             auto* meta = static_cast<uint64_t*>(comp.graph->get_output(static_cast<size_t>(node_id)));
             if (!meta || len >= meta[0]) continue;
-            meta[4] = std::min<uint64_t>(meta[4], len);
             meta[0] = len;
         }
     }

cactus-engine/tests/run.cpp

@@ -247,6 +247,9 @@ struct TokenPrinter {
     std::string pending;
     bool in_code_fence = false;
     bool color = false;
+    bool suppress_thinking_stream = false;
+    bool show_thinking = false;
+    std::string stream_carry;
     bool label_printed = false;
     std::thread spinner_thread;
     std::atomic<bool> spinner_running{false};
@@ -261,10 +264,12 @@ struct TokenPrinter {
         in_code_fence = false;
         label_printed = false;
         color = stdout_is_terminal();
+        suppress_thinking_stream = false;
+        stream_carry.clear();
     }
 
     void start_thinking() {
-        if (!color) return;
+        if (!color || spinner_running.load()) return;
         spinner_running = true;
         spinner_thread = std::thread([this]() {
             const char* frames[] = {"⠋", "⠙", "⠹", "⠸", "⠼", "⠴", "⠦", "⠧", "⠇", "⠏"};
@@ -434,25 +439,61 @@ struct TokenPrinter {
     }
 
     void on_token(const char* text) {
-        print_label();
         if (!saw_first) {
             first = std::chrono::steady_clock::now();
             saw_first = true;
         }
         ++count;
         if (!text) return;
-        pending += text;
+        feed(text);
+        std::cout << std::flush;
+    }
+
+    void feed(const std::string& chunk) {
+        static const std::vector<std::string> markers = {
+            "<|channel>", "<channel|>", "<|turn>", "<turn|>", "<|think|>", "<think|>",
+        };
+        stream_carry += chunk;
+        size_t i = 0;
+        while (i < stream_carry.size()) {
+            const std::string* matched = nullptr;
+            bool maybe_partial = false;
+            for (const auto& m : markers) {
+                if (stream_carry.compare(i, m.size(), m) == 0) { matched = &m; break; }
+                size_t avail = stream_carry.size() - i;
+                if (avail < m.size() && stream_carry.compare(i, avail, m, 0, avail) == 0) maybe_partial = true;
+            }
+            if (matched) {
+                if (*matched == "<|channel>") { suppress_thinking_stream = true; if (!show_thinking) start_thinking(); }
+                else if (*matched == "<channel|>") { suppress_thinking_stream = false; stop_thinking(); }
+                i += matched->size();
+                continue;
+            }
+            if (maybe_partial) break;
+            if (!suppress_thinking_stream || show_thinking) emit_visible(stream_carry[i]);
+            ++i;
+        }
+        stream_carry.erase(0, i);
+    }
+
+    void emit_visible(char c) {
+        print_label();
+        pending += c;
         size_t nl;
         while ((nl = pending.find('\n')) != std::string::npos) {
             render_line(pending.substr(0, nl));
             std::cout << "\n";
             pending.erase(0, nl + 1);
         }
-        std::cout << std::flush;
     }
 
     // Flush the trailing partial line once generation completes (no trailing newline).
     void finish() {
+        stop_thinking();
+        if (!stream_carry.empty() && (!suppress_thinking_stream || show_thinking)) {
+            for (char c : stream_carry) emit_visible(c);
+        }
+        stream_carry.clear();
         print_label();
         if (!pending.empty()) {
             render_line(pending);
@@ -893,6 +934,7 @@ int main(int argc, char** argv) {
     std::vector<uint8_t> current_pcm;
     TokenPrinter printer;
     g_printer = &printer;
+    printer.show_thinking = thinking;
     bool auto_send = !initial_prompt.empty() || !current_audio.empty() || !current_image.empty();
 
     while (true) {
@@ -990,7 +1032,7 @@ int main(int argc, char** argv) {
             + ",\"auto_handoff\":" + (auto_handoff ? "true" : "false")
             + ",\"confidence_threshold\":" + std::to_string(confidence_threshold)
             + ",\"cloud_timeout_ms\":" + std::to_string(cloud_timeout_ms)
-            + ",\"stop_sequences\":[\"<|im_end|>\",\"<end_of_turn>\"]}";
+            + ",\"stop_sequences\":[\"<|im_end|>\",\"<end_of_turn>\",\"<turn|>\"]}";
 
         if (!current_image.empty()) std::cout << "[image: " << current_image << "]\n";
         if (!current_audio.empty()) std::cout << "[audio: " << current_audio << "]\n";

cactus-graph/cactus_graph.h

@@ -1073,7 +1073,7 @@ CACTUS_FFI_EXPORT int cactus_graph_gaussian_topk(
 CACTUS_FFI_EXPORT int cactus_graph_moe_layer_gated(
     cactus_graph_t graph, cactus_node_t hidden, cactus_node_t routing_probs, cactus_node_t topk_indices,
     const cactus_node_t* w1_weights, const cactus_node_t* w3_weights, const cactus_node_t* w2_weights,
-    size_t num_experts, size_t num_experts_per_tok, bool normalize_routing, float epsilon, float routed_scaling_factor, cactus_node_t* out);
+    size_t num_experts, size_t num_experts_per_tok, bool normalize_routing, float epsilon, float routed_scaling_factor, int32_t activation, cactus_node_t* out);
 CACTUS_FFI_EXPORT int cactus_graph_moe_layer_ungated(
     cactus_graph_t graph, cactus_node_t hidden, cactus_node_t routing_probs, cactus_node_t topk_indices,
     const cactus_node_t* w1_weights, const cactus_node_t* w2_weights,

cactus-graph/src/builder.cpp

@@ -1220,7 +1220,12 @@ size_t CactusGraph::add_node(OpType op_type, const std::vector<size_t>& inputs,
     }
 
     Precision result_precision = params.output_precision;
-    if (op_type == OpType::PRECISION_CAST || op_type == OpType::EMBEDDING) {
+    if (op_type == OpType::PRECISION_CAST ||
+        op_type == OpType::EMBEDDING ||
+        op_type == OpType::TOPK ||
+        op_type == OpType::SCATTER_TOPK ||
+        op_type == OpType::SAMPLE ||
+        op_type == OpType::GAUSSIAN_TOPK) {
         result_precision = params.output_precision;
     } else if (!inputs.empty()) {
         result_precision = nodes_[node_index_map_[inputs[0]]]->output_buffer.precision;
@@ -1305,8 +1310,12 @@ size_t CactusGraph::altup_correct(size_t coefs, size_t innovation, const size_t*
 
 size_t CactusGraph::gaussian_topk(size_t input, float ppf) {
     const auto& in_buf = get_output_buffer(input);
+    if (in_buf.precision != Precision::FP16) {
+        throw std::runtime_error("gaussian_topk only supports FP16 input");
+    }
     OpParams params;
     params.scalar = ppf;
+    params.output_precision = in_buf.precision;
     return add_node(OpType::GAUSSIAN_TOPK, {input}, in_buf.shape, params);
 }
 

cactus-graph/src/graph_ffi.cpp

@@ -1341,7 +1341,7 @@ int cactus_graph_gaussian_topk(cactus_graph_t graph, cactus_node_t input, float
 
 int cactus_graph_moe_layer_gated(cactus_graph_t graph, cactus_node_t hidden, cactus_node_t routing_probs, cactus_node_t topk_indices,
                                  const cactus_node_t* w1_weights, const cactus_node_t* w3_weights, const cactus_node_t* w2_weights,
-                                 size_t num_experts, size_t num_experts_per_tok, bool normalize_routing, float epsilon, float routed_scaling_factor, cactus_node_t* out) {
+                                 size_t num_experts, size_t num_experts_per_tok, bool normalize_routing, float epsilon, float routed_scaling_factor, int32_t activation, cactus_node_t* out) {
     if (!graph || !w1_weights || !w3_weights || !w2_weights || !out) return fail_invalid("Invalid args to cactus_graph_moe_layer_gated");
     try {
         std::vector<size_t> w1(num_experts), w3(num_experts), w2(num_experts);
@@ -1350,7 +1350,7 @@ int cactus_graph_moe_layer_gated(cactus_graph_t graph, cactus_node_t hidden, cac
             w3[i] = static_cast<size_t>(w3_weights[i]);
             w2[i] = static_cast<size_t>(w2_weights[i]);
         }
-        *out = static_cast<cactus_node_t>(as_graph(graph)->graph.moe_layer(static_cast<size_t>(hidden), static_cast<size_t>(routing_probs), static_cast<size_t>(topk_indices), w1, w3, w2, num_experts, num_experts_per_tok, normalize_routing, epsilon, routed_scaling_factor));
+        *out = static_cast<cactus_node_t>(as_graph(graph)->graph.moe_layer(static_cast<size_t>(hidden), static_cast<size_t>(routing_probs), static_cast<size_t>(topk_indices), w1, w3, w2, num_experts, num_experts_per_tok, normalize_routing, epsilon, routed_scaling_factor, static_cast<Activation>(activation)));
         return 0;
     } catch (const std::exception& e) {
         last_error_message = e.what();

cactus-graph/src/ops_nn.cpp

@@ -83,6 +83,7 @@ void compute_matmul_node(GraphNode& node, const std::vector<std::unique_ptr<Grap
 namespace {
     thread_local std::vector<__fp16> moe_compact_hidden_buf;
     thread_local std::vector<__fp16> moe_gate_buf;
+    thread_local std::vector<__fp16> moe_gate_pad_buf;
     thread_local std::vector<__fp16> moe_up_buf;
     thread_local std::vector<__fp16> moe_expert_out_buf;
     thread_local std::vector<size_t> moe_expert_offsets_buf;
@@ -185,18 +186,21 @@ void compute_moe_layer_node(GraphNode& node, const std::vector<std::unique_ptr<G
     size_t* expert_offsets = moe_expert_offsets_buf.data(); 
     size_t* expert_tokens_flat = moe_expert_tokens_buf.data();  
 
+    auto expert_index = [&](float raw_idx) -> size_t {
+        if (!std::isfinite(raw_idx) || raw_idx < 0.0f) {
+            throw std::runtime_error("moe_layer got invalid expert index");
+        }
+        size_t idx = static_cast<size_t>(raw_idx + 0.5f);
+        if (idx >= num_experts) {
+            throw std::runtime_error("moe_layer got expert index out of range");
+        }
+        return idx;
+    };
+
     std::memset(expert_offsets, 0, (num_experts + 1) * sizeof(size_t));
     for (size_t tok = 0; tok < token_count; ++tok) {
         for (size_t k = 0; k < top_k; ++k) {
-            float raw_idx = topk_idx[tok * top_k + k];
-            if (!std::isfinite(raw_idx)) {
-                throw std::runtime_error("moe_layer got non-finite expert index");
-            }
-            size_t idx = static_cast<size_t>(raw_idx + 0.5f);
-            if (idx >= num_experts) {
-                throw std::runtime_error("moe_layer got expert index out of range");
-            }
-            expert_offsets[idx + 1]++;
+            expert_offsets[expert_index(topk_idx[tok * top_k + k]) + 1]++;
         }
     }
 
@@ -210,7 +214,7 @@ void compute_moe_layer_node(GraphNode& node, const std::vector<std::unique_ptr<G
 
     for (size_t tok = 0; tok < token_count; ++tok) {
         for (size_t k = 0; k < top_k; ++k) {
-            size_t idx = static_cast<size_t>(topk_idx[tok * top_k + k] + 0.5f);
+            size_t idx = expert_index(topk_idx[tok * top_k + k]);
             expert_tokens_flat[moe_write_cursors[idx]++] = tok;
         }
     }
@@ -220,7 +224,7 @@ void compute_moe_layer_node(GraphNode& node, const std::vector<std::unique_ptr<G
         for (size_t tok = 0; tok < token_count; ++tok) {
             float sum_probs = 0.0f;
             for (size_t k = 0; k < top_k; ++k) {
-                size_t idx = static_cast<size_t>(topk_idx[tok * top_k + k] + 0.5f);
+                size_t idx = expert_index(topk_idx[tok * top_k + k]);
                 sum_probs += routing_prob(tok, idx);
             }
             routing_denom[tok] = sum_probs + eps;
@@ -283,7 +287,23 @@ void compute_moe_layer_node(GraphNode& node, const std::vector<std::unique_ptr<G
             cactus_multiply_f16(gate, up, gate, selected_count * expert_intermediate_dim);
         }
 
-        moe_matmul(gate, selected_count, expert_intermediate_dim, w2_buffer, expert_out, hidden_dim);
+        const size_t w2_k = w2_buffer.shape.size() == 2 ? w2_buffer.shape[1] : 0;
+        if (w2_k < expert_intermediate_dim) {
+            throw std::runtime_error("moe_layer down-proj weight K smaller than expert intermediate dim");
+        }
+        const __fp16* w2_input = gate;
+        if (w2_k != expert_intermediate_dim) {
+            if (moe_gate_pad_buf.size() < selected_count * w2_k) moe_gate_pad_buf.resize(selected_count * w2_k);
+            std::memset(moe_gate_pad_buf.data(), 0, selected_count * w2_k * sizeof(__fp16));
+            for (size_t i = 0; i < selected_count; ++i) {
+                std::memcpy(moe_gate_pad_buf.data() + i * w2_k,
+                            gate + i * expert_intermediate_dim,
+                            expert_intermediate_dim * sizeof(__fp16));
+            }
+            w2_input = moe_gate_pad_buf.data();
+        }
+
+        moe_matmul(w2_input, selected_count, w2_k, w2_buffer, expert_out, hidden_dim);
 
         for (size_t i = 0; i < selected_count; ++i) {
             const size_t tok = selected_tokens[i];
@@ -934,6 +954,3 @@ void compute_groupnorm_node(GraphNode& node, const std::vector<std::unique_ptr<G
         }
     }
 }
-
-
-

python/cactus/bindings/cactus.py

@@ -608,7 +608,7 @@ class cactus_tensor_info_t(ctypes.Structure):
 _lib.cactus_graph_moe_layer_gated.argtypes = [
     cactus_graph_t, cactus_node_t, cactus_node_t, cactus_node_t,
     ctypes.POINTER(cactus_node_t), ctypes.POINTER(cactus_node_t), ctypes.POINTER(cactus_node_t),
-    ctypes.c_size_t, ctypes.c_size_t, ctypes.c_bool, ctypes.c_float, ctypes.c_float, ctypes.POINTER(cactus_node_t)
+    ctypes.c_size_t, ctypes.c_size_t, ctypes.c_bool, ctypes.c_float, ctypes.c_float, ctypes.c_int32, ctypes.POINTER(cactus_node_t)
 ]
 _lib.cactus_graph_moe_layer_gated.restype = ctypes.c_int
 _bind_optional(
@@ -2560,7 +2560,7 @@ def gaussian_topk(self, x, ppf):
         return self._tensor_from_node(out.value)
 
     def moe_layer_gated(self, hidden, routing_probs, topk_indices, w1_weights, w3_weights, w2_weights,
-                        num_experts, num_experts_per_tok, normalize_routing=True, epsilon=1e-6, routed_scaling_factor=1.0):
+                        num_experts, num_experts_per_tok, normalize_routing=True, epsilon=1e-6, routed_scaling_factor=1.0, activation=0):
         hidden = self._ensure_tensor(hidden)
         routing_probs = self._ensure_tensor(routing_probs)
         topk_indices = self._ensure_tensor(topk_indices)
@@ -2572,7 +2572,7 @@ def moe_layer_gated(self, hidden, routing_probs, topk_indices, w1_weights, w3_we
             self.h, cactus_node_t(hidden.id), cactus_node_t(routing_probs.id), cactus_node_t(topk_indices.id),
             w1, w3, w2, ctypes.c_size_t(int(num_experts)), ctypes.c_size_t(int(num_experts_per_tok)),
             ctypes.c_bool(bool(normalize_routing)), ctypes.c_float(float(epsilon)),
-            ctypes.c_float(float(routed_scaling_factor)), ctypes.byref(out)
+            ctypes.c_float(float(routed_scaling_factor)), ctypes.c_int32(int(activation)), ctypes.byref(out)
         )
         if rc != 0:
             raise RuntimeError(_err("graph_moe_layer_gated failed"))

python/cactus/cli/__init__.py

@@ -209,6 +209,7 @@ def create_parser():
     --torch-dtype <dtype>              float16 | float32 | bfloat16
     --trust-remote-code                allow HF remote code during the build
     --local-files-only                 require model/processor to be local
+    --low-memory-load                  graph capture with meta tensors
     --allow-unconverted-weights        debug-only: skip the CQ-weights check
     --execute-after-transpile          run a reference execution after building
     --graph-filename <name>            override saved graph filename
@@ -354,6 +355,10 @@ 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("--skip-model-load", action="store_true",
+                                help="Convert directly from checkpoint tensors without loading the full HF model object")
+    convert_parser.add_argument("--low-memory-load", action="store_true",
+                                help="Avoid loading checkpoint tensors during graph capture")
     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",

python/cactus/cli/convert.py

@@ -65,6 +65,7 @@ def cmd_convert(args):
             token=args.token,
             reconvert=args.reconvert,
             output_dir=output_dir,
+            skip_model_load=bool(getattr(args, "skip_model_load", False)),
         )
         if getattr(args, "weights_only", False):
             return 0
@@ -83,6 +84,7 @@ 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
+    from cactus.transpile.component_plan import infer_component_plan_from_output
 
     extra_args = []
     if args.weights_dir:
@@ -100,11 +102,21 @@ def cmd_transpile(args):
 
     image_files = list(args.image_file or [])
     audio_file = args.audio_file
-    if not image_files or not audio_file:
+    weights_dir = Path(args.weights_dir).expanduser() if args.weights_dir else get_weights_dir(args.model_id)
+    from .model import _AUDIO_TASKS
+    plan = infer_component_plan_from_output(str(weights_dir), model_id=args.model_id)
+    default_needs_image = bool(
+        plan.needs_image if plan is not None else args.task == "multimodal_causal_lm_logits"
+    )
+    default_needs_audio = bool(
+        plan.needs_audio if plan is not None
+        else (args.task in _AUDIO_TASKS or args.task == "multimodal_causal_lm_logits")
+    )
+    if (default_needs_image and not image_files) or (default_needs_audio and not audio_file):
         default_images, default_audio = _default_multimodal_assets()
-        if not image_files:
+        if default_needs_image and not image_files:
             image_files = default_images
-        if not audio_file and default_audio:
+        if default_needs_audio and not audio_file and default_audio:
             audio_file = default_audio
 
     for img in image_files:
@@ -125,6 +137,8 @@ def cmd_transpile(args):
         extra_args.append("--trust-remote-code")
     if args.local_files_only:
         extra_args.append("--local-files-only")
+    if getattr(args, "low_memory_load", False):
+        extra_args.append("--low-memory-load")
     if args.artifact_dir:
         extra_args.extend(["--artifact-dir", args.artifact_dir])
     if args.graph_filename: