#!/usr/bin/env python3 """Cached byte generation for two- and three-level nested Mamba checkpoints.""" from __future__ import annotations import argparse import base64 import gc import io import json import time from pathlib import Path from typing import Dict, List, Optional import torch from nested_inference_tools import generate_bytes, load_nested_checkpoint def parse_args() -> argparse.Namespace: parser = argparse.ArgumentParser(description=__doc__) parser.add_argument( "--checkpoint", default=".", help="Hugging Face model directory (default: current directory) or legacy last.pt.", ) prompt = parser.add_mutually_exclusive_group() prompt.add_argument("--prompt", default="", help="UTF-8 prompt text.") prompt.add_argument("--prompt-file", help="Read prompt bytes from this file.") parser.add_argument("--max-new-bytes", type=int, default=256) parser.add_argument("--temperature", type=float, default=0.8) parser.add_argument("--top-p", type=float, default=0.9) parser.add_argument("--top-k", type=int, default=0) parser.add_argument("--repeat-penalty", type=float, default=1.05) parser.add_argument("--repeat-window", type=int, default=256) parser.add_argument("--greedy", action="store_true") parser.add_argument("--seed", type=int, default=1234) parser.add_argument( "--precision", choices=["fp16", "bf16", "fp32"], default="bf16", help=( "Model weight/activation precision. BF16 is the safe default for " "Mamba-2's long cached rollouts; use FP16 only for checkpoints and " "GPUs verified to remain finite." ), ) parser.add_argument("--device", default="cuda:0", help="Single-device inference target.") parser.add_argument("--fine-device", default=None) parser.add_argument("--nested-devices", default=None) parser.add_argument("--tertiary-device", default=None) parser.add_argument("--use-saved-placement", action="store_true") parser.add_argument("--output", help="Optional raw-byte output file.") parser.add_argument("--stats-json", help="Optional JSON statistics output.") parser.add_argument( "--image-output-dir", help=( "Extract complete generated P6 PPM images, convert them to PNG, and " "write them to this directory. Images are also embedded in --html-output." ), ) parser.add_argument("--html-output", "--output-html", dest="html_output", help="Write a self-contained prompt/response report with toggleable L2/L3 influence coloring.") return parser.parse_args() def _ppm_token(data: bytes, position: int) -> tuple[bytes, int]: """Read one whitespace/comment-delimited PPM header token.""" size = len(data) while position < size: if data[position] in b" \t\r\n": position += 1 continue if data[position] == ord("#"): newline = data.find(b"\n", position) if newline < 0: raise ValueError("unterminated PPM header comment") position = newline + 1 continue break start = position while position < size and data[position] not in b" \t\r\n#": position += 1 if position == start: raise ValueError("missing PPM header token") return data[start:position], position def extract_ppm_images(data: bytes) -> tuple[List[Dict[str, object]], List[str]]: """Extract complete P6 RGB images from an arbitrary generated byte stream.""" try: from PIL import Image except ImportError as error: return [], [f"Pillow is required to render generated images: {error}"] images: List[Dict[str, object]] = [] warnings: List[str] = [] search_from = 0 while len(images) < 32: start = data.find(b"P6", search_from) if start < 0: break search_from = start + 2 if start > 0 and data[start - 1] not in b" \t\r\n>": continue if start + 2 >= len(data) or data[start + 2] not in b" \t\r\n": continue try: width_token, position = _ppm_token(data, start + 2) height_token, position = _ppm_token(data, position) maximum_token, position = _ppm_token(data, position) width, height, maximum = int(width_token), int(height_token), int(maximum_token) if width < 1 or height < 1 or width * height > 16_777_216: raise ValueError(f"unsafe dimensions {width}x{height}") if maximum != 255: raise ValueError(f"unsupported maximum channel value {maximum}; expected 255") if position >= len(data) or data[position] not in b" \t\r\n": raise ValueError("missing whitespace before PPM pixel payload") # The delimiter is one whitespace unit. Treat CRLF as one unit so # the first pixel is never shifted on Windows-produced headers. pixel_start = position + 1 if data[position] == ord("\r") and pixel_start < len(data) and data[pixel_start] == ord("\n"): pixel_start += 1 pixel_bytes = width * height * 3 pixel_end = pixel_start + pixel_bytes if pixel_end > len(data): warnings.append( f"Incomplete PPM at byte {start}: {width}x{height} needs " f"{pixel_bytes:,} RGB bytes, but only {max(0, len(data)-pixel_start):,} remain." ) continue image = Image.frombytes("RGB", (width, height), data[pixel_start:pixel_end]) encoded = io.BytesIO() image.save(encoded, format="PNG", optimize=True) images.append( { "width": width, "height": height, "start": start, "end": pixel_end, "png": encoded.getvalue(), } ) search_from = pixel_end except (TypeError, ValueError) as error: warnings.append(f"Invalid PPM candidate at byte {start}: {error}.") return images, warnings def image_report_payload(images: List[Dict[str, object]], warnings: List[str]) -> Dict[str, object]: return { "images": [ { "width": int(item["width"]), "height": int(item["height"]), "start": int(item["start"]), "end": int(item["end"]), "data_uri": "data:image/png;base64," + base64.b64encode(item["png"]).decode("ascii"), } for item in images ], "warnings": list(warnings), } def write_extracted_images( directory: Path, images: List[Dict[str, object]], *, prefix: str ) -> List[str]: directory.mkdir(parents=True, exist_ok=True) paths: List[str] = [] for index, item in enumerate(images, 1): path = directory / f"{prefix}_{index:03d}_{item['width']}x{item['height']}.png" path.write_bytes(item["png"]) paths.append(str(path)) return paths def attributed_utf8_segments( data: bytes, attributions: List[Dict[str, object]] ) -> List[Dict[str, object]]: """Group byte attribution into valid UTF-8 characters without losing data.""" segments: List[Dict[str, object]] = [] index = 0 while index < len(data): first = data[index] width = ( 1 if first < 0x80 else 2 if 0xC2 <= first <= 0xDF else 3 if 0xE0 <= first <= 0xEF else 4 if 0xF0 <= first <= 0xF4 else 1 ) chunk = data[index : index + width] try: text = chunk.decode("utf-8", "strict") except UnicodeDecodeError: width = 1 chunk = data[index : index + 1] text = chunk.decode("utf-8", "replace") values = attributions[index : index + width] count = max(1, len(values)) segments.append( { "text": text, "bytes": list(chunk), "level2_active": any(bool(item.get("level2_active")) for item in values), "level3_active": any(bool(item.get("level3_active")) for item in values), "level2_delta_logp": sum( float(item.get("level2_delta_logp", 0.0)) for item in values ) / count, "level3_delta_logp": sum( float(item.get("level3_delta_logp", 0.0)) for item in values ) / count, "hierarchy_delta_logp": sum( float(item.get("hierarchy_delta_logp", 0.0)) for item in values ) / count, } ) index += width return segments def make_generation_html( prompt: bytes, generated: bytes, attributions: List[Dict[str, object]], stats: Dict[str, object], comparisons: Optional[List[Dict[str, object]]] = None, image_report: Optional[Dict[str, object]] = None, ) -> str: segments = attributed_utf8_segments(generated, attributions) payload = json.dumps( { "prompt": prompt.decode("utf-8", "replace"), "segments": segments, "stats": stats, "comparisons": comparisons or [], "image_report": image_report or {"images": [], "warnings": []}, }, separators=(",", ":"), ).replace(" Nested Mamba generation influence

