feat(phyai): humming quant kernels

.gitignore

@@ -102,6 +102,7 @@ Thumbs.db
 # claude & codex
 .claude/settings.local.json
 .agents/
+.mimocode/
 skills-lock.json
 
 # profile
@@ -110,3 +111,6 @@ skills-lock.json
 
 # For some OpenCode agents
 docs/compose/
+
+# Auto Tune and Auto Research Agent
+.auto_research/

CLAUDE.md

@@ -14,7 +14,7 @@ This file provides guidance to Claude Code (claude.ai/code) when working with co
 | `phyai-model-optimizer/` | Placeholder, no source yet. |
 | `phyai-utils-tools/` | Placeholder, no source yet. |
 
-Hard pins: `torch==2.11`, `flashinfer-python==v0.6.12`, `transformers==5.8.1`. Don't bump these casually — green-context and flashinfer behaviour is tied to these versions.
+Hard pins: `torch==2.11`, `flashinfer-python==v0.6.14`, `transformers==5.8.1`. Don't bump these casually — green-context and flashinfer behaviour is tied to these versions.
 
 ## Common commands
 
@@ -52,6 +52,12 @@ CPU is the default for tests — `phyai/tests/conftest.py` autouses a fixture th
 - C/C++: clang-format from `.clang-format` (column 128, 2-space indent, `PointerAlignment: Left`). clang-tidy from `.clang-tidy` (google + modernize + performance, `WarningsAsErrors: '*'`, identifier naming enforced: classes `CamelCase`, variables `lower_case`, globals `UPPER_CASE`). C++20 (`add_compile_options(-std=c++20)` in `phyai-ext/CMakeLists.txt`).
 - Python: `ruff-format` via pre-commit. Public-by-default — `_` prefix only for genuine implementation details. Singletons exposed via `get_*()` getters, not module-level instances.
 
+## Code Comment conventions
+
+All comments should be self-contained. Do not explain how you did something or explain what this code block did; just explain whywe did this. And, also adds your name for all of the comments, like # note(foo). If you don't know user's name, just ask them.
+
+Comments should be concise. If not explain why, ust delete it. In this sense, most of the AI comments should be removed.
+
 ## More MUST FOLLOW Conventions provided by human
 
 - all log function in phyai package should use phyai.utils' logging api. U judge using `this_rank_log` or `all_rank_log`

benchmark/bench_n_batch.py

@@ -78,6 +78,11 @@ class BenchSpec:
     name: str
     step_callable: Callable[[], Any]
     teardown_callable: Callable[[], None]
+    sample_count: int | None = None
+
+    def __post_init__(self) -> None:
+        if self.sample_count is not None and self.sample_count <= 0:
+            raise ValueError(f"sample_count must be positive, got {self.sample_count}")
 
 
 @dataclass
@@ -201,7 +206,14 @@ def run(self) -> list[BenchResult]:
             spec = self.setup_fn(bs)
             try:
                 extras = self.extras_fn(bs, spec) if self.extras_fn else {}
-                result = self._run_one(spec, batch_size=bs, extras=extras)
+                result = self._run_one(
+                    spec,
+                    batch_size=bs,
+                    throughput_batch_size=(
+                        bs if spec.sample_count is None else spec.sample_count
+                    ),
+                    extras=extras,
+                )
                 results.append(result)
                 _print_result(result)
             finally:
@@ -226,6 +238,7 @@ def _run_one(
         spec: BenchSpec,
         *,
         batch_size: int,
+        throughput_batch_size: int,
         extras: dict[str, Any],
     ) -> BenchResult:
         step_callable = spec.step_callable
