mplang-nightly 0.1.dev203__tar.gz → 0.1.dev266__tar.gz

{mplang_nightly-0.1.dev203 → mplang_nightly-0.1.dev266}/.gitignore

@@ -208,6 +208,7 @@ __marimo__/
 
 # mplang
 app.log
+.mpl/
 
 # Ignore pylint config files (we use ruff instead)
 .pylintrc

{mplang_nightly-0.1.dev203 → mplang_nightly-0.1.dev266}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: mplang-nightly
-Version: 0.1.dev203
+Version: 0.1.dev266
 Summary: Multi-Party Programming Language
 Author-email: SecretFlow Team <secretflow-contact@service.alipay.com>
 License:                                  Apache License
@@ -205,15 +205,21 @@ License:                                  Apache License
            See the License for the specific language governing permissions and
            limitations under the License.
 License-File: LICENSE
-Requires-Python: >=3.10
+Requires-Python: <3.13,>=3.11
+Requires-Dist: coincurve>=20.0.0
+Requires-Dist: cryptography>=43.0.0
 Requires-Dist: duckdb>=1.0.0
 Requires-Dist: fastapi
-Requires-Dist: httpx>=0.27.0
+Requires-Dist: flax>=0.12.0
+Requires-Dist: httpx<1.0.0,>=0.27.0
+Requires-Dist: jax[cpu]==0.8.0
 Requires-Dist: lightphe<0.1.0,>=0.0.15
+Requires-Dist: numpy>=2.0.0
 Requires-Dist: pandas>=2.0.0
 Requires-Dist: protobuf<6.0,>=5.0
 Requires-Dist: pyarrow>=14.0.0
-Requires-Dist: spu==0.9.4.dev20250827
+Requires-Dist: pyyaml>=6.0
+Requires-Dist: spu>=0.10.0.dev20251208
 Requires-Dist: sqlglot>=23.0.0
 Requires-Dist: tenseal==0.3.16
 Requires-Dist: typing-extensions
@@ -242,7 +248,7 @@ multiple parties in a synchronous, SPMD (Single Program, Multiple Data) fashion.
 
 ### Installation
 
-You'll need a modern Python environment (3.10+). We recommend using `uv` for fast installation.
+You'll need a modern Python environment (3.11+). We recommend using `uv` for fast installation.
 
 ```bash
 # Install uv (if not already installed)

{mplang_nightly-0.1.dev203 → mplang_nightly-0.1.dev266}/README.md

@@ -20,7 +20,7 @@ multiple parties in a synchronous, SPMD (Single Program, Multiple Data) fashion.
 
 ### Installation
 
-You'll need a modern Python environment (3.10+). We recommend using `uv` for fast installation.
+You'll need a modern Python environment (3.11+). We recommend using `uv` for fast installation.
 
 ```bash
 # Install uv (if not already installed)

{mplang_nightly-0.1.dev203/examples → mplang_nightly-0.1.dev266/examples/v1}/stax_nn/stax_nn.py

@@ -27,11 +27,11 @@ import yaml
 from jax.example_libraries import stax
 from sklearn.metrics import accuracy_score
 
-import mplang as mp
+import mplang.v1 as mp
 
 parser = argparse.ArgumentParser(description="distributed driver.")
 parser.add_argument("--model", default="network_a", type=str)
-parser.add_argument("-c", "--config", default="examples/conf/3pc.yaml", type=str)
+parser.add_argument("-c", "--config", default="examples/v1/conf/3pc.yaml", type=str)
 parser.add_argument("-e", "--epoch", default=5, type=int)
 parser.add_argument("-b", "--batch_size", default=128, type=int)
 parser.add_argument("-o", "--optimizer", default="SGD", type=str)

mplang_nightly-0.1.dev266/examples/v1/xgboost/bench_fhe_hist.py