Nested Mamba generation influence

fine onlyL2 dominantL3 dominantdominant branch reduced selected-byte probability

Color is a decoder counterfactual, not merely an activation marker. Positive delta means the cached dynamic level increased the selected byte's log probability relative to replacing that level with its beginning-of-document state. Hover text for exact values.

PROMPT SHARED BY ALL ROLLOUTS

Level 1 only

Levels 1 + 2

Levels 1 + 2 + 3

Generated image rendering

Complete P6 RGB byte sequences found in each prompt-plus-rollout stream are validated and embedded here as PNG.

Level 1 only

Levels 1 + 2

Levels 1 + 2 + 3

""" def main() -> None: args = parse_args() if args.max_new_bytes < 0 or args.temperature <= 0 or not 0 < args.top_p <= 1: raise ValueError("max-new-bytes must be non-negative, temperature positive, and top-p in (0, 1]") if args.prompt_file: prompt = Path(args.prompt_file).expanduser().read_bytes() else: prompt = args.prompt.encode("utf-8") load_started = time.perf_counter() loaded = load_nested_checkpoint( args.checkpoint, precision=args.precision, device=args.device, fine_device=args.fine_device, nested_devices=args.nested_devices, tertiary_device=args.tertiary_device, use_saved_placement=args.use_saved_placement, ) torch.cuda.synchronize() load_seconds = time.perf_counter() - load_started generation_started = time.perf_counter() generated, stream = generate_bytes( loaded, prompt, max_new_bytes=args.max_new_bytes, temperature=args.temperature, top_p=args.top_p, top_k=args.top_k, repeat_penalty=args.repeat_penalty, repeat_window=args.repeat_window, greedy=args.greedy, seed=args.seed, collect_hierarchy_attribution=bool(args.html_output), ) torch.cuda.synchronize() generation_seconds = time.perf_counter() - generation_started combined = prompt + generated full_images, full_image_warnings = extract_ppm_images(combined) full_image_report = image_report_payload(full_images, full_image_warnings) print(combined.decode("utf-8", "replace")) attributions = list(stream.get("generated_attribution", [])) stats = { "checkpoint": str(loaded.checkpoint_path), "checkpoint_step": loaded.checkpoint_step, "trained_tokens": loaded.trained_tokens, "precision": loaded.precision, "model_parallel": loaded.model_parallel, "prompt_bytes": len(prompt), "generated_bytes": len(generated), "load_seconds": load_seconds, "generation_seconds_including_prefill": generation_seconds, "generated_bytes_per_second_including_prefill": ( len(generated) / generation_seconds if generation_seconds else 0.0 ), "fine_patches": int(stream["completed_patches"]), "level2_pools": int(stream["completed_nested_patches"]), "level3_pools": int(stream.get("completed_tertiary_patches", 0)), "hierarchy_attribution": bool(args.html_output), "rendered_images": len(full_images), "image_render_warnings": full_image_warnings, } if attributions: stats.update( { "attributed_bytes": len(attributions), "level2_active_generated_bytes": sum( bool(item.get("level2_active")) for item in attributions ), "level3_active_generated_bytes": sum( bool(item.get("level3_active")) for item in attributions ), "mean_level2_delta_logp": sum( float(item.get("level2_delta_logp", 0.0)) for item in attributions ) / len(attributions), "mean_level3_delta_logp": sum( float(item.get("level3_delta_logp", 0.0)) for item in attributions ) / len(attributions), "mean_hierarchy_delta_logp": sum( float(item.get("hierarchy_delta_logp", 0.0)) for item in attributions ) / len(attributions), } ) comparisons: List[Dict[str, object]] = [] del stream gc.collect() if args.html_output: for mode, label in (("level1", "Level 1 only"), ("level12", "Levels 1 + 2")): comparison_started = time.perf_counter() comparison_bytes, comparison_stream = generate_bytes( loaded, prompt, max_new_bytes=args.max_new_bytes, temperature=args.temperature, top_p=args.top_p, top_k=args.top_k, repeat_penalty=args.repeat_penalty, repeat_window=args.repeat_window, greedy=args.greedy, seed=args.seed, hierarchy_mode=mode, ) torch.cuda.synchronize() comparison_seconds = time.perf_counter() - comparison_started comparisons.append( { "mode": mode, "label": label, "text": comparison_bytes.decode("utf-8", "replace"), "generated_bytes": len(comparison_bytes), "seconds": comparison_seconds, "fine_patches": int(comparison_stream["completed_patches"]), "level2_pools": int( comparison_stream["completed_nested_patches"] ), "level3_pools": int( comparison_stream.get("completed_tertiary_patches", 0) ), "image_report": image_report_payload( *extract_ppm_images(prompt + comparison_bytes) ), } ) del comparison_stream gc.collect() # Keep base64 PNG payloads in the HTML only; terminal/stats JSON should # stay compact even when a rollout contains a large rendered image. stats["comparison_rollouts"] = [ {key: value for key, value in item.items() if key != "image_report"} for item in comparisons ] if args.output: Path(args.output).expanduser().write_bytes(combined) if args.image_output_dir: image_paths = write_extracted_images( Path(args.image_output_dir).expanduser(), full_images, prefix="full" ) stats["image_output_files"] = image_paths for item in comparisons: report = item.get("image_report") or {} comparison_images = [] for encoded in report.get("images", []): raw = base64.b64decode(str(encoded["data_uri"]).split(",", 1)[1]) comparison_images.append({**encoded, "png": raw}) stats.setdefault("comparison_image_output_files", {})[item["mode"]] = ( write_extracted_images( Path(args.image_output_dir).expanduser(), comparison_images, prefix=str(item["mode"]), ) ) print(json.dumps(stats, indent=2)) if args.stats_json: Path(args.stats_json).expanduser().write_text( json.dumps(stats, indent=2), encoding="utf-8" ) if args.html_output: html_path = Path(args.html_output).expanduser() html_path.parent.mkdir(parents=True, exist_ok=True) html_path.write_text( make_generation_html( prompt, generated, attributions, stats, comparisons=comparisons, image_report=full_image_report, ), encoding="utf-8", ) print(f"wrote hierarchy influence report: {html_path}") if __name__ == "__main__": main()