@@ -272,6 +285,7 @@ def _run_one(
             bench_name=self.bench_name,
             spec_name=spec.name,
             batch_size=batch_size,
+            throughput_batch_size=throughput_batch_size,
             n_warmup=self.n_warmup,
             n_timed=self.n_timed,
             latencies_ms=latencies_ms,
@@ -319,6 +333,7 @@ def _build_result(
     bench_name: str,
     spec_name: str,
     batch_size: int,
+    throughput_batch_size: int,
     n_warmup: int,
     n_timed: int,
     latencies_ms: list[float],
@@ -333,8 +348,8 @@ def _build_result(
     stdev = float(statistics.stdev(latencies_ms)) if len(latencies_ms) > 1 else 0.0
     mn = float(arr.min())
     mx = float(arr.max())
-    # Throughput in samples/s = batch_size / mean_latency_seconds.
-    throughput = batch_size / (mean / 1000.0) if mean > 0 else float("inf")
+    # note(chenghua): Throughput uses the actual request size, which can be below max batch size.
+    throughput = throughput_batch_size / (mean / 1000.0) if mean > 0 else float("inf")
     merged_extras = {**extras, "spec_name": spec_name}
     return BenchResult(
         run_name=run_name,

benchmark/bench_n_batch_ws1_pi05.py

@@ -1,7 +1,7 @@
 """End-to-end pi0.5 ws1 (single-card) latency benchmark, swept over batch sizes.
 
 Builds the pi0.5 engine via the standard plugin path once per batch
-size, feeds a dummy request (random pixels + one-token "prompt"), and
+size, feeds a deterministic dummy request, and
 hands it to the generic :class:`NBatchBenchRunner` from
 :mod:`bench_n_batch` for timing + optional Nsight Systems / Perfetto
 profile capture.
@@ -49,6 +49,7 @@
 from __future__ import annotations
 
 import argparse
+import json
 from pathlib import Path
 from typing import Any
 
@@ -85,8 +86,15 @@ def make_dummy_request(
     plugin_cfg: PI05Config,
     device: torch.device,
     dtype: torch.dtype,
+    lang_len: int = 1,
+    seed: int = 0,
 ) -> PI05Request:
-    """Random pixels + single-token prompt PI05Request for ``batch_size`` robots."""
+    """Deterministic random request for ``batch_size`` robots."""
+    if not 1 <= lang_len <= plugin_cfg.tokenizer_max_length:
+        raise ValueError(
+            f"lang_len must be in [1, {plugin_cfg.tokenizer_max_length}], got {lang_len}"
+        )
+    generator = torch.Generator(device=device).manual_seed(seed)
     pixel_values = torch.rand(
         batch_size,
         num_images,
@@ -95,12 +103,13 @@ def make_dummy_request(
         plugin_cfg.vision.image_size,
         dtype=dtype,
         device=device,
+        generator=generator,
     )
     input_ids = torch.zeros(
         batch_size, plugin_cfg.tokenizer_max_length, dtype=torch.int64, device=device
     )
-    input_ids[:, 0] = 2  # any non-pad token id
-    lang_lens = torch.ones(batch_size, dtype=torch.int64, device=device)
+    input_ids[:, :lang_len] = 2  # any non-pad token id
+    lang_lens = torch.full((batch_size,), lang_len, dtype=torch.int64, device=device)
     return PI05Request(
         pixel_values=pixel_values,
         input_ids=input_ids,
@@ -116,6 +125,9 @@ def make_setup_fn(
     use_cuda_graph: bool,
     num_images: int = 3,
     vision_params_dtype: torch.dtype | None = None,
+    actual_batch_size: int | None = None,
+    lang_len: int = 1,
+    seed: int = 0,
 ):
     """Build the per-batch-size ``setup_fn`` closure for :class:`NBatchBenchRunner`.
 
@@ -131,6 +143,12 @@ def make_setup_fn(
     ]
 
     def setup_fn(batch_size: int) -> bnb.BenchSpec:
+        actual_B = batch_size if actual_batch_size is None else actual_batch_size
+        if not 1 <= actual_B <= batch_size:
+            raise ValueError(
+                f"actual_batch_size must be in [1, max_batch_size={batch_size}], "
+                f"got {actual_B}"
+            )
         engine = Engine(
             EngineArgs(
                 plugin="pi05",
@@ -147,16 +165,19 @@ def setup_fn(batch_size: int) -> bnb.BenchSpec:
             )
         )
         request = make_dummy_request(
-            batch_size=batch_size,
+            batch_size=actual_B,
             num_images=num_images,
             plugin_cfg=plugin_cfg,
             device=device,
             dtype=dtype,
+            lang_len=lang_len,
+            seed=seed,
         )
         return bnb.BenchSpec(
             name="ws1_pi05",
             step_callable=lambda: engine.step(request),
             teardown_callable=engine.close,
+            sample_count=actual_B,
         )
 
     return setup_fn
@@ -167,20 +188,73 @@ def make_extras_fn(
     dtype_name: str,
     device_target: str,
     use_cuda_graph: bool,
+    plugin_cfg: PI05Config,
+    actual_batch_size: int | None,
+    lang_len: int,
+    seed: int,
+    quantization: dict[str, Any],
 ):
     def extras_fn(batch_size: int, spec: bnb.BenchSpec) -> dict[str, Any]:
+        actual_B = batch_size if actual_batch_size is None else actual_batch_size
+        lang_buckets = sorted(
+            {b for b in (16, 48, 112) if b < plugin_cfg.tokenizer_max_length}
+            | {plugin_cfg.tokenizer_max_length}
+        )
+        lang_bucket = next(
+            (bucket for bucket in lang_buckets if bucket >= lang_len),
+            plugin_cfg.tokenizer_max_length,
+        )
         return {
             "model": "pi05",
             "scheduler": "ws1",
             "dtype": dtype_name,
             "device": device_target,
             "use_cuda_graph": use_cuda_graph,
             "max_batch_size": batch_size,
+            "actual_batch_size": actual_B,
+            "lang_len": lang_len,
+            "lang_bucket": lang_bucket,
+            "chunk_size": plugin_cfg.chunk_size,
+            "seed": seed,
+            "quantization": quantization,
         }
 
     return extras_fn
 
 
+def checkpoint_quantization_metadata(checkpoint: Path) -> dict[str, Any]:
+    """Return portable quantization metadata without recording host paths."""
+    config = json.loads((checkpoint / "config.json").read_text(encoding="utf-8"))
+    quant = config.get("quantization_config") or {}
+    groups = []
+    for name, group in sorted((quant.get("config_groups") or {}).items()):
+        weight = group.get("weights") or {}
+        activation = group.get("input_activations") or {}
+        groups.append(
+            {
+                "name": name,
+                "targets": len(group.get("targets") or []),
+                "weight_num_bits": weight.get("num_bits"),
+                "weight_type": weight.get("type"),
+                "weight_humming_dtype": weight.get("humming_dtype"),
+                "weight_strategy": weight.get("strategy"),
+                "weight_block_structure": weight.get("block_structure"),
+                "activation_num_bits": activation.get("num_bits"),
+                "activation_type": activation.get("type"),
+                "activation_strategy": activation.get("strategy"),
+                "activation_group_size": activation.get("group_size"),
+                "activation_dynamic": activation.get("dynamic"),
+            }
+        )
+    return {
+        "quant_method": quant.get("quant_method"),
+        "format": quant.get("format"),
+        "pack_format": quant.get("pack_format"),
+        "status": quant.get("quantization_status"),
+        "groups": groups,
+    }
+
+
 def main() -> None:
     parser = argparse.ArgumentParser(
         description=__doc__,
@@ -219,6 +293,19 @@ def main() -> None:
         default=3,
         help="Number of cameras per robot (default 3, the pi05_base contract).",
     )
+    parser.add_argument(
+        "--actual-batch-size",
+        type=int,
+        default=None,
+        help="Request batch size; defaults to each swept max batch size.",
+    )
+    parser.add_argument(
+        "--lang-len",
+        type=int,
+        default=1,
+        help="Real prompt length before scheduler bucket padding (default 1).",
+    )
+    parser.add_argument("--seed", type=int, default=0)
     parser.add_argument(
         "--vision-dtype",
         choices=("bfloat16", "float32"),
@@ -241,6 +328,8 @@ def main() -> None:
 
     dtype = _DTYPES[args.dtype]
     use_cuda_graph = not args.no_cuda_graph and args.device == "cuda"
+    plugin_cfg = load_config(args.checkpoint, PI05Config)
+    quantization = checkpoint_quantization_metadata(args.checkpoint)
 
     # Install whatever profiler the CLI requested. NoOp is the default
     # when --profile-backend is "none" (or rank is excluded).
@@ -254,11 +343,19 @@ def main() -> None:
         use_cuda_graph=use_cuda_graph,
         num_images=args.num_images,
         vision_params_dtype=(torch.float32 if args.vision_dtype == "float32" else None),
+        actual_batch_size=args.actual_batch_size,
+        lang_len=args.lang_len,
+        seed=args.seed,
     )
     extras_fn = make_extras_fn(
         dtype_name=args.dtype,
         device_target=args.device,
         use_cuda_graph=use_cuda_graph,
+        plugin_cfg=plugin_cfg,
+        actual_batch_size=args.actual_batch_size,
+        lang_len=args.lang_len,
+        seed=args.seed,
+        quantization=quantization,
     )
 
     runner = bnb.NBatchBenchRunner(

benchmark/pi05/model_flops.py

@@ -143,10 +143,10 @@ def stage_flops(
     )
     vision += gemm_flop(image_tokens, dims.l_hidden, dims.v_hidden)
 
-    # --- LLM prefix: n_per_sample tokens, full self-attention, 18 layers. ---
+    # note(chenghua): The final language layer executes only its input norm and K/V projection.
     n_ps = n_per_sample(lang_len, dims, num_images)
     llm_prefix = 0.0
-    for _ in range(dims.l_layers):
+    for _ in range(max(dims.l_layers - 1, 0)):
         llm_prefix += transformer_layer_flop(
             n_ps,
             dims.l_hidden,
@@ -157,6 +157,9 @@ def stage_flops(
             kv_len=n_ps,
             gated_mlp=True,
         )
+    if dims.l_layers:
+        kv_dim = dims.l_kv_heads * dims.l_head_dim
+        llm_prefix += gemm_flop(n_ps, 2 * kv_dim, dims.l_hidden)
 
     # --- Expert one Euler step: chunk_size queries vs (prefix + suffix) kv. ---
     e_kv_len = n_ps + dims.chunk_size
@@ -208,14 +211,16 @@ def attn_mlp_params(hidden, heads, kv_heads, head_dim, inter, *, gated):
     v_params += dims.v_hidden * (dims.num_channels * dims.patch_size**2)  # patch embed
     v_params += dims.v_hidden * dims.l_hidden  # projector
 
-    l_params = dims.l_layers * attn_mlp_params(
+    l_params = max(dims.l_layers - 1, 0) * attn_mlp_params(
         dims.l_hidden,
         dims.l_heads,
         dims.l_kv_heads,
         dims.l_head_dim,
         dims.l_intermediate,
         gated=True,
     )
+    if dims.l_layers:
+        l_params += dims.l_hidden * (2 * dims.l_kv_heads * dims.l_head_dim)
 
     e_params = dims.e_layers * attn_mlp_params(
         dims.e_hidden,