@@ -0,0 +1,615 @@
+# Copyright 2025 Ant Group Co., Ltd.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""
+Microbenchmark for FHE(BFV) histogram path in SecureBoost.
+
+This script measures the time to compute PP-side cumulative bucket sums
+via encrypted ct·ct dot products using TenSEAL/SEAL BFV vector backend.
+
+It provides two modes:
+- classic: separate g/h ciphertexts + ct·ct dot (current training path)
+- interleaved: interleave g/h into one ct, do one ct·ct mul + two ct·pt dots (even/odd)
+
+Usage examples:
+    uv run -q python examples/v1/xgboost/bench_fhe_hist.py --world-size 2 --m 4096 --n-total 16 --n-ap 4 --k 16 --t 4 --reps 3 --mode classic
+    uv run -q python examples/v1/xgboost/bench_fhe_hist.py --world-size 2 --m 4096 --n-total 16 --n-ap 4 --k 16 --t 4 --reps 3 --mode interleaved
+"""
+
+from __future__ import annotations
+
+import argparse
+import time
+from functools import partial
+
+import jax
+import jax.numpy as jnp
+import numpy as np
+from examples.xgboost.sgb import (
+    DEFAULT_FXP_BITS,
+    batch_feature_wise_bucket_sum_fhe_vector,
+    build_bins_equi_width,
+    compute_bin_indices,
+    compute_gh,
+    compute_init_pred,
+    quantize_gh,
+)
+
+import mplang.v1 as mp
+from mplang.v1.ops import fhe
+
+
+def _gen_data(n_samples: int, n_total_features: int, n_features_ap: int, seed: int):
+    rng = np.random.default_rng(seed)
+    X = rng.normal(size=(n_samples, n_total_features)).astype(np.float32)
+    # make a simple linear label with noise
+    w = rng.normal(size=(n_total_features,)).astype(np.float32)
+    z = X @ w + 0.1 * rng.normal(size=(n_samples,)).astype(np.float32)
+    p = 1 / (1 + np.exp(-z))
+    y = (p > 0.5).astype(np.float32)
+
+    X_ap = X[:, :n_features_ap]
+    X_pp = X[:, n_features_ap:]
+    return X_ap, X_pp, y
+
+
+@mp.function
+def _bench_once(
+    ap_id: int,
+    pp_ids: list[int],
+    X_ap: np.ndarray,
+    X_pp_splits: list[np.ndarray],
+    y: np.ndarray,
+    k: int,
+    t: int,
+    reps: int,
+    mode: str,
+    include_precompute: bool,
+    breakdown: bool,
+):
+    # Place data
+    X_ap_j = mp.run_jax_at(ap_id, lambda x: x, jnp.array(X_ap, dtype=jnp.float32))
+    X_pp_j = [
+        mp.run_jax_at(pp, lambda x: x, jnp.array(xpp, dtype=jnp.float32))
+        for pp, xpp in zip(pp_ids, X_pp_splits, strict=True)
+    ]
+    y_j = mp.run_jax_at(ap_id, lambda x: x, jnp.array(y, dtype=jnp.float32))
+
+    # Binning per party
+    build_bins_vmapped = jax.vmap(partial(build_bins_equi_width, max_bin=k), in_axes=1)
+    compute_indices_vmapped = jax.vmap(compute_bin_indices, in_axes=(1, 0), out_axes=1)
+
+    ap_bins = mp.run_jax_at(ap_id, build_bins_vmapped, X_ap_j)
+    _ = mp.run_jax_at(ap_id, compute_indices_vmapped, X_ap_j, ap_bins)
+
+    pp_bins = [
+        mp.run_jax_at(pp, build_bins_vmapped, X_pp_j[i]) for i, pp in enumerate(pp_ids)
+    ]
+    pp_idx = [
+        mp.run_jax_at(pp, compute_indices_vmapped, X_pp_j[i], pp_bins[i])
+        for i, pp in enumerate(pp_ids)
+    ]
+
+    # AP GH + quantize + encrypt
+    init_pred = mp.run_jax_at(ap_id, compute_init_pred, y_j)
+    logits0 = mp.run_jax_at(ap_id, lambda p, m=y_j.shape[0]: p * jnp.ones(m), init_pred)
+    GH = mp.run_jax_at(ap_id, compute_gh, y_j, logits0)
+
+    fxp_scale = 1 << DEFAULT_FXP_BITS
+    Q = mp.run_jax_at(ap_id, quantize_gh, GH, fxp_scale)
+    qg = mp.run_jax_at(ap_id, lambda a: a[:, 0].astype(jnp.int64), Q)
+    qh = mp.run_jax_at(ap_id, lambda a: a[:, 1].astype(jnp.int64), Q)
+
+    priv_ctx, pub_ctx, _ = mp.run_at(ap_id, fhe.keygen, scheme="BFV")
+
+    # Prepare ciphertext(s)
+    g_ct = None  # type: ignore[assignment]
+    h_ct = None  # type: ignore[assignment]
+    gh_ct = None  # type: ignore[assignment]
+    if mode in ("classic", "classic_cached"):
+        g_ct = mp.run_at(ap_id, fhe.encrypt, qg, pub_ctx)
+        h_ct = mp.run_at(ap_id, fhe.encrypt, qh, pub_ctx)
+    elif mode in (
+        "interleaved",
+        "interleaved_cached",
+        "interleaved_fused",
+        "interleaved_fused_cached",
+    ):
+        # Interleave qg and qh into one vector: [g0,h0,g1,h1,...]
+        def _interleave(a, b):
+            m = a.shape[0]
+            out = jnp.empty((m * 2,), dtype=jnp.int64)
+            out = out.at[0::2].set(a)
+            out = out.at[1::2].set(b)
+            return out
+
+        qgh = mp.run_jax_at(ap_id, _interleave, qg, qh)
+        gh_ct = mp.run_at(ap_id, fhe.encrypt, qgh, pub_ctx)
+    else:
+        raise ValueError(f"Unknown mode: {mode}")
+    rng = mp.run_jax_at(
+        ap_id,
+        lambda m: jnp.array(
+            np.random.default_rng(0).integers(0, t, size=m), dtype=jnp.int64
+        ),
+        y_j.shape[0],
+    )
+
+    def mk_subgroup_map(bt_level, group_size):
+        group_indices = jnp.arange(group_size)[:, None]
+        return (group_indices == bt_level).astype(jnp.int8)
+
+    # Precompute subgroup maps per-PP once (rng fixed) and parity selectors once if needed
+    subgroup_maps = []
+    for pp in pp_ids:
+        pub_ctx_pp = mp.p2p(ap_id, pp, pub_ctx)
+        rng_pp = mp.p2p(ap_id, pp, rng)
+        subgroup_map_pp = mp.run_jax_at(pp, mk_subgroup_map, rng_pp, t)
+        subgroup_maps.append((pp, pub_ctx_pp, subgroup_map_pp))
+
+    even_sel = None
+    odd_sel = None
+    if mode in (
+        "interleaved",
+        "interleaved_cached",
+    ):
+
+        def _build_parity_selectors(m_samples):
+            n = m_samples * 2
+            even = jnp.zeros((n,), dtype=jnp.int64).at[0::2].set(1)
+            odd = jnp.zeros((n,), dtype=jnp.int64).at[1::2].set(1)
+            return even, odd
+
+        even_sel, odd_sel = mp.run_jax_at(ap_id, _build_parity_selectors, y_j.shape[0])
+
+    # Optional: precompute and encrypt all bucket masks per-PP for cached modes
+    cached_masks = None
+    pre_dt = mp.run_jax_at(ap_id, lambda: jnp.array(0.0, dtype=jnp.float64))
+    if mode in ("interleaved_cached", "classic_cached"):
+        # Helper function to duplicate mask to interleaved length (used only for interleaved mode)
+        def _dup2(mask):
+            n = mask.shape[0]
+            out = jnp.empty((n * 2,), dtype=jnp.int64)
+            out = out.at[0::2].set(mask)
+            out = out.at[1::2].set(mask)
+            return out
+
+        use_interleave = mode == "interleaved_cached"
+
+        cached_masks = []  # list per PP: [ [list per group: [list per feature: [mask_ct per bucket]]] ]
+        tpre0 = mp.run_jax_at(ap_id, lambda: jnp.array(time.time(), dtype=jnp.float64))
+        for i, (pp, pub_ctx_pp, subgroup_map_pp) in enumerate(subgroup_maps):
+            feature_size = pp_idx[i].shape[1]
+            grp_masks = []
+            for grp in range(t):
+                gom = mp.run_jax_at(pp, lambda m, idx: m[idx], subgroup_map_pp, grp)
+
+                def create_masked_order_map(m, om):
+                    mask_expanded = jnp.expand_dims(m, axis=1)
+                    mask_full = jnp.broadcast_to(mask_expanded, om.shape)
+                    return jnp.where(mask_full == 1, om, -1)
+
+                gom_map = mp.run_jax_at(pp, create_masked_order_map, gom, pp_idx[i])
+
+                feat_masks = []
+                for feature_idx in range(feature_size):
+                    # Build all bucket masks at once: (k, M)
+                    def build_bucket_masks(gom_, f_idx, num_buckets):
+                        def mask_for_b(b_idx, gom_i, f_i):
+                            fb = gom_i[:, f_i]
+                            valid_and_in_bucket = (fb >= 0) & (fb <= b_idx)
+                            return valid_and_in_bucket.astype(jnp.int64)
+
+                        bs = jnp.arange(num_buckets, dtype=jnp.int64)
+                        return jax.vmap(mask_for_b, in_axes=(0, None, None))(
+                            bs, gom_, f_idx
+                        )
+
+                    bucket_masks = mp.run_jax_at(
+                        pp, build_bucket_masks, gom_map, feature_idx, k
+                    )
+                    # Encrypt each bucket mask (with optional duplication for interleaved mode)
+                    masks_ct = []
+                    for b in range(k):
+                        row_b = mp.run_jax_at(pp, lambda M, bi: M[bi], bucket_masks, b)
+                        # Apply _dup2 transformation only for interleaved mode
+                        if use_interleave:
+                            mask_to_encrypt = mp.run_jax_at(pp, _dup2, row_b)
+                        else:
+                            mask_to_encrypt = row_b
+                        mask_ct_pp = mp.run_at(
+                            pp, fhe.encrypt, mask_to_encrypt, pub_ctx_pp
+                        )
+                        mask_ct_ap = mp.p2p(pp, ap_id, mask_ct_pp)
+                        masks_ct.append(mask_ct_ap)
+                    feat_masks.append(masks_ct)
+                grp_masks.append(feat_masks)
+            cached_masks.append(grp_masks)
+        tpre1 = mp.run_jax_at(ap_id, lambda: jnp.array(time.time(), dtype=jnp.float64))
+        pre_dt = mp.run_jax_at(ap_id, lambda a, b: a - b, tpre1, tpre0)
+
+    # Run reps and time compute + decrypt assembly across all PPs
+    times_total = []
+    times_comp = []
+    times_dec = []
+    for rep_i in range(reps):
+        t0 = mp.run_jax_at(ap_id, lambda: jnp.array(time.time(), dtype=jnp.float64))
+        comp_parts = []
+        dec_parts = []
+        for i, (pp, pub_ctx_pp, subgroup_map_pp) in enumerate(subgroup_maps):
+            if mode == "classic":
+                assert g_ct is not None and h_ct is not None
+                tcomp0 = mp.run_jax_at(
+                    ap_id, lambda: jnp.array(time.time(), dtype=jnp.float64)
+                )
+                g_lists, h_lists = batch_feature_wise_bucket_sum_fhe_vector(
+                    g_ct,
+                    h_ct,
+                    subgroup_map_pp,
+                    pp_idx[i],
+                    k,
+                    t,
+                    rank=pp,
+                    ap_rank=ap_id,
+                )
+                tcomp1 = mp.run_jax_at(
+                    ap_id, lambda: jnp.array(time.time(), dtype=jnp.float64)
+                )
+                comp_parts.append(
+                    mp.run_jax_at(ap_id, lambda a, b: a - b, tcomp1, tcomp0)
+                )
+                tdec0 = mp.run_jax_at(
+                    ap_id, lambda: jnp.array(time.time(), dtype=jnp.float64)
+                )
+                for grp in range(t):
+                    enc_g_list = g_lists[grp]
+                    enc_h_list = h_lists[grp]
+                    dec_g = [
+                        mp.run_at(ap_id, fhe.decrypt, ct, priv_ctx) for ct in enc_g_list
+                    ]
+                    dec_h = [
+                        mp.run_at(ap_id, fhe.decrypt, ct, priv_ctx) for ct in enc_h_list
+                    ]
+
+                    def _stack(*xs):
+                        return jnp.stack(xs)
+
+                    _ = mp.run_jax_at(ap_id, _stack, *dec_g)
+                    _ = mp.run_jax_at(ap_id, _stack, *dec_h)
+                tdec1 = mp.run_jax_at(
+                    ap_id, lambda: jnp.array(time.time(), dtype=jnp.float64)
+                )
+                dec_parts.append(mp.run_jax_at(ap_id, lambda a, b: a - b, tdec1, tdec0))
+            elif mode == "classic_cached":
+                assert g_ct is not None and h_ct is not None
+                assert cached_masks is not None
+                feature_size = pp_idx[i].shape[1]
+                g_lists = [[] for _ in range(t)]
+                h_lists = [[] for _ in range(t)]
+                tcomp0 = mp.run_jax_at(
+                    ap_id, lambda: jnp.array(time.time(), dtype=jnp.float64)
+                )
+                for grp in range(t):
+                    feat_masks = cached_masks[i][grp]
+                    for feature_idx in range(feature_size):
+                        for bucket_idx in range(k):
+                            mask_ct_ap = feat_masks[feature_idx][bucket_idx]
+                            g_sum_ct = mp.run_at(ap_id, fhe.dot, g_ct, mask_ct_ap)
+                            h_sum_ct = mp.run_at(ap_id, fhe.dot, h_ct, mask_ct_ap)
+                            g_lists[grp].append(g_sum_ct)
+                            h_lists[grp].append(h_sum_ct)
+                tcomp1 = mp.run_jax_at(
+                    ap_id, lambda: jnp.array(time.time(), dtype=jnp.float64)
+                )
+                comp_parts.append(
+                    mp.run_jax_at(ap_id, lambda a, b: a - b, tcomp1, tcomp0)
+                )
+                tdec0 = mp.run_jax_at(
+                    ap_id, lambda: jnp.array(time.time(), dtype=jnp.float64)
+                )
+                for grp in range(t):
+                    enc_g_list = g_lists[grp]
+                    enc_h_list = h_lists[grp]
+                    dec_g = [
+                        mp.run_at(ap_id, fhe.decrypt, ct, priv_ctx) for ct in enc_g_list
+                    ]
+                    dec_h = [
+                        mp.run_at(ap_id, fhe.decrypt, ct, priv_ctx) for ct in enc_h_list
+                    ]
+
+                    def _stack(*xs):
+                        return jnp.stack(xs)
+
+                    _ = mp.run_jax_at(ap_id, _stack, *dec_g)
+                    _ = mp.run_jax_at(ap_id, _stack, *dec_h)
+                tdec1 = mp.run_jax_at(
+                    ap_id, lambda: jnp.array(time.time(), dtype=jnp.float64)
+                )
+                dec_parts.append(mp.run_jax_at(ap_id, lambda a, b: a - b, tdec1, tdec0))
+            else:
+                assert gh_ct is not None
+                # even_sel/odd_sel were built once before reps
+
+                def _dup2(mask):
+                    n = mask.shape[0]
+                    out = jnp.empty((n * 2,), dtype=jnp.int64)
+                    out = out.at[0::2].set(mask)
+                    out = out.at[1::2].set(mask)
+                    return out
+
+                feature_size = pp_idx[i].shape[1]
+                g_lists = [[] for _ in range(t)]
+                h_lists = [[] for _ in range(t)]
+                tcomp0 = mp.run_jax_at(
+                    ap_id, lambda: jnp.array(time.time(), dtype=jnp.float64)
+                )
+                for grp in range(t):
+                    if mode in ("interleaved_cached",) and cached_masks is not None:
+                        # Use precomputed encrypted masks
+                        feat_masks = cached_masks[i][grp]
+                        for feature_idx in range(feature_size):
+                            for bucket_idx in range(k):
+                                mask_ct_ap = feat_masks[feature_idx][bucket_idx]
+                                prod_ct = mp.run_at(ap_id, fhe.mul, gh_ct, mask_ct_ap)
+                                g_sum_ct = mp.run_at(ap_id, fhe.dot, prod_ct, even_sel)
+                                h_sum_ct = mp.run_at(ap_id, fhe.dot, prod_ct, odd_sel)
+                                g_lists[grp].append(g_sum_ct)
+                                h_lists[grp].append(h_sum_ct)
+                    else:
+                        # Build on the fly
+                        gom = mp.run_jax_at(
+                            pp, lambda m, idx: m[idx], subgroup_map_pp, grp
+                        )
+
+                        def create_masked_order_map(m, om):
+                            mask_expanded = jnp.expand_dims(m, axis=1)
+                            mask_full = jnp.broadcast_to(mask_expanded, om.shape)
+                            return jnp.where(mask_full == 1, om, -1)
+
+                        gom_map = mp.run_jax_at(
+                            pp, create_masked_order_map, gom, pp_idx[i]
+                        )
+                        for feature_idx in range(feature_size):
+                            for bucket_idx in range(k):
+
+                                def create_bucket_mask(gom_, f_idx, b_idx):
+                                    fb = gom_[:, f_idx]
+                                    valid_and_in_bucket = (fb >= 0) & (fb <= b_idx)
+                                    return valid_and_in_bucket.astype(jnp.int64)
+
+                                bucket_mask = mp.run_jax_at(
+                                    pp,
+                                    create_bucket_mask,
+                                    gom_map,
+                                    feature_idx,
+                                    bucket_idx,
+                                )
+                                inter_mask = mp.run_jax_at(pp, _dup2, bucket_mask)
+                                mask_ct_pp = mp.run_at(
+                                    pp, fhe.encrypt, inter_mask, pub_ctx_pp
+                                )
+                                mask_ct_ap = mp.p2p(pp, ap_id, mask_ct_pp)
+
+                                prod_ct = mp.run_at(ap_id, fhe.mul, gh_ct, mask_ct_ap)
+                                g_sum_ct = mp.run_at(ap_id, fhe.dot, prod_ct, even_sel)
+                                h_sum_ct = mp.run_at(ap_id, fhe.dot, prod_ct, odd_sel)
+
+                                g_lists[grp].append(g_sum_ct)
+                                h_lists[grp].append(h_sum_ct)
+                tcomp1 = mp.run_jax_at(
+                    ap_id, lambda: jnp.array(time.time(), dtype=jnp.float64)
+                )
+                comp_parts.append(
+                    mp.run_jax_at(ap_id, lambda a, b: a - b, tcomp1, tcomp0)
+                )
+                tdec0 = mp.run_jax_at(
+                    ap_id, lambda: jnp.array(time.time(), dtype=jnp.float64)
+                )
+                for grp in range(t):
+                    enc_g_list = g_lists[grp]
+                    enc_h_list = h_lists[grp]
+                    dec_g = [
+                        mp.run_at(ap_id, fhe.decrypt, ct, priv_ctx) for ct in enc_g_list
+                    ]
+                    dec_h = [
+                        mp.run_at(ap_id, fhe.decrypt, ct, priv_ctx) for ct in enc_h_list
+                    ]
+
+                    def _stack(*xs):
+                        return jnp.stack(xs)
+
+                    _ = mp.run_jax_at(ap_id, _stack, *dec_g)
+                    _ = mp.run_jax_at(ap_id, _stack, *dec_h)
+                tdec1 = mp.run_jax_at(
+                    ap_id, lambda: jnp.array(time.time(), dtype=jnp.float64)
+                )
+                dec_parts.append(mp.run_jax_at(ap_id, lambda a, b: a - b, tdec1, tdec0))
+
+        t1 = mp.run_jax_at(ap_id, lambda: jnp.array(time.time(), dtype=jnp.float64))
+        dt = mp.run_jax_at(ap_id, lambda a, b: a - b, t1, t0)
+        # Optionally include precompute cost once (first repetition only)
+        if include_precompute and rep_i == 0:
+            dt = mp.run_jax_at(ap_id, lambda x, add: x + add, dt, pre_dt)
+        # Sum parts across PPs
+
+        def _sum_vec(*xs):
+            s = xs[0]
+            for x in xs[1:]:
+                s = s + x
+            return s
+
+        comp_sum = (
+            mp.run_jax_at(ap_id, _sum_vec, *comp_parts)
+            if comp_parts
+            else mp.run_jax_at(ap_id, lambda: jnp.array(0.0, dtype=jnp.float64))
+        )
+        dec_sum = (
+            mp.run_jax_at(ap_id, _sum_vec, *dec_parts)
+            if dec_parts
+            else mp.run_jax_at(ap_id, lambda: jnp.array(0.0, dtype=jnp.float64))
+        )
+        times_total.append(dt)
+        times_comp.append(comp_sum)
+        times_dec.append(dec_sum)
+
+    # Stack per-rep durations into a vector at AP for robust fetch
+    def _stack_times(*xs):
+        return jnp.stack(xs)
+
+    total_vec = mp.run_jax_at(ap_id, _stack_times, *times_total)
+    if not breakdown:
+        return total_vec
+    comp_vec = mp.run_jax_at(ap_id, _stack_times, *times_comp)
+    dec_vec = mp.run_jax_at(ap_id, _stack_times, *times_dec)
+
+    def _stack3(a, b, c):
+        return jnp.stack([a, b, c], axis=0)
+
+    return mp.run_jax_at(ap_id, _stack3, total_vec, comp_vec, dec_vec)
+
+
+def main():
+    parser = argparse.ArgumentParser(description="FHE histogram microbenchmark")
+    parser.add_argument(
+        "--world-size", type=int, default=2, help="Total parties (AP=1+PPs)"
+    )
+    parser.add_argument("--m", type=int, default=4096, help="Samples")
+    parser.add_argument("--n-total", type=int, default=16, help="Total features")
+    parser.add_argument("--n-ap", type=int, default=4, help="AP feature count")
+    parser.add_argument("--k", type=int, default=16, help="Bins per feature")
+    parser.add_argument("--t", type=int, default=4, help="Groups (nodes at level)")
+    parser.add_argument("--reps", type=int, default=3, help="Repetitions")
+    parser.add_argument("--seed", type=int, default=0, help="Random seed")
+    parser.add_argument(
+        "--mode",
+        type=str,
+        default="classic",
+        choices=[
+            "classic",
+            "classic_cached",
+            "interleaved",
+            "interleaved_cached",
+        ],
+        help="Histogram mode to benchmark",
+    )
+    parser.add_argument(
+        "--include-precompute",
+        action="store_true",
+        help="Include precompute (mask generation+encryption) time in the first repetition",
+    )
+    parser.add_argument(
+        "--breakdown",
+        action="store_true",
+        help="Report timing breakdown (total, compute, decrypt)",
+    )
+    args = parser.parse_args()
+
+    assert args.world_size >= 2, "world-size must be >= 2"
+    pp_parties = args.world_size - 1
+    assert args.n_total >= args.n_ap
+
+    # Split PP features evenly
+    n_pp_total = args.n_total - args.n_ap
+    n_pp_each = [n_pp_total // pp_parties] * pp_parties
+    n_pp_each[-1] += n_pp_total - sum(n_pp_each)
+
+    X_ap, X_pp_all, y = _gen_data(args.m, args.n_total, args.n_ap, args.seed)
+    offset = 0
+    X_pp_splits = []
+    for c in n_pp_each:
+        X_pp_splits.append(X_pp_all[:, offset : offset + c])
+        offset += c
+
+    sim = mp.Simulator.simple(args.world_size)
+
+    ap_id = 0
+    pp_ids = list(range(1, args.world_size))
+
+    print("\n=== FHE Histogram Microbenchmark ===")
+    print(
+        f"world-size={args.world_size} (AP+{pp_parties} PPs), m={args.m}, n_total={args.n_total} (AP={args.n_ap}, PP={n_pp_total}), k={args.k}, t={args.t}, reps={args.reps}, mode={args.mode}"
+    )
+
+    out = mp.evaluate(
+        sim,
+        _bench_once,
+        ap_id,
+        pp_ids,
+        X_ap,
+        X_pp_splits,
+        y,
+        args.k,
+        args.t,
+        args.reps,
+        args.mode,
+        args.include_precompute,
+        args.breakdown,
+    )
+    times_raw = mp.fetch(sim, out)
+
+    # Expected: [times_at_ap, None, ...] in 2PC; extract first non-None
+    if isinstance(times_raw, list) and len(times_raw) >= 1 and times_raw[-1] is None:
+        times_nodes = times_raw[0]
+    else:
+        times_nodes = times_raw
+
+    if args.breakdown:
+        times_arr = np.asarray(times_nodes, dtype=float)
+        # Expect shape (3, reps): [total, compute, decrypt]
+        if times_arr.ndim == 1:
+            # Fallback if flattened; try to split into 3 roughly equal parts
+            n = times_arr.size
+            r = n // 3
+            total = times_arr[:r]
+            comp = times_arr[r : 2 * r]
+            dec = times_arr[2 * r : 3 * r]
+        else:
+            total, comp, dec = (
+                times_arr[0].ravel(),
+                times_arr[1].ravel(),
+                times_arr[2].ravel(),
+            )
+
+        print(f"Per-rep total (s): {', '.join(f'{t:.4f}' for t in total.tolist())}")
+        print(
+            f"Per-rep compute-only (s): {', '.join(f'{t:.4f}' for t in comp.tolist())}"
+        )
+        print(
+            f"Per-rep decrypt-only (s): {', '.join(f'{t:.4f}' for t in dec.tolist())}"
+        )
+        print(
+            f"Averages — total: {float(total.mean()):.4f}s, compute: {float(comp.mean()):.4f}s, decrypt: {float(dec.mean()):.4f}s"
+        )
+    else:
+        # Convert to numpy array of floats (handle scalar, list, or numpy array)
+        if isinstance(times_nodes, list):
+            # elements are likely [val, None] pairs; take first
+            times_arr = np.array(
+                [
+                    float(np.array(e[0]))
+                    if isinstance(e, (list, tuple))
+                    else float(np.array(e))
+                    for e in times_nodes
+                ],
+                dtype=float,
+            )
+        else:
+            times_arr = np.asarray(times_nodes, dtype=float).ravel()
+        avg = float(times_arr.mean())
+        print(f"Per-rep time (s): {', '.join(f'{t:.4f}' for t in times_arr.tolist())}")
+        print(f"Average time (s): {avg:.4f}")
+
+
+if __name__ == "__main__":
+    main()