diff --git a/README.md b/README.md index 937bf5b53..5c5da7226 100644 --- a/README.md +++ b/README.md @@ -1102,6 +1102,7 @@ examples/ - **Cosmos3-Nano text-to-video** (`config="cosmos3_video"`) — RTX 5090 BF16/FP8 denoise and complete benchmark workflow; [usage and performance](docs/cosmos3_video_usage.md) - **Cosmos3-Edge AV inverse dynamics and Reasoner** (`config="cosmos3_edge"`) — Jetson AGX Thor official baseline, 6.60x no-cache AV denoise, and NVFP4 multimodal chat decode; [complete usage and performance](docs/cosmos3_edge_thor.md) - **Qwen3-VL-8B** — RTX 5090 NVFP4/FP8 multimodal path, RTX 4090 official-FP8 path, and Jetson BF16 paths for Thor and Orin; [RTX 5090 usage](docs/qwen3_vl_nvfp4.md), [RTX 4090 usage](docs/qwen3_vl_fp8_sm89.md), [Jetson Thor usage](docs/qwen3_vl_thor.md), [Jetson Orin usage](docs/qwen3_vl_rtx_bf16.md) +- **Chameleon-7B** — Jetson Thor dynamic-FP8 prefill (~120 ms E2E, ~30 tok/s decode) and Jetson Orin INT8/QuaRot-INT4 path; [usage](docs/chameleon_usage.md), [Thor notes](docs/chameleon_thor_sm110.md), [Orin SM87 notes](docs/chameleon7b_rtx_sm87.md) - **MiniMax-Remover** — FP8 transformer + NVFP4 VAE video inpainting; [usage and performance](docs/minimax_remover_usage.md) - **MelBandRoformer** — kernelized FP8 audio source separation; [usage and performance](docs/melband_roformer_usage.md) - **OmniVoice TTS** — BF16/FP4 acceleration and HTTP serving; [serving quickstart](serving/omnivoice_agent/README.md) diff --git a/USAGE.md b/USAGE.md index 47e2414bd..32b94bfdf 100644 --- a/USAGE.md +++ b/USAGE.md @@ -681,6 +681,39 @@ once and reused, so per-chunk cost is just the infer row. This path does not change CMake targets, C++ bindings, or existing Pi0/Pi0.5/GROOT N1.6 runtime dispatch. +### Chameleon-7B (Thor) + +Standalone Chameleon-7B (text + image) is a direct-instantiation Thor +frontend — it is registered in `_PIPELINE_MAP` but is **not** dispatched by +`flash_rt.load_model` (same pattern as Qwen3-VL). The VQGAN image +tokenizer defaults to the generic eager Chameleon path; if compatible +TensorRT engines exist in the deployment, it is recommended to opt in +explicitly (`use_trt_vqgan=True`). A Jetson Orin (SM87) INT8/QuaRot +frontend is also available; see [`docs/chameleon_usage.md`](docs/chameleon_usage.md), +[`docs/chameleon_thor_sm110.md`](docs/chameleon_thor_sm110.md) and +[`docs/chameleon7b_rtx_sm87.md`](docs/chameleon7b_rtx_sm87.md). + +```python +from flash_rt.frontends.torch.chameleon_thor import ChameleonTorchFrontendThor + +fe = ChameleonTorchFrontendThor( + "/path/to/Chameleon_7B_mGPT", + use_fp8=True, # dynamic per-tensor FP8 (recommended default) + use_cuda_graph=True, + use_trt_vqgan=False, # generic default = eager VQGAN; set True when engines exist + target_size=512, +) + +out = fe.prefill("Describe the image.", [pil_image]) +# out["logits"]: (65536,) fp32 with mask_image_logits applied +``` + +FA4 attention is an explicit opt-in fast path +(`use_fa4_attn=True` / `FLASHRT_CHAMELEON_FA4_ATTN=1`, needs the +`thor-fa4` pip extra); measured transformer-prefill-only ≈ **104 ms** at +Se≈1056 vs ≈111 ms with CUTLASS FMHA. E2E (TRT VQGAN + FA4) ≈ **120 ms** +vs HF BF16 ≈ 403 ms transformer-only (~3.4×). + ### Wan2.2 TI2V-5B Wan2.2 TI2V-5B is exposed as an RTX SM120 official-pipeline baseline: diff --git a/benchmarks/chameleon_thor_latency.py b/benchmarks/chameleon_thor_latency.py new file mode 100644 index 000000000..c72521c8d --- /dev/null +++ b/benchmarks/chameleon_thor_latency.py @@ -0,0 +1,181 @@ +#!/usr/bin/env python3 +"""Chameleon-7B (Thor sm_110) latency benchmark. + +Measures standalone Chameleon-7B prefill latency on real images with clean +stage separation: + +- ``--reuse-input-ids``: build real-image input ids once, then time only + embed + backbone + lm_head (transformer-prefill-only, HF-comparable). +- Default: full ``prefill()`` E2E including VQGAN tokenization. +- ``--use-trt-vqgan``: explicit TensorRT VQGAN opt-in (recommended when + compatible engines exist; the generic default stays eager). +- FA4 attention: enable with ``FLASHRT_CHAMELEON_FA4_ATTN=1`` (needs the + ``thor-fa4`` pip extra; prints whether it is active). + +Latency is wall-clock P50 (per CONTRIBUTING.md: quickstart --benchmark +style; CUDA-graph replayed latency is what the pipeline measures inside the +graph). Every result row records device, VQGAN backend, FA4 state, Se, +fp8/fp16 and graph settings for reproducible reporting. + +Usage: + + python benchmarks/chameleon_thor_latency.py \ + --checkpoint /path/to/Chameleon_7B_mGPT \ + --image-dir /path/to/images \ + --iters 20 --warmup 5 +""" +from __future__ import annotations + +import argparse +import json +import pathlib +import time +from typing import Dict, List + +import torch + + +def _stats(xs: List[float]) -> Dict[str, float]: + a = sorted(xs) + n = len(a) + return {"mean": sum(a) / n, "p50": a[n // 2], "min": a[0], "max": a[-1]} + + +def _load_images(image_dir: pathlib.Path, max_images: int): + from PIL import Image + + paths = sorted( + p for p in image_dir.iterdir() + if p.suffix.lower() in (".jpg", ".jpeg", ".png", ".bmp")) + if not paths: + raise FileNotFoundError(f"No real images under {image_dir}") + paths = paths[:max_images] + return [Image.open(p).convert("RGB") for p in paths], [str(p) for p in paths] + + +def _pad_ids(input_ids: List[int], pad_id: int = 1): + real_len = len(input_ids) + padded = list(input_ids) + rem = len(padded) % 16 + if rem: + padded.extend([pad_id] * (16 - rem)) + return padded, real_len + + +def _prefill_once(fe, prompt, images, cached_ids, use_graph: bool) -> Dict[str, float]: + """One timed prefill; returns stage latencies in ms.""" + times: Dict[str, float] = {} + t0 = time.perf_counter() + ids = fe.encode_prompt(prompt, images) if cached_ids is None else cached_ids + torch.cuda.synchronize() + times["encode_ms"] = (time.perf_counter() - t0) * 1000.0 + + padded, real_len = _pad_ids(ids) + fe._real_len = real_len + fe.Se = len(padded) + fe._last_input_ids = padded + if fe._use_autotune: + fe._autotune_gemms(fe.Se) + torch.cuda.synchronize() + t1 = time.perf_counter() + + fe._embed_ids(padded) + torch.cuda.synchronize() + t2 = time.perf_counter() + + if use_graph: + fe._capture_graph(fe.Se) + fe._infer_graph.replay() + else: + fe._run_backbone(fe.Se) + torch.cuda.synchronize() + t3 = time.perf_counter() + + fe._project_last() + torch.cuda.synchronize() + t4 = time.perf_counter() + + times["prepare_ms"] = (t1 - t0) * 1000.0 - times["encode_ms"] + times["embed_ms"] = (t2 - t1) * 1000.0 + times["backbone_ms"] = (t3 - t2) * 1000.0 + times["lm_head_ms"] = (t4 - t3) * 1000.0 + times["transformer_ms"] = times["embed_ms"] + times["backbone_ms"] + times["lm_head_ms"] + times["total_ms"] = (t4 - t0) * 1000.0 + return times + + +def main() -> None: + ap = argparse.ArgumentParser(description=__doc__) + ap.add_argument("--checkpoint", required=True) + ap.add_argument("--image-dir", required=True, + help="Directory of real input images") + ap.add_argument("--prompt", default="Describe the image.") + ap.add_argument("--max-images", type=int, default=1) + ap.add_argument("--target-size", type=int, default=512) + ap.add_argument("--use-trt-vqgan", action="store_true", + help="Use TensorRT VQGAN if compatible engines exist " + "(recommended when engines are available; default is eager VQGAN)") + ap.add_argument("--reuse-input-ids", action="store_true", + help="time transformer prefill only (VQGAN excluded)") + ap.add_argument("--no-graph", action="store_true") + ap.add_argument("--use-fp16", action="store_true", + help="FP16 reference path instead of dynamic FP8") + ap.add_argument("--iters", type=int, default=20) + ap.add_argument("--warmup", type=int, default=5) + ap.add_argument("--output", default=None, help="JSON output path") + args = ap.parse_args() + + from flash_rt.frontends.torch.chameleon_thor import ChameleonTorchFrontendThor + from flash_rt.hardware.thor import fa4_backend + + images, image_paths = _load_images(pathlib.Path(args.image_dir), args.max_images) + use_graph = not args.no_graph + fp8 = not args.use_fp16 + + fe = ChameleonTorchFrontendThor( + args.checkpoint, use_fp8=fp8, use_cuda_graph=use_graph, + target_size=args.target_size, use_trt_vqgan=args.use_trt_vqgan) + + cached_ids = None + if args.reuse_input_ids: + cached_ids = fe.encode_prompt(args.prompt, images) + torch.cuda.synchronize() + + for _ in range(args.warmup): + _prefill_once(fe, args.prompt, images, cached_ids, use_graph) + + rows: Dict[str, List[float]] = {} + for _ in range(args.iters): + t = _prefill_once(fe, args.prompt, images, cached_ids, use_graph) + for k, v in t.items(): + rows.setdefault(k, []).append(v) + + result = { + "model": "chameleon-7b", + "device": torch.cuda.get_device_name(0), + "sm_count": torch.cuda.get_device_properties(0).multi_processor_count, + "checkpoint": args.checkpoint, + "image_paths": image_paths, + "prompt": args.prompt, + "target_size": args.target_size, + "fp8": fp8, + "cuda_graph": use_graph, + "vqgan_backend": fe.vqgan_backend, + "fa4_attn": fe.fa4_attn_active, + "fa4_status": fa4_backend.status(), + "reuse_input_ids": bool(args.reuse_input_ids), + "Se": int(fe.Se), + "latency_ms": {k: _stats(v) for k, v in rows.items()}, + } + for k, v in result["latency_ms"].items(): + print(f"[chameleon] {k:14s} p50={v['p50']:8.1f} ms mean={v['mean']:8.1f}") + print(f"[chameleon] device={result['device']} vqgan={fe.vqgan_backend} " + f"fa4={fe.fa4_attn_active} fp8={fp8} graph={use_graph} Se={fe.Se}") + + if args.output: + pathlib.Path(args.output).write_text(json.dumps(result, indent=2)) + print(f"[chameleon] wrote {args.output}") + + +if __name__ == "__main__": + main() diff --git a/benchmarks/fp4_chameleon_layer16.py b/benchmarks/fp4_chameleon_layer16.py new file mode 100644 index 000000000..54bc0745f --- /dev/null +++ b/benchmarks/fp4_chameleon_layer16.py @@ -0,0 +1,177 @@ +#!/usr/bin/env python3 +"""Verify FP4 Gate+Up substitution at Chameleon layer-16 FFN shape. + +Compares three paths driven by identical fp16 weights/activation: + + REF : pure fp16 matmul + silu*mul + matmul (fp32 accumulate) + FP8 : full FP8 path (used by current chameleon_forward) + MIX : FP4 Gate+Up + (existing) silu_mul_split_fp8_fp16 + FP8 Down + ALL4: FP4 Gate+Up + fp16 silu*mul + FP4 Down (upper bound) + +For each path: cosine similarity vs REF + microbenchmark latency. +""" +import pytest + +torch = pytest.importorskip("torch") +fp4 = pytest.importorskip( + "flash_rt.flash_rt_fp4", + reason="flash_rt_fp4 requires an NVFP4 (sm_120+) build") +import numpy as np +import flash_rt.flash_rt_kernels as fvk +from flash_rt.executors.fp4_utils import ( + quant_weight_nvfp4, FP4ActScratch, quant_act_nvfp4, fp4_gemm, pick_variant, +) + + +def fp16_t(*shape, scale=1.0): + return (torch.randn(*shape, dtype=torch.float16, device='cuda') * scale).contiguous() + + +def cuda_time(fn, iters=100, warmup=20): + s = torch.cuda.current_stream() + for _ in range(warmup): fn() + s.synchronize() + e0 = torch.cuda.Event(enable_timing=True); e1 = torch.cuda.Event(enable_timing=True) + e0.record() + for _ in range(iters): fn() + e1.record(); s.synchronize() + return e0.elapsed_time(e1) / iters * 1000 # μs + + +def amax_scale(t: torch.Tensor) -> float: + return max(t.abs().max().item() / 448.0, 1e-9) + + +def make_scale_buf(scale: float) -> torch.Tensor: + return torch.tensor([scale], dtype=torch.float32, device='cuda') + + +def quant_fp8(W: torch.Tensor, scale: float): + out = torch.empty_like(W, dtype=torch.uint8) + sb = make_scale_buf(scale) + fvk.quantize_fp8_static_fp16(W.data_ptr(), out.data_ptr(), + sb.data_ptr(), W.numel(), 0) + return out, sb + + +def cos_vs(a, b): + return torch.nn.functional.cosine_similarity( + a.flatten().float().unsqueeze(0), + b.flatten().float().unsqueeze(0)).item() + + +def main(): + print(f"FP4 enabled: {fp4.has_nvfp4()}; variants: {fp4.cutlass_fp4_gemm_num_variants()}") + + Se, D, Dff = 1216, 4096, 11008 + + torch.manual_seed(0) + W_g = fp16_t(Dff, D, scale=0.02) + W_u = fp16_t(Dff, D, scale=0.02) + W_d = fp16_t(D, Dff, scale=0.02) + X = fp16_t(Se, D, scale=1.0) + + # ---- REF ---- + gate_ref = (X.float() @ W_g.float().T).half() + up_ref = (X.float() @ W_u.float().T).half() + h_ref = (torch.nn.functional.silu(gate_ref.float()) * up_ref.float()).half() + out_ref = (h_ref.float() @ W_d.float().T).half() + print(f"REF: |gate|max={gate_ref.abs().max():.2f} |up|max={up_ref.abs().max():.2f}" + f" |h|max={h_ref.abs().max():.2f} |out|max={out_ref.abs().max():.2f}") + + # ---- Pre-compute calibrated scales (per-tensor amax/448) ---- + s_x = amax_scale(X) + s_wg = amax_scale(W_g) + s_wu = amax_scale(W_u) + s_wd = amax_scale(W_d) + s_h = amax_scale(h_ref) # post-silu*up → fp8 input to Down + print(f"scales: x={s_x:.3e} w_g={s_wg:.3e} w_u={s_wu:.3e} w_d={s_wd:.3e} h={s_h:.3e}") + + gemm = fvk.GemmRunner() + + # FP8 weights + activation + # NB: fp8_nn_dev is NN (no transpose), so B must be [K, N] row-major. + # We store HF-style W as [N, K]; transpose before fp8 quant. + Wg_fp8, sg = quant_fp8(W_g.t().contiguous(), s_wg) # [D, Dff] + Wu_fp8, su = quant_fp8(W_u.t().contiguous(), s_wu) # [D, Dff] + Wd_fp8, sd = quant_fp8(W_d.t().contiguous(), s_wd) # [Dff, D] + sx_buf = make_scale_buf(s_x); sh_buf = make_scale_buf(s_h) + X_fp8 = torch.empty(Se, D, dtype=torch.uint8, device='cuda') + fvk.quantize_fp8_static_fp16(X.data_ptr(), X_fp8.data_ptr(), + sx_buf.data_ptr(), Se*D, 0) + + gate_out = torch.empty(Se, Dff, dtype=torch.float16, device='cuda') + up_out = torch.empty(Se, Dff, dtype=torch.float16, device='cuda') + gu_fp8 = torch.empty(Se, Dff, dtype=torch.uint8, device='cuda') + out_fp8 = torch.empty(Se, D, dtype=torch.float16, device='cuda') + + def run_fp8(): + gemm.fp8_nn_dev(X_fp8.data_ptr(), Wg_fp8.data_ptr(), gate_out.data_ptr(), + Se, Dff, D, sx_buf.data_ptr(), sg.data_ptr(), 0) + gemm.fp8_nn_dev(X_fp8.data_ptr(), Wu_fp8.data_ptr(), up_out.data_ptr(), + Se, Dff, D, sx_buf.data_ptr(), su.data_ptr(), 0) + fvk.silu_mul_split_fp8_fp16(gate_out.data_ptr(), up_out.data_ptr(), + gu_fp8.data_ptr(), Se*Dff, + sh_buf.data_ptr(), 0) + gemm.fp8_nn_dev(gu_fp8.data_ptr(), Wd_fp8.data_ptr(), out_fp8.data_ptr(), + Se, D, Dff, sh_buf.data_ptr(), sd.data_ptr(), 0) + + run_fp8(); torch.cuda.synchronize() + cos_fp8 = cos_vs(out_fp8, out_ref) + fp8_us = cuda_time(run_fp8) + + # ---- MIX (FP4 Gate+Up, FP8 Down) ---- + qg = quant_weight_nvfp4(W_g) + qu = quant_weight_nvfp4(W_u) + sc_x = FP4ActScratch(max_M=Se, K=D) + var_gu = pick_variant(Dff, D) + out_mix = torch.empty(Se, D, dtype=torch.float16, device='cuda') + + def run_mix(): + quant_act_nvfp4(X, sc_x, Se, stream=0) + fp4_gemm(sc_x, qg, gate_out, Se, Dff, D, variant_idx=var_gu, stream=0) + fp4_gemm(sc_x, qu, up_out, Se, Dff, D, variant_idx=var_gu, stream=0) + fvk.silu_mul_split_fp8_fp16(gate_out.data_ptr(), up_out.data_ptr(), + gu_fp8.data_ptr(), Se*Dff, + sh_buf.data_ptr(), 0) + gemm.fp8_nn_dev(gu_fp8.data_ptr(), Wd_fp8.data_ptr(), out_mix.data_ptr(), + Se, D, Dff, sh_buf.data_ptr(), sd.data_ptr(), 0) + + run_mix(); torch.cuda.synchronize() + cos_mix = cos_vs(out_mix, out_ref) + mix_us = cuda_time(run_mix) + + # ---- ALL-FP4 (Gate+Up+Down all FP4, fp16 silu*mul) ---- + qd = quant_weight_nvfp4(W_d) + sc_h = FP4ActScratch(max_M=Se, K=Dff) + var_dn = pick_variant(D, Dff) + h_buf = torch.empty(Se, Dff, dtype=torch.float16, device='cuda') + out_all4 = torch.empty(Se, D, dtype=torch.float16, device='cuda') + + def run_all4(): + quant_act_nvfp4(X, sc_x, Se, stream=0) + fp4_gemm(sc_x, qg, gate_out, Se, Dff, D, variant_idx=var_gu, stream=0) + fp4_gemm(sc_x, qu, up_out, Se, Dff, D, variant_idx=var_gu, stream=0) + # fp16 silu*up via torch (bench-only) + torch.mul(torch.nn.functional.silu(gate_out), up_out, out=h_buf) + quant_act_nvfp4(h_buf, sc_h, Se, stream=0) + fp4_gemm(sc_h, qd, out_all4, Se, D, Dff, variant_idx=var_dn, stream=0) + + run_all4(); torch.cuda.synchronize() + cos_all4 = cos_vs(out_all4, out_ref) + all4_us = cuda_time(run_all4) + + print() + print("="*72) + print("Chameleon layer-16 FFN block (Se=1216, D=4096, Dff=11008)") + print("="*72) + fmt = " {:14s} cos_vs_ref = {:.6f} {:7.1f} μs speedup={}" + print(fmt.format("FP8 baseline", cos_fp8, fp8_us, "1.00x")) + print(fmt.format("MIX (FP4 GU)", cos_mix, mix_us, f"{fp8_us/mix_us:.2f}x")) + print(fmt.format("ALL-FP4", cos_all4, all4_us, f"{fp8_us/all4_us:.2f}x")) + delta = fp8_us - mix_us + print(f"\n Per-layer MIX saves {delta:6.1f} μs → 32 layers ≈ {delta*32/1000:5.2f} ms") + + +if __name__ == '__main__': + main() diff --git a/docs/benchmark_comparison.md b/docs/benchmark_comparison.md index 50fad3f50..3294eead4 100644 --- a/docs/benchmark_comparison.md +++ b/docs/benchmark_comparison.md @@ -255,6 +255,27 @@ TRT-aligned FP4 loop comparison, using the same quantization scheme. | 25 | ~304 ms | **97.5 ms** | **~3.1x** | | 50 | ~608 ms | **155.8 ms** | **~3.9x** | +## Chameleon-7B + +Baseline: HF `transformers` BF16 eager, transformer-only (input ids +built by FlashRT, model forward timed). FlashRT rows are the same harness +(real image `hand_1.jpg`, prompt "Describe the image.", target_size=512, +Se≈1053-1072, wall-clock P50). VQGAN backend and FA4 state are recorded per +row — generic default is eager VQGAN + CUTLASS FMHA; TRT VQGAN and FA4 are +explicit opt-ins. + +| FlashRT FP8 | VQGAN | FA4 | Latency | Speedup vs HF | +|---|---:|---:|---:|---:| +| transformer-only | (n/a) | on | **104.2 ms** | **3.9×** | +| transformer-only | (n/a) | off | 111.2 ms | 3.6× | +| E2E | TRT opt-in | on | **120.2 ms** | **3.4×** | +| E2E | eager | on | 177.3 ms | 2.3× | +| E2E | eager | off | ~190 ms | ~2.1× | + +| Baseline (HF BF16) | Latency | +|---|---:| +| transformer-only | 402.9 ms | + ## Qwen3-8B LLM rows list the baseline and FlashRT measurements without speedup. diff --git a/docs/chameleon7b_rtx_sm87.md b/docs/chameleon7b_rtx_sm87.md new file mode 100644 index 000000000..410b129d0 --- /dev/null +++ b/docs/chameleon7b_rtx_sm87.md @@ -0,0 +1,722 @@ +# Chameleon-7B VLM on Jetson AGX Orin (SM87) — FlashRT adaptation + +> **Status: Phase 1 complete — Gate 1 PASSES.** Production tier is +> **INT8 W8A8 + Hadamard (QuaRot at 8 bits)**: greedy output is **bit-identical +> to the HF bf16 reference for 16/16 tokens**, worst layer cosine 0.9986, +> last-row logit cosine 0.999968, at **21.07 tok/s decode and 273.8 ms warm +> prefill** (ISL=1032) — both at their measured ceilings (§4.2). Both defects found during bring-up were fixed by +> **quantization-method changes, not precision fallbacks**: a basis rotation for +> the massive-activation outliers (§4.5) and a ported clamp for the L31 FP16 +> overflow (§4.6). +> +> This is the authoritative document for **upstream Chameleon-7B as an image+text +> → text VLM** on Orin SM87. + +--- + +## 0. Conclusion first + +**Shipped configuration** — `ChameleonTorchFrontendRtxSm87`, all defaults: + +| | | +|---|---| +| LLM GEMMs (Q/K/V/O, gate, up) | **INT8 W8A8 + Hadamard rotation** (QuaRot at 8 bits), per-row dynamic activation scales | +| FFN down | INT8 W8A8, per-row dynamic (K=11008 is not a power of two) | +| lm_head | INT8 W8A8 | +| residual / QK-LayerNorm / RoPE / attention / KV cache | FP16, with `ffn_down_clamp=60000` on the last 4 layers | +| attention | FA2 fp16 causal, `split_kv_bias=4` on decode | +| VQ-GAN encoder | FP16 convs, **fp32** codebook distance/argmin | + +**Result** (ISL=1032 = one 512² image + prompt, OSL=16, warm): + +| metric | value | +|---|---| +| greedy output vs HF bf16 reference | **bit-identical, 16/16 tokens** | +| worst per-layer residual cosine (L0..L31) | **0.9986** | +| last-row logit cosine | **0.999968** | +| decode | **21.07 tok/s** (47.5 ms/token) | +| prefill (LLM) | **273.8 ms** — GEMM at 91 % of the achievable CUTLASS ceiling | +| steady GPU memory | **7.6 GB** | + +**Why it works, in one line:** the Chameleon backbone's massive-activation +channels are fixed by a **basis rotation** rather than by more bits, finer +granularity, smoothing, or a per-layer precision fallback — and because the +rotation preserves per-row scales it reuses the stock CUTLASS INT8 GEMMs, so it +costs nothing (§4.5). + +## 1. Platform + +| Field | Value | +|---|---| +| Device | NVIDIA Jetson AGX Orin 64 GB (`torch.cuda.get_device_properties`: `Orin`, 61.4 GB) | +| GPU | SM **8.7** (Ampere), **16 SMs**, L2 4 MB | +| Memory | LPDDR5X unified, 204.8 GB/s spec | +| FP8 / FP4 | **not native** (Ada sm89+/Hopper/Blackwell only) → INT8/INT4 is the only low-bit route | +| CUDA / torch | 12.2 / 2.3.0 | +| Build | `cmake -B build -S . -DGPU_ARCH=87 -DFA2_ARCH_NATIVE_ONLY=ON -DFA2_HDIMS='128;256' -DFA2_DTYPES='fp16;bf16'` | + +**Measured bandwidth (this is the load-bearing calibration).** Three different numbers, and +picking the wrong one produces >100 % "efficiency" nonsense: + +| probe | result | use for | +|---|---|---| +| D2D copy, 512 MB, read+write | **124.5–126.0 GB/s** | copy-bound ops | +| single-stream vectorized reduce (`roofline.py --measure-bw`) | 99 GB/s | nothing — undersaturated | +| **best achieved by a real weight-streaming kernel** (int8 gate GEMM @ M=1) | **173.3 GB/s** = 85 % of spec | **the decode roofline denominator** | + +A D2D copy measures 124.5 GB/s in this container. Treat 173 GB/s +(kernel-achieved, read-dominated) as the decode ceiling. + +> ⚠️ `/sys/devices/gpu.0/devfreq/*/cur_freq` is **not readable in this container**, so clocks +> cannot be locked or even observed. Every number below is warm (≥30 warmup iters) and a median +> over ≥50 iters. Cross-config *ratios* are trustworthy; absolute values carry DVFS uncertainty +> (idle 306 MHz vs 1300.5 MHz loaded — principle #16). + +## 2. Phase 0 — probe verdicts + +### 2.1 R1 — FA2 split-KV is a silent no-op at 32 Q heads ⚠️ **and the fix is one Python argument** + +`csrc/attention/fa2_wrapper_causal.cu:41-43,152-158`: + +``` +num_splits = fa2_num_splits_heuristic_causal(batch*num_heads_q*num_m_blocks, num_sms*2, ...) + → if (batch_nheads_mblocks >= 0.8f * num_SMs) return 1; +``` + +Chameleon decode: `1*32*1 = 32` vs `0.8 * (16*2) = 25.6` → **`num_splits = 1`**. Passing the +accumulators does nothing. A split-KV win works *only* when the model has few +enough Q heads (e.g. **16**) for the heuristic to engage splitting. + +`num_sms` is a pure heuristic knob in this wrapper, so biasing it selects the split count. +Measured (q=1, kv=1040, 32 Q heads, head_dim 128, fp16): + +| `num_sms` passed | latency | speedup | max abs diff vs no-split | +|---|---|---|---| +| 0 (baseline, no accum) | 204.9 µs | 1.00× | — | +| 16 (**real SM count**) | 195.2 µs | 1.05× | **0.000e+00** ← proves `num_splits=1` | +| 32 | 149.8 µs | 1.37× | 1.221e-04 | +| **64** | **141.8 µs** | **1.44×** | 1.221e-04 | +| 128 | 154.8 µs | 1.32× | 1.221e-04 (over-split) | + +**Verdict: ship a `split_kv_bias` backend parameter, default 4× (`num_sms=64`).** 1.44× on +decode attention, pure Python, graph-safe. The 1.221e-04 delta is fp16 accumulation-order noise +(fp16 eps ≈ 9.8e-4 at magnitude 1), not an error. + +### 2.2 R2/R3 — M=1 GEMM: INT8 needs no GEMV; INT4 needs a small-M tile + +All shapes at M=1, achieved GB/s = weight bytes / time, normalized to the **173.3 GB/s** +kernel-achieved read ceiling. + +| shape | INT8 variant | µs | GB/s | %ceil | INT4 variant | µs | GB/s | %ceil | +|---|---|---|---|---|---|---|---|---| +| Q/K/V/O 4096×4096 | fp16out | 144.2 | 116.4 | 67 % | fp16out | 121.4 | 69.1 | **40 %** ⚠️ | +| gate 11008×4096 | bf16out | 260.2 | 173.3 | **100 %** | bf16out | 178.5 | 126.3 | 73 % | +| up+silu 11008×4096 | silu_gated | 273.3 | 165.0 | 95 % | silu_gated | 165.7 | 136.1 | 79 % | +| down 4096×11008 | fp16out | 267.3 | 168.7 | 97 % | fp16out | 170.7 | 132.1 | 76 % | +| lm_head 65536×4096 | bf16out | 1743.8 | 153.9 | 89 % | *(stays int8)* | — | — | — | +| **per token (GEMM only)** | | **47.70 ms** | | | | **33.45 ms** | | | +| **→ tok/s (GEMM only)** | | **21.0** | | | | **29.9** | | | + +**Verdict 1 — INT8: ship CUTLASS as-is, do not write a GEMV.** Four of five shapes are at +89–100 % of the achieved read ceiling. Measured 47.70 ms vs the predicted 44 ms weight floor — +**measured ≈ predicted, so the bottleneck model is correct** (principle #15). + +**Verdict 2 — INT4 delivers only 1.43×, not 2×.** Root cause confirmed in source: +`csrc/gemm/cutlass_sm80_int4_rowwise.cu:61` defines exactly one tile (`GemmShape<128,128,128>`) +with **no `M<=64` dispatcher**, whereas INT8 dispatches `M<=64 → 64×128` +(`cutlass_sm80_int8_rowwise_fp16out.cu:330-333`). At M=1 INT4 therefore wastes a 128-row tile — +visible as Q/K/V/O at **40 %** of ceiling (69.1 GB/s) versus INT8's 67 % (116.4 GB/s) on the +same shape with half the bytes. → **`cutlass_sm80_int4_rowwise_t64x128.cu` is justified**; +predicted recovery 33.45 → ~27 ms/token (~37 tok/s GEMM-only). + +**Verdict 3 — the planned "M=1 up-projection split" lever is DEAD. Dropped before writing any +code.** The hypothesis was that `cutlass_int8_silu_gated_bf16out` (128×128 tile only, +`cutlass_sm80_int8_silu_gated.cu:54`) would lose ~93 µs/layer at M=1 versus a +`bf16out` (64-tile) + `silu_mul_qwen36_bf16` split. Measured: **273.3 µs vs 260.3 µs = 13 µs**, +i.e. 0.4 ms/token ≈ 0.9 % — and the split adds a `silu_mul` launch plus an 11008-element bf16 +round trip that roughly cancels it. The prediction was **7× too optimistic**; both GEMMs are +already bandwidth-bound. (Principle #13: microbenchmark before writing the kernel.) + +### 2.3 R4 — HF reference: **PASS** + +Stock transformers 4.57.1 has `ChameleonForConditionalGeneration`, but its `ChameleonLayerNorm` +builds weights of shape `(num_heads, head_dim) = (32,128)` +(`transformers/models/chameleon/modeling_chameleon.py:187-202, 281-282`) while this Lumina-mGPT +checkpoint stores `(1,128)` — so it **cannot be loaded directly**. + +Working recipe (also: 4.57 *rejects* `state_dict=` together with a checkpoint path, so a +naked-constructor pattern is required — exactly the one used by the HF reference builder in +`scripts/chameleon_orin_check.py`): + +1. read both shards, `repeat_interleave(32, dim=0)` the **128** `self_attn.{q,k}_norm.{weight,bias}` tensors; +2. `torch.set_default_dtype(torch.bfloat16)`; `ChameleonForConditionalGeneration(cfg)` with `cfg._attn_implementation = "eager"`; +3. `load_state_dict(sd, strict=False)` → **0 missing, 0 unexpected** (548 tensors); `.eval().cuda()` → 13.1 GB. + +Verified in the same run: `mask_image_logits` is live — logits over ids **4..8195** come back at +`-3.390e+38` = `finfo(bf16).min`, while the text range max is `-13.375`. Since +`model_parallel_size == 1`, the expansion is a pure broadcast, so this reference is numerically +equivalent to official `facebook/chameleon-7b`. + +### 2.4 R5 — VQ-GAN codebook argmin precision + +`ChameleonVQVAE._from_config` + the 129 `model.vqmodel.*` tensors load with **0 missing / 0 +unexpected** (confirming the checkpoint is encoder-only and so is the HF module). Codebook index +match on a deterministic 512×512 input, versus a full-fp32 reference: + +| configuration | index match | +|---|---| +| fp32 convs + fp32 argmin | 100.00 % | +| **fp16 convs + fp16 argmin** | 98.14 % | +| **fp16 convs + fp32 distance/argmin** | 99.02 % | + +The fp32-argmin fix helps (`z²+e²−2ez` is cancellation-prone in fp16; +`modeling_chameleon.py:850-861`) and costs <0.1 ms. This probe used random noise; +divergence on **real images** can be much higher (~92 % has been observed), so +re-measure on real content at Phase 5 before declaring the fix sufficient. + +### 2.5 R8 — decode-graph primitives are graph-safe: **PASS** + +Captured `index_select(emb, 0, tok, out=x)` → lm_head stand-in → `mask_view.fill_(bf16_min)` → +`argmax(out=)` → `tok.copy_(out_tok)`: capture succeeded (no `code=13`), the **stale-value test +passed** (changing the seed token changed the embedding output, maxdiff 4.85 — i.e. the graph +re-reads `tok` rather than baking it), and the logit mask survived replay. So no fp16 +embedding-lookup kernel is needed. + +### 2.6 R6 — deferred + +Original `original_tokenizers/vqgan.{yaml,ckpt}` vs HF `model.vqmodel.*` equivalence only gates +the **TRT engine** track (a TRT engine built from the checkpoint VQ-VAE weights must produce the same tokens as +the safetensors weights). Deferred to Phase 5. + +## 3. Design + +### 3.1 Zero new CUDA kernels for the QK-Norm / RoPE / KV path + +Three source facts combine to make the existing prefill kernel cover decode too: + +1. CUTLASS int8/int4 GEMM output row stride is hard-wired to `N` + (`cutlass_sm80_int8_rowwise_fp16out.cu:169-171`), so a `[32, max_seq, 32, 128]` fp16 KV cache — + whose per-layer slab is a contiguous `[max_seq, 4096]` with row stride exactly 4096 == N — is + a **legal GEMM destination**. `AlignmentC=8` (16 B) is satisfied by both the layer and row offsets. +2. `qk_norm_rope_fused_fp16` is in-place with implicit row stride `dim=128` and derives position + as `seq_pos = row / num_heads` (`qk_norm_rope_fused.cu:56-57, 65-66`) — so at `seq_len=1` every + row maps to row 0 of whatever cos/sin pointer it is handed. +3. The RoPE tables are C-contiguous `[max_seq, 128]` fp16 (`ChameleonTorchFrontendRtxSm87` + builds them that way), so position `pos` is `data_ptr() + pos*128*2` bytes. + +Therefore: + +* **prefill** — point the K and V GEMMs at `Kcache + li*layer_stride` / `Vcache + li*layer_stride`, + then call `qk_norm_rope_fused_fp16` unchanged (V needs no transform); +* **decode** — point them at `+ pos*4096*2` and call the same kernel with `seq_len=1` and cos/sin + pre-offset by `pos*128*2`. + +Also required: **`Se` must not be even-padded** (e.g. for FP8 GEMM alignment) — +with a real KV cache the pad row is junk that decode *will* attend to, and +CUTLASS constrains only `K`. + +Attention correctness: FA2 causal is **bottom-right aligned** +(`fa2_wrapper_causal.cu:126-138`), so `q=1, kv=N` attends all N keys. The cuBLAS fallback +`attention_mha_causal_fp16` is **top-left aligned** (`softmax.cu:182-191` masks with +`q = row % S_q`) and is therefore *silently wrong* at q=1 — the Chameleon backend must **raise** +rather than degrade to it. + +### 3.2 Precision policy + +| Component | Default (lossless tier) | Opt-in tier | +|---|---|---| +| Q/K/V/O, gate/up, down | INT8 W8A8, per-output-row weight scale, dynamic per-row act | QuaRot W4A4 (`use_int4`), down via block-diagonal `H_128` (`use_int4_down`) | +| lm_head (65536×4096) | INT8 (268 MB/token = 1.74 ms = 3.7 % of budget) | stays INT8 — never int4 | +| residual / RMSNorm / QK-LayerNorm / RoPE / attention / KV cache | FP16 | unchanged | +| VQ-GAN encoder | FP16 convs + **fp32 distance/argmin** | TRT FP16 (Phase 5) | + +Decode always uses **dynamic per-row** activation quant — never the prefill static calibration, +which was fitted at M=Se and does not describe a single decode row. + +### 3.3 Token contract + +`[BOS 0] + n_img × ([8197 ] + [8711 ]×1024 + [8196 ]) + text + [8710 sep]`, +so `S = 1 + n_img*1026 + n_text + 1`. Image token id = **VQ codebook index + 4**, exactly, for +all 8192 codes; the 1024 tokens are a raster scan of the 32×32 latent grid. + +> ⚠️ **Trap:** do not hardcode `: 8710, : 8720`. Both are +> **wrong** for upstream Chameleon (8710 is the `sep_token`). Likewise, special +> ids 65536-65539 are out of range for `vocab_size=65536`. This is one of three +> reasons the Chameleon frontend is standalone rather than a subclass — see §4. +> +> ⚠️ `config.json` says `bos_token_id: 1`, which is **stale** (``); `tokenizer.json` gives +> ` = 0` and that is what the processor emits. + +### 3.4 Why the frontend is standalone (not a subclass) + +Three of the most attractive inheritable helpers from a VLA-style frontend are +*actively wrong* for upstream Chameleon: its `_preprocess_image` is bicubic/384/`x*2-1` where +Chameleon needs PIL **LANCZOS**/512/`u8*0.0078-1.0` → `[-1, +0.989]`; its `_vqgan_encode` emits a +grid+newline token layout instead of a bare 1024 raster; its `_load_tokenizer` / +`_init_special_token_ids` produce the wrong ids above. Genuinely reusable are the quantizers and +`_split_fused_llm_weights` — extracted to `flash_rt/frontends/torch/_chameleon_quant.py`. + +## 4. Phase 1 — implementation and Gate 1 + +### 4.1 Shipped + +| file | role | +|---|---| +| `flash_rt/frontends/torch/_chameleon_quant.py` | checkpoint-agnostic INT8 / QuaRot-INT4 weight quantizers + the fused-projection split | +| `flash_rt/frontends/torch/_chameleon_spec.py` | weight spec — an inlined, Chameleon-specific `_llm_block()` (no dependency on any other model's spec) + `embed`/`norm`/`lm_head` singletons | +| `flash_rt/hardware/rtx/attn_backend_chameleon.py` | `ChameleonAttnBackend` — real per-layer FP16 KV cache, prefill + decode, `split_kv_bias` | +| `flash_rt/models/chameleon/pipeline_rtx.py` | one `chameleon_forward` serving both prefill (`pos=None`) and decode (`S=1`, `pos` set) | +| `flash_rt/frontends/torch/chameleon_rtx_sm87.py` | `ChameleonTorchFrontendRtxSm87` — `set_prompt` / `prefill` / `decode_step` / `generate` | +| `scripts/chameleon_orin_check.py` | Gate-1 harness (HF reference + graph safety + overflow guard) | +| `flash_rt/hardware/__init__.py`, `flash_rt/api.py` | dispatch entry + `_SM87_ALLOWED` + chat-VLM redirect | + +**Zero new CUDA kernels**, as predicted in §3.1. + +### 4.2 Measured performance (warm p50 over 10 iters after 2 discarded) + +| quantity | measured | predicted (§5) | verdict | +|---|---|---|---| +| **decode, ISL=1032** | **47.5 ms/token = 21.07 tok/s** | 53.7 ms = 18.6 tok/s | **beats prediction** -> bottleneck model correct | +| **prefill LLM, ISL=1032** | **273.8 ms** (min 270.6) | ~255 ms | **within 7 %** | +| prefill, plain INT8 (no rotation) | 282.2 ms | — | rotation is free at prefill too | +| prefill, **first call** | 460 ms | — | **1.68x cold-start penalty** — CUTLASS workspace `cudaMalloc` + JIT | +| load / steady memory | 29 s / **7.6 GB** | — | fp16 originals freed after quantization | + +> WARNING: an earlier revision of this doc claimed "prefill 496 ms, 1.9x worse +> than predicted — unexplained". That was **our measurement error**: the Gate-1 +> harness calls `prefill()` exactly once, so the number included first-call +> CUTLASS workspace allocation and JIT. **There is no prefill gap.** Principle +> #16 exists for exactly this reason — warm before every measurement. + +### 4.2.1 Per-kernel breakdown (torch.profiler, S=1032, post-warmup) + +Total GPU 281.1 ms (measured before the clamp restriction in §4.2.3): + +| kernel | ms | % | calls | +|---|---|---|---| +| CUTLASS INT8 GEMM — Q/K/V/O `(1032,4096,4096)` | 72.71 | 25.9 % | 128 | +| CUTLASS INT8 GEMM — up + fused SiLU-gate | 57.81 | 20.6 % | 32 | +| CUTLASS INT8 GEMM — gate (bf16 out) | 49.36 | 17.6 % | 32 | +| CUTLASS INT8 GEMM — down (t256x128) | 47.61 | 16.9 % | 32 | +| FA2 fp16 causal | 14.07 | 5.0 % | 32 | +| `residual_add_rms_norm_fht` (rotation fused into the norm) | 11.52 | 4.1 % | 63 | +| `qk_norm_rope_fused_fp16` | 9.36 | 3.3 % | 32 | +| `quantize_int8_rowwise_vec8` (bf16, pre-down) | 6.38 | 2.3 % | 32 | +| `clamp_inplace_fp16` | 6.23 | 2.2 % | 32 | +| `fht_int8_quant` (pre-O) | 3.90 | 1.4 % | 32 | +| lm_head + tail | 2.17 | 0.8 % | 4 | + +=> **GEMM 229.2 ms (81.5 %)**, elementwise tail 37.8 ms (13.5 %), attention 14.07 ms (5.0 %). + +### 4.2.2 The real GEMM ceiling is 64.4 TOPS, not 84.8 — GEMM tuning is spent + +The often-quoted 84.8 TOPS figure is the **raw `mma.s8` issue rate** from a +register-only probe. What CUTLASS actually achieves on its best-case shape is +lower: big-square probes measure **58.7 TOPS at 4096^3 and 64.4 TOPS at +8192^3**. Against that realistic ceiling: + +| shape | ms/call | TOPS | vs 64.4 ceiling | +|---|---|---|---| +| Q/K/V/O `(1032,4096,4096)` | 0.567 | 61.1 | **95 %** | +| gate / up `(1032,11008,4096)` | 1.694 | 54.9 | 85 % | +| down `(1032,4096,11008)` | 1.540 | 60.4 | 94 % | +| **whole LLM** | **229.2** | **58.7** | **91 %** | + +The prefill GEMMs are at **91 % of the achievable CUTLASS ceiling**, and the +isolated probe reproduces the in-pipeline time to within 0.5 % (0.567 vs +0.568 ms on Q/K/V/O) — so there is no pipeline overhead left to recover. This +confirms that the GEMM ladder (swizzle Id4 / stages-5 / t256x128) is spent. +Only `gate/up` at 85 % shows slack, and tile sweeps there measured 256x128 as +"only ~2 % better, not worth a 4th instantiation". + +> WARNING: **the roofline probe itself had a DVFS bug**, found here. Whichever +> shape was measured *first* was penalised by clock ramp: Q/K/V/O reported +> **28.3 TOPS** measured first versus **61.1** for the identical shape after +> adding a 3-second saturating pre-ramp, and the per-shape TOPS ascended purely +> in measurement order (28.3 -> 53.8 -> 60.5). `_ramp_clocks()` now runs before +> any timing in the roofline script +> (`scripts/bench/orin_int8_roofline.py`). Any earlier per-shape number from that +> script is suspect. + +### 4.2.3 Clamp restricted to the last 4 layers: -6.3 ms + +`clamp_inplace_fp16` cost 6.23 ms (2.2 %) across all 32 layers, but the measured +magnitudes (§4.6) grow monotonically with depth and L28 is 1616 — **37x below the +60000 clamp** — so early layers can never reach it. Restricting it to the last +`ffn_down_clamp_last_n` layers (default 4): + +| | before | after | +|---|---|---| +| prefill warm p50 | 280.1 ms | **273.8 ms** | +| decode | 20.96 tok/s | **21.07 tok/s** | +| L31 / final / logit cosine, greedy text | 0.999722 / 0.999447 / 0.999968 / 16-of-16 | **bit-identical** | + +The Gate-1 harness reports per-layer clamp saturation, so a checkpoint that +violates the monotonicity assumption is detectable; set +`FLASHRT_CHAMELEON_DOWN_CLAMP_LAST_N=32` if that ever happens. + + +### 4.3 Gate 1 — PASS + +Real image (`FlashRT.png`), `"Describe this image."`, ISL=1032, OSL=16, +tier **INT8+Hadamard**: + +| check | result | gate | verdict | +|---|---|---|---| +| **greedy text identical to HF** | **16/16 tokens** | 16/16 | **PASS** | +| worst layer cosine (L0..L31) | **0.9986** | ≥0.97 | PASS | +| L31 / final-norm cosine | **0.999722 / 0.999447** | — | PASS | +| last-row logit cosine | **0.999968** | ≥0.999 | PASS | +| graph safety (capture + stale-value) | cos 0.9884 between two seed tokens | not frozen | PASS | +| FP16 residual finiteness | no inf/nan | finite | PASS | +| argmax, image positions (1026) | 89.6 % exact / **95.0 % tie-adjusted** | — | informational (§4.7) | +| argmax, text positions (6) | 5/6 | — | **not binding** — n=6 is too small; one near-tie flip moves it 17 points | + +Both engines produce `"The image is a logo for the company Flexsteel. The logo is a"`. + +### 4.4 Root cause: the row-0 massive activation + +Probing every layer 20-31 against the reference localizes the failure precisely. +It is **not** spread across the tensor — it is **row 0, the BOS/attention-sink +token**, in the last four layers: + +| layer | cosine (all rows) | worst-row cosine | worst row | FlashRT max\|x\| | ref max\|x\| | +|---|---|---|---|---|---| +| L24 | 0.9966 | 0.982 | row4 | 1993 | 2512 | +| L27 | 0.9922 | 0.994 | row4 | 1990 | 2512 | +| **L28** | 0.850 | **0.688** | **row0** | 240 | 1632 | +| **L31** | 0.691 | **−0.384** | **row0** | 3502 | 23936 | + +(ISL=20 text-only.) The reference's L31 row-0 norm is **42971 vs a median of +10456**, concentrated in a few channels — **d632 = 23936**, then d808, d1282, +d2669. FlashRT has 1225 in d632. + +Mechanism: per-row INT8 activation quantization sets `scale = amax/127` from that +outlier, so the other ~4090 channels of row 0 round to zero and the row's +direction is destroyed (cosine goes *negative*). This is the documented +Chameleon massive-activation zone (the skill's backbone profile names d671/d579 +for the L15→L19 band; here it is d632 in the L28→L31 band), and per principle +#3/#17 the fix is **basis rotation, not smoothing**. + +### 4.5 The tier ladder: rotation × bit-width (W8A8+Hadamard wins) + +Two independent error sources act here, and each shipped tier only fixed one: + +* **outlier conditioning** — a row whose amax is set by a massive-activation + channel loses its other ~4090 channels to rounding. Fixed by a **basis + rotation**, not by finer granularity or smoothing (principle #17). +* **quantization noise** — the resolution left for the other 1031 ordinary rows. + Fixed by **more bits**. + +Measured on the same prompt, all three tiers: + +| ISL | tier | L24 | L28 | L31 | final | last-row logit cos | greedy prefix | +|---|---|---|---|---|---|---|---| +| 7 | INT8 (per-row) | 0.9987 | 0.699 | 0.508 | 0.794 | 0.998850 | 8/12 | +| 7 | INT4 (QuaRot) | 0.9985 | **0.981** | **0.960** | 0.955 | 0.999484 | 8/12 | +| 1032 | INT8 (per-row) | 0.9952 | 0.9953 | 0.9983 | 0.9969 | 0.999916 | 8/16 | +| 1032 | INT4+down | 0.9334 | 0.9313 | — | — | 0.998881 | **0/16** | +| **1032** | **INT8+Hadamard** | **0.9989** | **0.9989** | **0.99972** | **0.99945** | **0.999968** | **16/16** | + +> This may appear to **contradict an earlier prefill-only conclusion** ("both +> INT4 tiers beat INT8 at every layer probe on every frame"). That measurement +> is not wrong — it was taken on a *prefill-only workload at fixed short Se*; +> the verdict is ISL-dependent, and a VLM's production ISL sits in the opposite +> regime. + +So the INT8-vs-INT4 verdict *inverts with sequence length* — at short ISL the +sink row is 1/7 of the tensor and rotation dominates; at long ISL it is 1/1032 +and 4-bit noise dominates. That inversion is the tell that the two tiers were +each solving half the problem. **Rotating at 8 bits solves both and strictly +dominates**, which is why it is the default. + +Cost: **one new device-side pack function** (`quant_int8`) inside the existing +`csrc/kernels/fht_int4.cu`, plus templating its three kernels on the output +width — the norm and the radix-16 register FHT are shared verbatim with the INT4 +path. **No new GEMM**: because the rotation keeps plain per-row scales, the +unmodified `cutlass_int8_rowwise_*` kernels consume the rotated activations +directly. Measured decode **20.96 tok/s vs 19.9** for plain INT8, i.e. no +throughput cost (the FHT rides inside an already-optimized fused norm kernel; +the difference is within the ±3 % process-to-process variance this platform +shows). + +The weight side folds offline (`W_rot = H·W/√K`, `quantize_int8_hadamard`); the +activation side is fused into the norm (`rms_norm_fht_int8_fp16`, +`residual_add_rms_norm_fht_int8_fp16`, `fht_int8_quant_fp16`). The FFN **down** +projection stays plain INT8: K=11008 is not a power of two and its input is the +un-rotated BF16 SiLU output. + +**Why not the alternatives** (principle #17's measured ladder on this backbone): +SmoothQuant reached only 0.641 and outlier-splitting 0.970 on the A4 variant of +this problem, while group-128 / block-scaled schemes need a bespoke GEMM whose +hand-written ceiling on 16-SM Orin measured just 41 TOPS. A per-layer FP16 +fallback would also have worked, but it is checkpoint-specific tuning that +permanently costs throughput — the rotation is free and generalizes. + +### 4.6 SOLVED — the FP16 overflow, via `ffn_down_clamp` + +With a real image the reference's L31 residual reaches **max|x| = 89088**, above +FP16's 65504, so FlashRT stored `inf` and the final RMSNorm turned that row's +logits into `nan`. It affects **both** precision tiers — it is a property of the +residual *dtype*, not of the quantization. + +The first instinct (a BF16 residual stream, ~1 new kernel) was **wrong** — the +answer is a clamp. This port empirically confirmed that a clamp is sufficient. + +**Why a clamp is sufficient** — measured per-layer magnitudes in the bf16 +reference (ISL=1032). The explosion is confined to **exactly one layer**: + +| quantity | L28 | L29 | L30 | **L31** | +|---|---|---|---|---| +| residual | 1616 | 1720 | 2032 | **266240** | +| o_proj output | 76 | 78 | 80 | 1056 | +| down **input** (gu) | 1128 | 1528 | 6080 | **151552** | +| down **output** | 1120 | 1032 | 1880 | **264192** | + +Because the pre-L31 residual is only ~2032, clamping the down **output** at +60000 leaves the residual at ~62000 < 65504. So one `clamp_inplace_fp16` +(already in `flash_rt_kernels`, CUDA-Graph safe) removes the overflow with +**zero new kernels and no dtype change**. + +We do **not** need to clamp the down *input*: ours is BF16 +(`cutlass_int8_silu_gated_bf16out`), whose range absorbs 151552 without issue. +The clamp is applied on every layer because L0-L30 are three +orders of magnitude below it and therefore untouched; cost is 32 extra +elementwise launches (<0.3 % of the decode budget, ~0.7 % of prefill). + +**Result** (ISL=1032, INT8, real image): + +| | before | after | +|---|---|---| +| L31 cosine | `nan` (inf) | **0.998266** | +| final-norm cosine | `nan` | **0.996923** | +| L31 max\|x\| | `inf` | 60160 (saturating at the clamp, as intended) | +| last-row logit cosine | 0.999916 | 0.999916 | +| greedy prefix vs HF | 8/16 | 8/16 | + +Exposed as `ffn_down_clamp` (default 60000, env `FLASHRT_CHAMELEON_DOWN_CLAMP`). + +⚠️ **The clamp did not change the text divergence** (still 8/16). That confirms +the overflow was confined to the sink row's post-L31 residual, which feeds only +that row's final norm — so it was never the cause of the divergence. The +remaining gap is ordinary INT8 error at high-confidence text decisions and is +still open; see §6 for the ranked options. + +Also note the **gate itself was wrong** at first: an absolute +"max|x| < 30000" threshold fails by construction on a backbone whose reference +legitimately runs at 2.6e5. The correct gate is **finiteness**, with clamp +saturation reported as information. + +### 4.7 Measurement-hygiene finding: argmax-over-all-positions is meaningless here + +1024 of the 1032 teacher-forced positions are **image** positions. At those the +model predicts a next token while all 8192 image ids are masked out of the +logits (§3.3), so the winner is an arbitrary low-confidence text token — median +reference top1−top2 gap **0.250** on a logit scale of ~20, versus **0.895** at +the 6 text positions. An unsplit "argmax match = 87.21 %" therefore says almost +nothing about generation quality. The harness now reports image and text +positions separately and gates only on text positions, and additionally +classifies a mismatch as a **BF16 tie** when the reference's top-2 gap is within +one BF16 ULP (58 of the 132 mismatches were ties). + +### 4.8 Two harness traps worth remembering + +* **A forward hook that returns a value replaces the module output.** Using + `dict.setdefault(...)` inside a `register_forward_hook` lambda returns the + stored tensor, which silently substituted a *CPU* tensor for + `model.model.norm`'s output and crashed `lm_head` with a device mismatch. + Always `return None`. +* **Launching on stream 0 while another stream is capturing silently drops the + kernels from the graph.** The first graph-safety run reported + `stale-value: FAIL (frozen)` with cos exactly 1.0000 — not because anything + was baked, but because only the torch ops got captured and none of the `fvk` + kernels did. `decode_step` now takes an explicit `stream` argument. + +## 5. Roofline ladder — predicted vs finally measured (principle #15) + +Decode reads 6.745 G params/token (32 layers 6.476 G + lm_head 0.268 G). **MHA +with 32 KV heads makes the KV cache 4x heavier than a GQA model** — 0.524 MB per +token of context, so 0.55 GB at S=1040 and **2.15 GB at S=4096, where KV would +dominate an int4 tier.** + +| tier | weight floor @173 GB/s | + KV @S~1040 | GEMM-only µbench | predicted total | **finally measured** | +|---|---|---|---|---|---| +| **int8 (+Hadamard, shipped)** | 39.0 ms | +3.2 ms | 47.70 ms | 53.7 ms = 18.6 tok/s | **47.5 ms = 21.07 tok/s** | +| int4 (as built) | 19.5 ms | +3.2 ms | 33.45 ms | 39.5 ms = 25.3 tok/s | not shipped (loses on precision, §4.5) | + +Decode came in **13 % better than predicted** — the prediction charged full price +for attention and the elementwise tail, but `split_kv_bias` (§2.1) and the fused +FHT norm absorb part of it. A prediction that is close *and* slightly pessimistic +is the sign the bottleneck model is right (a large gap in either direction would +mean the model of the bottleneck is wrong, not that there is tuning left). + +**Prefill**: predicted ~255 ms by scaling a measured Se=1214 prefill to +Se=1032; **measured 273.8 ms warm** (within 7 %), of which GEMM is 229.2 ms at +**91 % of the achievable CUTLASS ceiling** (§4.2.2). Image tokenize adds ~53 ms +(PyTorch VQ-GAN; ~27 ms with a 512x512 TRT engine, not built). + +> Superseded numbers, kept so they don't re-mislead: an earlier revision of this +> section predicted **18.6 tok/s** decode and this doc once reported **496 ms** +> prefill and an **84.8 TOPS** GEMM target. Current values: **21.07 tok/s**, +> **273.8 ms**, and a **64.4 TOPS** achievable ceiling. See §4.2 for why the +> 496 ms was a cold-start artifact and §4.2.2 for the ceiling correction. + +## 6. Ranked lever menu + +### Precision — status: closed + +| # | lever | outcome | +|---|---|---| +| 1 | **W8A8 + Hadamard (QuaRot at 8 bits)** | **DONE — this closed it.** greedy 8/16 → **16/16**, worst layer 0.9946 → 0.9986, last-row logit 0.999916 → 0.999968, at no throughput cost. Default tier. | +| 2 | `ffn_down_clamp` | **DONE** — removed the L31 FP16 `inf` (§4.6) | +| ~~3~~ | ~~Tier-3 FP16 fallback for L31~~ | **not needed** — the rotation fixed the same layer at 8 bits. A per-layer precision fallback is checkpoint-specific tuning and costs throughput permanently; prefer the quantization method. | +| ~~4~~ | ~~AWQ / SmoothQuant per-K smoothing~~ | **not needed** — and principle #17's measured ladder on this backbone puts smoothing (0.641) far below rotation (0.9914). Kept only as a fallback if a future checkpoint defeats rotation. | +| 5 | ISL-adaptive tier selection | **obsolete** — W8A8+Hadamard wins at both short and long ISL, so there is nothing to switch between | +| ~~6~~ | ~~BF16 residual stream~~ | **superseded by the clamp** — would have cost a new `qk_norm_rope_fused_bf16` kernel to fix what one existing elementwise kernel already fixes | + +**Decode (M=1, weight-bandwidth-bound)** + +| # | lever | predicted | effort | status | +|---|---|---|---|---| +| 1 | `use_int4` / `use_int4_down` | **1.43×** (measured, not 2× — §2.2) | trivial, already built | Phase 2 | +| 2 | `cutlass_sm80_int4_rowwise_t64x128.cu` + `M<=64` dispatcher | int4 33.45 → ~27 ms = **+20 %** | ~150 lines + 1 CMake line | Phase 3 — **justified by §2.2** | +| 3 | `split_kv_bias = 4` (`num_sms=64`) | attention **1.44×** = +1.7 ms/token (+3–5 %), more at long S | 1 Python arg | Phase 1 — **measured (§2.1)** | +| 4 | Per-position decode CUDA graph | 0–15 % throughput, −6 s startup | medium | Phase 4 | +| 5 | INT8 Q/K/V/O at 67 % of ceiling (small-N tail: 4096/128 = 32 tiles on 16 SMs) | up to +12 % if it reached 100 % | high (hand GEMV) | open | +| 6 | Devpos kernel + fp16 seqused-splitkv FA2 → *one* decode graph | 0 % throughput; removes capture cost + `max_new_tokens` cap | high (1 `.cu` + FA2 rebuild) | deferred | +| 7 | Reduced lm_head (drop rows 4..8195) | +0.5–1 % | low | deferred | +| 8 | INT8 KV cache | +3.8 % @1040, **+12 % @4096** | high (needs a 32Q/32KV variant; break-even measured) | S≥4096 only | +| ~~9~~ | ~~M=1 up-projection split~~ | ~~+7 %~~ → **measured 0.9 %, net ≈0** | — | **DEAD (§2.2)** | +| 10 | Speculative decode | 1.5–2× | N/A — no draft model | — | + +**Prefill / TTFT (FLOPs-bound — levers do not transfer)** + +| # | lever | predicted | status | +|---|---|---|---| +| 1 | `use_int4` / `use_int4_down` | LLM ~265 → ~165 ms (**−38 % TTFT**) | already built | +| 2 | 512×512 TRT VQ-GAN engine | 53 → 27 ms (−9 % TTFT) | `scripts/build_vqgan_trt.py`, inputs present in `original_tokenizers/` | +| 3 | GEMM-util / elementwise fusion | **≈0** — measured spent at 74 % of the 84.8 TOPS mma peak | closed | +| 4 | Per-Se prefill CUDA graph | 0–5 %, and a full capture per new prompt length; `Se` cannot be bucketed (padding poisons the KV cache) | **reject for a VLM** | + +## 7. Dead-ends (measured — do not re-walk) + +| direction | result | one-line reason | +|---|---|---| +| FA2 split-KV with the real `num_sms=16` | **bit-identical, 1.05×** | `32 >= 0.8*32` → `num_splits=1`; the heuristic disables itself at 32 Q heads | +| M=1 up-projection split (`bf16out` + `silu_mul`) | 13 µs/layer ≈ 0.9 %, net ≈0 | both GEMMs already bandwidth-bound; the extra launch + bf16 round trip cancels it | +| INT4 at M=1 expecting 2× | **1.43×** | single 128×128 tile, no `M<=64` dispatcher | +| **`use_int4` / `use_int4_down` as the VLM default** | **L24/L28 cosine 0.933 vs INT8's 0.995; greedy prefix 0/16 vs 8/16** | at production ISL the sink row is 1/1032 of the sequence, so INT8's per-row damage is diluted and 4-bit noise on the other 1031 rows dominates (§4.5). INT4 still wins at short ISL — the verdict is ISL-dependent | +| **unsplit argmax match as a precision metric** | "87.21 %" says nothing | 1024/1032 teacher-forced positions are image positions whose logits are fully masked → arbitrary low-confidence winners (§4.7) | +| forward hook capturing via `dict.setdefault` | CPU/CUDA device crash in `lm_head` | a hook returning non-None **replaces** the module output | +| launching pipeline kernels on `stream=0` during graph capture | `stale-value: FAIL (frozen)`, cos exactly 1.0000 | kernels on the default stream are silently *not* recorded; only the torch ops were captured | +| `attention_mha_causal_fp16` for decode | silently wrong | top-left-aligned causal mask; at `S_q=1` only column 0 survives | +| stock `from_pretrained(..., state_dict=...)` | `ValueError` in 4.57 | use naked ctor + `load_state_dict` | +| stock transformers loading this ckpt unmodified | shape mismatch on 128 tensors | `ChameleonLayerNorm` wants `(32,128)`, ckpt has `(1,128)` | +| **measuring prefill on the first call** | 460-496 ms vs 273.8 ms warm | 1.68x cold-start penalty from CUTLASS workspace `cudaMalloc` + JIT; produced a phantom "1.9x prefill gap" | +| **roofline probe without a clock pre-ramp** | first shape measured reads 28.3 TOPS vs 61.1 warm | Orin DVFS ramps 306 -> 1300 MHz; per-shape TOPS ascend in measurement order | +| clamping the down output on all 32 layers | 6.23 ms (2.2 %) for no effect on 28 of them | magnitudes grow monotonically with depth; L28 is 37x below the clamp | +| trusting 84.8 TOPS as the GEMM target | it is the **raw mma issue rate**, not achievable | CUTLASS peaks at 64.4 TOPS big-square on this part; we are at 91 % of *that* | +| single-stream reduce as a bandwidth probe | 99 GB/s | undersaturated; use a real weight-streaming kernel (173 GB/s) | +| int8 `sum(dtype=int64)` as a read-BW probe | 9.3 GB/s | ALU-bound reduction, not a bandwidth measurement | + +## 8. Reproduction + +```bash +# Gate 1: correctness vs the HF bf16 reference + graph safety + fp16 health +PYTHONPATH=. python3 scripts/chameleon_orin_check.py \ + --checkpoint /path/to/Chameleon_7B_mGPT \ + --image FlashRT.png --prompt "Describe this image." --steps 16 +# ... --text-only fast smoke, no image +# ... --int4 / --int4-down alternative tiers +# ... --vq-fp16-argmin measure VQ index drift instead of avoiding it + +# M=1 decode roofline (no checkpoint) — the "do we need a GEMV?" gate +# (roofline probe script: scripts/bench/orin_int8_roofline.py) +python3 scripts/bench/orin_int8_roofline.py --decode + +# large-M prefill roofline at the production shape +python3 scripts/bench/orin_int8_roofline.py --M 1032 +``` + +Minimal use: + +```python +from flash_rt.frontends.torch.chameleon_rtx_sm87 import ChameleonTorchFrontendRtxSm87 +f = ChameleonTorchFrontendRtxSm87("/path/to/Chameleon_7B_mGPT") # int8+hadamard +f.set_prompt("Describe this image.", images=[pil_img]) +print(f.generate(max_new_tokens=32)) +``` + +`load_model(config="chameleon")` deliberately raises with this snippet: it is a +chat VLM (`set_prompt` + `generate`), not the VLA `predict()` surface. + +**Build** (the INT8 FHT kernels ride in the existing `ENABLE_SM80_INT8_CUTLASS` +source list, so no CMake change is needed): + +```bash +cmake -B build -S . -DGPU_ARCH=87 -DFA2_ARCH_NATIVE_ONLY=ON \ + -DFA2_HDIMS='128;256' -DFA2_DTYPES='fp16;bf16' +cmake --build build --target flash_rt_kernels -j6 +``` + +## 9. Constructor parameters (classified: required / recommended / experimental / informational) + +`ChameleonTorchFrontendRtxSm87(checkpoint_dir, **kwargs)`. There is no CLI/server +entry yet (out of scope this round), so these are the frontend kwargs. + +**Required** + +| param | notes | +|---|---| +| `checkpoint_dir` | path to the Chameleon-7B checkpoint. Backbone dims and the special token ids are **hard-asserted** against its `config.json` | + +**Recommended (safe defaults; tune for deployment)** + +| param | default | notes | +|---|---|---| +| `max_seq` | 2048 | sizes the KV cache (`2 × 32 × max_seq × 4096 × 2 B` = 2.15 GB at 2048) and the RoPE tables. Hard-checked against `max_position_embeddings=4096`. `S = 1 + n_img*1026 + n_text + 1` | +| `use_hadamard` | **`True`** | the W8A8+QuaRot tier. **Do not disable in production** — plain per-row INT8 reproduces only 8/16 reference tokens (§4.5). Kept switchable for A/B only | +| `split_kv_bias` | `4` | multiplies the `num_sms` passed to FA2 so split-KV actually engages at 32 Q heads (§2.1). `1` disables | +| `ffn_down_clamp` | `60000` | **correctness requirement**, not a knob (§4.6). Env: `FLASHRT_CHAMELEON_DOWN_CLAMP` | +| `ffn_down_clamp_last_n` | `4` | layers to clamp, counted from the end. `32` = all layers (safe but costs 2.2 % of prefill). Env: `FLASHRT_CHAMELEON_DOWN_CLAMP_LAST_N` | +| `vq_argmin_fp32` | `True` | fp32 codebook distance/argmin; costs <0.1 ms and lifts index match 98.1 % → 99.0 % (§2.4) | +| `free_fp16_weights` | `True` | drop the 13 GB fp16 originals after quantization | + +**Experimental (off by default)** + +| param | default | notes | +|---|---|---| +| `use_int4` | `False` | QuaRot W4A4 on the six K=4096 projections. Wins only at very short ISL; **below `use_hadamard` at production ISL** (§4.5) | +| `use_int4_down` | `False` | additionally int4 the FFN down via block-H128. **Not recommended** — worst tier measured at production ISL (0/16 greedy) | +| `probe_layers` | `None` | list of layer indices to snapshot post-residual hidden states for `snapshot_probe()`; zero cost when `None` | + +**Informational** + +| param | notes | +|---|---| +| `precision_tier` / `precision_spec()` / `get_model_info()` | report the resolved configuration; `timing` carries `prompt_ms` / `prefill_ms` / `decode_tok_s` | + +**Not accepted**: `use_fp8`, `use_fp4`, +`use_fp8_attn`, `use_awq_v_proj`, `num_views`, `action_dim`, +`action_chunk_size`, `state_dim` — SM87 has no FP8/FP4 tensor cores and this is +not a VLA. Unknown kwargs are swallowed by `**_ignored`. + +## 10. Tier status (all four verified to run and produce finite logits) + +| tier | flag | verdict | +|---|---|---| +| **int8+hadamard** | default | **production** — Gate 1 PASS, 16/16 greedy match, 21.07 tok/s | +| int8 plain | `use_hadamard=False` | works; loses the outlier conditioning (8/16 greedy) — kept for A/B | +| int4 (QuaRot) | `use_int4=True` | works; best at very short ISL, below int8+hadamard at production ISL | +| int4+down | `use_int4_down=True` | works but **not recommended** — worst at production ISL (§4.5) | + +## 11. Hardware gate and generation boundary + +- `ChameleonTorchFrontendRtxSm87` fail-fasts on non-Orin hardware: it checks + `torch.cuda.get_device_capability()` before checkpoint loading / weight + quantization / large CUDA allocation and raises on anything other than SM87. + The documented development override is `FLASHRT_CHAMELEON_SM87_FORCE=1` + (skips the probe only; kernels still need the real hardware at runtime). +- `generate(...)` defines `max_new_tokens` explicitly: negative values raise + `ValueError`, zero returns an empty result (no prefill, no decode), and + values above remaining `max_seq` capacity are clipped with a warning. diff --git a/docs/chameleon_thor_sm110.md b/docs/chameleon_thor_sm110.md new file mode 100644 index 000000000..dec09a4a0 --- /dev/null +++ b/docs/chameleon_thor_sm110.md @@ -0,0 +1,343 @@ +# Chameleon-7B on Thor SM110 + +**Platform**: Jetson AGX Thor (SM110, aarch64) · CUDA 13.0 · transformers 4.43+ +**Model**: Standalone Chameleon-7B (LLM backbone + VQGAN image tokenizer, **no ActionHead / ActionVAE**) +**Production path**: All 32 layers with **runtime dynamic per-tensor FP8** (implemented in the Chameleon-specific `flash_rt/models/chameleon/pipeline_thor.py::chameleon_forward`) + generic eager Chameleon VQGAN default + cuBLASLt per-shape autotune + L31 selective clamp; TensorRT VQGAN is explicit opt-in only +**Version**: v1.4 (2026-08, added KV-cache incremental decode `generate_greedy`: 30.4 tok/s, token-exact vs full-prefix recompute oracle) + +--- + +## 0. Summary + +- **Asset path**: `/path/to/Chameleon_7B_mGPT` (note the actual directory name is `mGPT`, not `mGP`). Contains weight shards, tokenizer, and `original_tokenizers/vqgan.{yaml,ckpt}`. +- **HF direct loading fails**: The current `transformers` `ChameleonForConditionalGeneration.from_pretrained` errors or silently misloads on this checkpoint because `q_norm`/`k_norm` shapes are legacy `[1,128]` (not the newer `[32,128]`). **Workaround**: The production path uses FlashRT's own declarative `WeightLoader` (bypassing HF `from_pretrained` entirely); alternatively, as in `scripts/check_chameleon_thor_precision.py`, use the bare `ChameleonForConditionalGeneration` constructor + `load_state_dict(strict=False)` for the HF reference model (the script imports `ChameleonForConditionalGeneration` directly from `transformers`; use `--skip-hf` to skip the HF reference comparison). +- **Precision validated (real images, not synthetic token ids)**: + - FlashRT FP16 vs HF BF16 (last-token logits cosine, after mask_image_logits): **0.9999997**, greedy next-token exact match. + - FlashRT dynamic FP8 vs FlashRT FP16: **0.99999999**, greedy next-token exact match, top-10 overlap 1.0. +- **VQGAN backend policy (framework positioning)**: For **generic/standard Chameleon**, FlashRT preserves framework generality — VQGAN **defaults to eager** Chameleon tokenization (`use_trt_vqgan=False`), with no default dependency on TensorRT engines, ensuring the framework's own capabilities run independently. **If the deployment environment has compatible TRT engines, explicitly opt in** (`use_trt_vqgan=True` or script `--use-trt-vqgan`; measured VQGAN 74.9→17.3 ms, TRT E2E ~121 ms vs eager ~190 ms). Output JSON records the actual backend (`eager`/`trt`). +- **Latest end-to-end performance (real image `hand_1.jpg`, prompt "Describe the image.", target_size=512, stage-aware benchmark, including §4.11 fused kernels + §4.12 FA4)**: + + | Scope | VQGAN backend | FlashRT FP8 p50/mean | Notes | + |---|---|--:|---| + | Default E2E | eager | **~190 ms** | VQGAN 74.9 ms dominates; eager bottleneck is VQGAN without TRT | + | Explicit opt-in E2E | TRT | **121.1 / 121.2 ms** | TRT VQGAN 17.5 ms + transformer 103.5 ms (with FA4) | + | transformer-prefill-only (FA4) | eager ids reused | **101.9 / 102.0 ms** | HF-comparable, excludes VQGAN, 50 iter | + + > **2026-08-05 re-measurement (single hot window, 20 iter, `benchmarks/chameleon_thor_latency.py`)**: + > transformer-only FA4 off **111.2 ms** / FA4 on **104.2 ms** (−7.0); E2E eager+FA4 **177.3 ms**; + > E2E TRT+FA4 **120.2 ms**. Differences from the table above are within thermal noise (±5%); + > PR-facing docs (`docs/chameleon_usage.md`, `docs/benchmark_comparison.md`, USAGE.md) use the + > re-measured values. + + Roofline conclusion (see §4.10-4.12): At Se=1056/1072 the theoretical workload is ~**14.3-14.5 TFLOP**; at 240 TFLOP/s the optimistic compute floor is ~**59-60 ms**. Per-shape GEMM micro-benchmarks (§4.11) confirm GEMM tactics are already near the measured Thor ceiling (32-layer GEMM-only ≈61.9 ms), so the gap to floor is primarily non-GEMM work. §4.11 fused RMSNorm/SwiGLU+amax (117.5→110.9 ms), §4.12 FA4 attention (110.9→**101.9 ms**, 58.3% of 240 TFLOP/s, **1.71× floor**). Remaining headroom is in O-projection quantization (no natural fusion point) and KV-cache incremental decode (landed in §4.13). +- **FA4 attention (explicit opt-in)**: Following upstream PR [`flashrt-project/FlashRT#163`](https://github.com/flashrt-project/FlashRT/pull/163) (GROOT N1.7 Thor NVFP4+FA4, single-view 51.6→29.9 ms, 1.70×). Chameleon shapes (Se=1056, 32 heads, HD=128, causal) measured FA4 vs in-repo CUTLASS causal FMHA: **2.75× faster** (450.5→163.9 µs/layer), output cos=0.99999994; integrated transformer-only FP8 **-8.4 ms**. Requires `pip install .[thor-fa4]` (nvidia-cutlass-dsl==4.5.1 + quack-kernels==0.4.1), enabled via `FLASHRT_CHAMELEON_FA4_ATTN=1` or constructor arg `use_fa4_attn=True`, with automatic CUTLASS FMHA fallback when unavailable. +- **KV-cache incremental decode (2026-08, see §4.13)**: `generate_greedy` now uses one prefill + M=1 incremental decode (`chameleon_decode_step`), steady-state **30.4 tok/s** (32.9 ms/token), wall-clock ~**2.8×** vs full-prefix recompute; token-exact vs eager full-prefix recompute oracle (32-token generation 38/38). Added bottom-right aligned causal FMHA symbol `fmha_fp16_causal_br` (decode SQ=1`=2, override via `eos_token_id`) or `max_seq`. Requires + the dynamic-FP8 path (`use_fp8=True`); the eager full-recompute path is + retained as `_generate_greedy_recompute` for oracle comparisons. + Benchmark via `scripts/bench_chameleon_thor.py --generate-greedy N`. +- The TRT VQGAN path uses a square `target_size×target_size` bicubic + resize while eager uses aspect-preserving `var_center_crop` — token + counts can differ slightly between backends (expected behavior + difference, not a bug). + +## 9. VQ-VAE implementation + +The eager image tokenizer uses the Apache-2.0 Transformers +`ChameleonVQVAE` implementation. FlashRT loads only the +`model.vqmodel.*` tensors from the user-provided Transformers checkpoint; +no separately licensed VQGAN source is vendored or included in the FlashRT wheel. diff --git a/docs/stable_api.md b/docs/stable_api.md index 0806d56df..8efe602ed 100644 --- a/docs/stable_api.md +++ b/docs/stable_api.md @@ -29,7 +29,7 @@ def load_model( autotune: int = 3, # 0=off, 3=default, 5+=thorough recalibrate: bool = False, weight_cache: bool = True, # JAX only - config: str = "pi05", # "pi05" | "pi0" | "groot" | "groot_n17" | "pi0fast" | "motus" | "wan22_ti2v_5b" | "cosmos3_video" | "cosmos3_edge" + config: str = "pi05", # "pi05" | "pi0" | "groot" | "groot_n17" | "pi0fast" | "motus" | "wan22_ti2v_5b" | "cosmos3_video" | "cosmos3_edge" | "chameleon" device=None, # reserved # Pi0-FAST-specific: decode_cuda_graph: bool = False, @@ -162,6 +162,13 @@ Returns a `VLAModel` wrapping the appropriate frontend for the detected frontend; `rtx_sm89` resolves directly to its dedicated SM89 frontend. `use_fp16=True, use_fp8=False` requests the explicit RTX reference frontend for the selected hardware. +- `config="chameleon"` is a chat-style VLM and is not served through + `load_model`'s VLA wrapper. Calling `load_model(config="chameleon")` + raises `NotImplementedError` with direct-construction instructions. + Construct the frontend explicitly: + `ChameleonTorchFrontendRtxSm87` (Jetson Orin SM87) or + `ChameleonTorchFrontendThor` (Jetson Thor SM110). + See `docs/chameleon_usage.md`. ### `flash_rt.VLAModel` @@ -296,6 +303,8 @@ based on `use_fp8` / `use_fp16`; `rtx_sm89` resolves directly to the dedicated SM89 frontend class. Wan2.2 TI2V-5B is registered for `(config="wan22_ti2v_5b", framework="torch", arch="rtx_sm120")`. +Chameleon-7B is registered for `(config="chameleon", framework="torch", +arch in {"rtx_sm87", "thor"})`. ### `_PIPELINE_MAP` diff --git a/examples/thor/README.md b/examples/thor/README.md index 5e1e2124b..1a87b54a9 100644 --- a/examples/thor/README.md +++ b/examples/thor/README.md @@ -91,6 +91,25 @@ Quality measurements, the `max_pixels` resolution knob, and the per-projection `wq_overrides` sweep surface are in [`docs/qwen3_vl_thor.md`](../../docs/qwen3_vl_thor.md). +## Chameleon-7B (multimodal chat) + +`chameleon_quickstart.py` runs standalone Chameleon-7B (image + text) with +the dynamic per-tensor FP8 backbone and CUDA-graph prefill: + +```bash +python examples/thor/chameleon_quickstart.py \ + --checkpoint /path/to/Chameleon_7B_mGPT \ + --image /path/to/image.jpg \ + --prompt "Describe the image." \ + --benchmark +``` + +Add `--use-trt-vqgan` when compatible TensorRT VQ-GAN engines exist +(`scripts/build_vqgan_trt.py`), and `FLASHRT_CHAMELEON_FA4_ATTN=1` for the +optional FA4 attention fast path. Measured ~190 ms E2E prefill (eager +VQGAN), ~120 ms with TRT VQGAN + FA4, and ~30 tok/s incremental decode. +Full details in [`docs/chameleon_usage.md`](../../docs/chameleon_usage.md). + ## Thor VLA performance ### Precision (Pi0.5, 2-view LIBERO) diff --git a/examples/thor/chameleon_quickstart.py b/examples/thor/chameleon_quickstart.py new file mode 100644 index 000000000..c071f28da --- /dev/null +++ b/examples/thor/chameleon_quickstart.py @@ -0,0 +1,95 @@ +#!/usr/bin/env python +"""Chameleon-7B (Thor sm_110) quickstart. + +Runs the standalone Chameleon-7B image+text frontend on a real image and +reports prefill latency. The frontend is a direct-instantiation class +(``ChameleonTorchFrontendThor``, registered in ``_PIPELINE_MAP`` but not +dispatched by ``flash_rt.load_model`` — same pattern as Qwen3-VL / Nex-N2 +/ LingBot). + +Build first (one shared module — the Chameleon kernels live inside +flash_rt_kernels; FA4 is optional): + + cmake -B build -S . -DGPU_ARCH=110 + cmake --build build -j + pip install -e ".[torch]" # add ,thor-fa4 for the FA4 fast path + +Run: + + python examples/thor/chameleon_quickstart.py \ + --checkpoint /path/to/Chameleon_7B_mGPT \ + --image /path/to/hand_1.jpg \ + --prompt "Describe the image." + +Expected on Thor (dynamic FP8, CUDA graph, target_size=512, eager VQGAN): +~190 ms/prefill E2E; with --use-trt-vqgan (engines present) ~120 ms; +with FA4 enabled (FLASHRT_CHAMELEON_FA4_ATTN=1) the transformer part drops +to ~104 ms. The script prints the actual VQGAN backend and FA4 status. +""" +import argparse +import time + +import torch + +from flash_rt.frontends.torch.chameleon_thor import ChameleonTorchFrontendThor +from flash_rt.hardware.thor import fa4_backend + + +def main() -> None: + ap = argparse.ArgumentParser(description="Chameleon-7B Thor quickstart") + ap.add_argument("--checkpoint", required=True, + help="Chameleon-7B dir (model-*-of-*.safetensors + config.json)") + ap.add_argument("--image", required=True, help="real image path (jpg/png)") + ap.add_argument("--prompt", default="Describe the image.") + ap.add_argument("--target-size", type=int, default=512) + ap.add_argument("--use-trt-vqgan", action="store_true", + help="Use TensorRT VQGAN if compatible engines exist " + "(recommended when engines are available; default is eager VQGAN)") + ap.add_argument("--use-fp16", action="store_true", + help="Use the FP16 reference path instead of dynamic FP8") + ap.add_argument("--no-graph", action="store_true", help="Disable CUDA Graph") + ap.add_argument("--iters", type=int, default=10, help="timed replays") + ap.add_argument("--benchmark", action="store_true", + help="report wall-clock prefill latency (P50)") + args = ap.parse_args() + + from PIL import Image + + image = Image.open(args.image).convert("RGB") + fa4 = fa4_backend.is_available() + print(f"[chameleon] FA4 attention available: {fa4} ({fa4_backend.status()})" + f"{'' if fa4 else ' <-- CUTLASS FMHA will be used; pip install .[thor-fa4]'}") + + fe = ChameleonTorchFrontendThor( + args.checkpoint, + use_fp8=not args.use_fp16, + use_cuda_graph=not args.no_graph, + target_size=args.target_size, + use_trt_vqgan=args.use_trt_vqgan, + ) + print(f"[chameleon] VQGAN backend: {fe.vqgan_backend} " + f"FA4 active: {fe.fa4_attn_active}") + + out = fe.prefill(args.prompt, [image]) + print(f"[chameleon] Se={out['Se']} (real_len={len([i for i in out['input_ids'] if i != 1])}) " + f"logits={tuple(out['logits'].shape)}") + top = int(torch.argmax(out["logits"]).item()) + print(f"[chameleon] greedy next-token id: {top}") + + if args.benchmark: + ts = [] + for _ in range(args.iters): + torch.cuda.synchronize() + t0 = time.perf_counter() + fe.prefill(args.prompt, [image]) + torch.cuda.synchronize() + ts.append((time.perf_counter() - t0) * 1000.0) + ts.sort() + p50 = ts[len(ts) // 2] + print(f"[chameleon] prefill P50: {p50:.1f} ms over {args.iters} iters " + f"(wall-clock, includes VQGAN; fp8={not args.use_fp16}, " + f"graph={not args.no_graph}, vqgan={fe.vqgan_backend})") + + +if __name__ == "__main__": + main() diff --git a/flash_rt/api.py b/flash_rt/api.py index 13d4e20e5..c3ebe5db5 100644 --- a/flash_rt/api.py +++ b/flash_rt/api.py @@ -510,6 +510,24 @@ def load_model(checkpoint, framework="torch", num_views=2, autotune=3, "use_fp4_decoder/use_fa4 are unsupported with framework=" "'jetson_pi'; use the Thor torch FP4/FA4 frontend instead") + # Chameleon-7B is a chat-style VLM, not a VLA: its frontends expose + # set_prompt(...) + generate() rather than predict(images, ...), so + # VLAModel's result['actions'] contract does not apply. Registered in + # _PIPELINE_MAP for discoverability only. + if config == "chameleon": + raise NotImplementedError( + "config='chameleon' is a chat-style VLM and is not served through " + "load_model's VLA wrapper. Construct it directly:\n" + " from flash_rt.frontends.torch.chameleon_rtx_sm87 import " + "ChameleonTorchFrontendRtxSm87 # Jetson Orin SM87\n" + " from flash_rt.frontends.torch.chameleon_thor import " + "ChameleonTorchFrontendThor # Jetson Thor SM110\n" + " f = ChameleonTorchFrontendRtxSm87('/path/to/Chameleon_7B_mGPT')\n" + " f.set_prompt('Describe this image.', images=[img])\n" + " print(f.generate(max_new_tokens=32))\n" + "See docs/chameleon7b_rtx_sm87.md, docs/chameleon_thor_sm110.md " + "and docs/chameleon_usage.md.") + if framework == "jetson_pi": if config not in ("pi0", "pi05", "llm", "mllm"): raise ValueError( diff --git a/flash_rt/frontends/torch/_chameleon_quant.py b/flash_rt/frontends/torch/_chameleon_quant.py new file mode 100644 index 000000000..0a1a12296 --- /dev/null +++ b/flash_rt/frontends/torch/_chameleon_quant.py @@ -0,0 +1,251 @@ +"""Low-bit weight quantizers for the Chameleon-7B GEMMs (Orin SM87). + +Checkpoint-agnostic and frontend-agnostic: every function takes plain tensor +lists and returns plain tensors, so the same code serves any Chameleon +frontend. + +Layout contract — get this wrong and the GEMM silently returns garbage: + +* The declarative weight spec hands us per-projection FP16 weights in + ``[K, N]`` row-major (``Cat``/``FusedGateUp`` followed by ``T()``). +* The CUTLASS SM80 rowwise GEMMs consume the **B** operand as ``[N, K]`` + ColumnMajor with an ``[N]`` FP32 ``RowBroadcast`` scale, so we transpose to + ``[N, K]`` and quantize each of the N output rows symmetrically. +* The **A** operand is ``[M, K]`` RowMajor with ``M`` consecutive FP32 + ``ColBroadcast`` scales, produced at runtime by the fused norm/quant kernels. + +INT4 additionally applies the QuaRot rotation ``W_rot = (H_K @ W)/sqrt(K)`` +offline; the matching activation rotation happens online in the fused +RMSNorm+FHT kernels. Values are packed 2/byte with the even index in the low +nibble (``cutlass::int4b_t`` order). +""" + +from __future__ import annotations + +import logging +from typing import Dict, List, Tuple + +import torch + +logger = logging.getLogger(__name__) + +INT8_QUANT_MAX = 127.0 +INT8_QUANT_EPS = 1e-12 +INT4_QUANT_MAX = 7.0 +INT4_QUANT_EPS = 1e-10 + +#: The seven per-layer GEMMs, in the order the pipeline consumes them. +PROJECTIONS = ("q", "k", "v", "o", "gate", "up", "d") + +#: Projections whose K == D (a power of two) and so can take a full-width +#: Hadamard rotation. ``d`` has K == Dff == 11008 and needs block-H128. +INT4_POW2_PROJECTIONS = ("q", "k", "v", "o", "gate", "up") + + +def split_fused_projections(qkv_w: List[torch.Tensor], + gu_w: List[torch.Tensor], + o_w: List[torch.Tensor], + d_w: List[torch.Tensor], + *, D: int, Dff: int) -> Dict[str, List[torch.Tensor]]: + """Materialize the seven per-projection ``[K, N]`` weight lists. + + ``qkv_w[li]`` is ``[D, 3D]`` row-major and ``gu_w[li]`` is ``[D, 2*Dff]``: + after ``Cat(dim=0) -> T().contiguous()`` the fused row stride is ``3D`` + (resp. ``2*Dff``), so a *byte-offset* split would read the projections + column-interleaved. Only a ``fused[:, lo:hi].contiguous()`` slice recovers + the original ``q_proj`` / ``k_proj`` / ``v_proj`` blocks. + """ + out: Dict[str, List[torch.Tensor]] = {k: [] for k in PROJECTIONS} + for li, (qkv, gu) in enumerate(zip(qkv_w, gu_w)): + if tuple(qkv.shape) != (D, 3 * D): + raise RuntimeError( + f"layer {li}: expected fused qkv_w {(D, 3 * D)}, " + f"got {tuple(qkv.shape)}") + if tuple(gu.shape) != (D, 2 * Dff): + raise RuntimeError( + f"layer {li}: expected fused gu_w {(D, 2 * Dff)}, " + f"got {tuple(gu.shape)}") + out["q"].append(qkv[:, 0:D].contiguous()) + out["k"].append(qkv[:, D:2 * D].contiguous()) + out["v"].append(qkv[:, 2 * D:3 * D].contiguous()) + out["gate"].append(gu[:, 0:Dff].contiguous()) + out["up"].append(gu[:, Dff:2 * Dff].contiguous()) + out["o"] = list(o_w) + out["d"] = list(d_w) + return out + + +def quantize_per_row_int8(w_kn: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]: + """Per-output-row symmetric INT8 of a ``[K, N]`` FP16 weight. + + Returns ``(q [N, K] int8, scale [N] fp32)``, both contiguous on the + weight's device. + """ + w_f32 = w_kn.float().transpose(0, 1).contiguous() # [N, K] + scale = torch.clamp(w_f32.abs().amax(dim=1) / INT8_QUANT_MAX, + min=INT8_QUANT_EPS).float().contiguous() # [N] + q = torch.clamp(torch.round(w_f32 / scale[:, None]), + -127, 127).to(torch.int8).contiguous() # [N, K] + return q, scale + + +def hadamard_gpu(n: int, device="cuda") -> torch.Tensor: + """Unnormalised Sylvester Hadamard ``H_n`` (fp32). ``n`` must be a power of 2.""" + H = torch.ones(1, 1, dtype=torch.float32, device=device) + while H.shape[0] < n: + H = torch.cat([torch.cat([H, H], 1), torch.cat([H, -H], 1)], 0) + return H + + +def _pack_int4_rows(w_rot_nk: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]: + """Per-output-row symmetric INT4 of an already-rotated ``[N, K]`` fp32 weight.""" + scale = (w_rot_nk.abs().amax(1) / INT4_QUANT_MAX).clamp_min(INT4_QUANT_EPS) + q = torch.clamp(torch.round(w_rot_nk / scale[:, None]), -7, 7).to(torch.int8) + lo = (q[:, 0::2] & 0xF).to(torch.uint8) + hi = (q[:, 1::2] & 0xF).to(torch.uint8) + return (lo | (hi << 4)).contiguous(), scale.float().contiguous() + + +class QuantizedWeights: + """Owns the quantized tensors and exposes the pointer dicts the pipeline wants. + + ``ptr[proj][li]`` is the packed weight pointer and ``scale_ptr[proj][li]`` + the matching ``[N]`` FP32 scale pointer. The tensors are kept alive in + ``_store`` for the object's lifetime — dropping this object invalidates + every pointer, including any already baked into a captured CUDA graph. + """ + + def __init__(self) -> None: + self._store: List[torch.Tensor] = [] + self.ptr: Dict[str, List[int]] = {k: [] for k in PROJECTIONS} + self.scale_ptr: Dict[str, List[int]] = {k: [] for k in PROJECTIONS} + self.precision: Dict[str, str] = {} + + def bytes(self) -> int: + return sum(t.numel() * t.element_size() for t in self._store) + + def _set(self, proj: str, li: int, q: torch.Tensor, s: torch.Tensor) -> None: + self._store.append(q) + self._store.append(s) + while len(self.ptr[proj]) <= li: + self.ptr[proj].append(0) + self.scale_ptr[proj].append(0) + self.ptr[proj][li] = q.data_ptr() + self.scale_ptr[proj][li] = s.data_ptr() + + def _drop(self, ptr: int) -> None: + if ptr: + self._store = [t for t in self._store if t.data_ptr() != ptr] + + +def quantize_int8_all(proj: Dict[str, List[torch.Tensor]], + *, num_layers: int) -> QuantizedWeights: + """INT8-quantize all seven projections for all layers.""" + qw = QuantizedWeights() + for li in range(num_layers): + for key in PROJECTIONS: + w = proj[key][li] + if w.dtype != torch.float16: + w = w.to(torch.float16) + qw._set(key, li, *quantize_per_row_int8(w)) + qw.precision.update({k: "int8" for k in PROJECTIONS}) + logger.info("INT8 quantized %d LLM GEMM weights (%d layers x %d proj), %.2f GB", + num_layers * len(PROJECTIONS), num_layers, len(PROJECTIONS), + qw.bytes() / 2 ** 30) + return qw + + +def quantize_int4_quarot(qw: QuantizedWeights, + proj: Dict[str, List[torch.Tensor]], + *, num_layers: int, D: int, + include_down: bool = False) -> QuantizedWeights: + """Replace the INT8 tensors of the rotatable projections with QuaRot INT4. + + ``include_down`` additionally rotates the FFN down projection with a + **block-diagonal** ``H_128`` because its K (11008) is not a power of two. + Unrotated per-row W4A4 measures cos 0.9494 and fails the gate, so the + rotation is mandatory rather than an optimization. + """ + Hm = hadamard_gpu(D) / (float(D) ** 0.5) + keys = list(INT4_POW2_PROJECTIONS) + Hb = None + if include_down: + keys.append("d") + Hb = hadamard_gpu(128) / (128.0 ** 0.5) + + for li in range(num_layers): + for key in keys: + w = proj[key][li].to(torch.float32) # [K, N] + if key == "d": + Kd = w.shape[0] + w_rot = (Hb.t() @ w.reshape(Kd // 128, 128, -1) + ).reshape(Kd, -1).t().contiguous() # [N, Kd] + else: + if w.shape[0] != D: + raise RuntimeError( + f"{key} layer {li}: expected K={D}, got {tuple(w.shape)}") + w_rot = (Hm @ w).t().contiguous() # [N, D] + qw._drop(qw.ptr[key][li]) + qw._set(key, li, *_pack_int4_rows(w_rot)) + qw.precision[key] = "int4" + del w, w_rot + del Hm, Hb + torch.cuda.empty_cache() + logger.info("QuaRot INT4: rotated+packed %d weights (%d layers x %d proj)%s", + num_layers * len(keys), num_layers, len(keys), + "" if include_down else "; down stays INT8") + return qw + + +def quantize_int8_hadamard(qw: QuantizedWeights, + proj: Dict[str, List[torch.Tensor]], + *, num_layers: int, D: int) -> QuantizedWeights: + """Re-quantize the rotatable projections as **Hadamard-rotated INT8** (W8A8+QuaRot). + + This configuration sits between two tiers: + + * plain per-row INT8 — 8-bit resolution, but *unconditioned*, so a row whose + amax is set by a massive-activation channel loses its remaining ~4090 + channels to rounding; + * QuaRot INT4 — conditioned by the rotation, but only 15 levels. + + Rotating at 8 bits gets both. The rotation is free at inference time: the + weight side folds offline here (``W_rot = H·W/sqrt(K)``) and the activation + side is fused into the norm kernels (``rms_norm_fht_int8_fp16`` and friends). + + Crucially it keeps **plain per-row scales**, so the unmodified + ``cutlass_int8_rowwise_*`` GEMMs are reused — the alternative outlier fixes + (group-128 / block-scaled) would each need a bespoke GEMM, and the measured + ceiling for a hand-written block-scaled s4 kernel on 16-SM Orin was only + 41 TOPS. Principle #17: pick the rotation that keeps you on the fast path. + + The FFN down projection is left as plain INT8: its K (11008) is not a power + of two, and its input is the un-rotated BF16 SiLU output. + """ + Hm = hadamard_gpu(D) / (float(D) ** 0.5) + for li in range(num_layers): + for key in INT4_POW2_PROJECTIONS: + w = proj[key][li].to(torch.float32) # [K, N] + if w.shape[0] != D: + raise RuntimeError( + f"{key} layer {li}: expected K={D}, got {tuple(w.shape)}") + w_rot = (Hm @ w).half() # [K, N] + qw._drop(qw.ptr[key][li]) + qw._set(key, li, *quantize_per_row_int8(w_rot)) + qw.precision[key] = "int8+hadamard" + del w, w_rot + del Hm + torch.cuda.empty_cache() + logger.info("W8A8+Hadamard: rotated %d weights (%d layers x %d proj); " + "down stays plain INT8", + num_layers * len(INT4_POW2_PROJECTIONS), num_layers, + len(INT4_POW2_PROJECTIONS)) + return qw + + +__all__ = [ + "INT8_QUANT_MAX", "INT4_QUANT_MAX", "PROJECTIONS", "INT4_POW2_PROJECTIONS", + "QuantizedWeights", "split_fused_projections", "quantize_per_row_int8", + "hadamard_gpu", "quantize_int8_all", "quantize_int8_hadamard", + "quantize_int4_quarot", +] diff --git a/flash_rt/frontends/torch/_chameleon_rtx_sm87_spec.py b/flash_rt/frontends/torch/_chameleon_rtx_sm87_spec.py new file mode 100644 index 000000000..c5671f802 --- /dev/null +++ b/flash_rt/frontends/torch/_chameleon_rtx_sm87_spec.py @@ -0,0 +1,85 @@ +"""Declarative weight spec for upstream Chameleon-7B (HF layout). + +Chameleon-7B LLM — 32 layers, MHA (num_kv_heads=32, no interleave). +``attention_bias`` and ``mlp_bias`` are both ``false`` in the checkpoint +config, so no bias items are emitted; the per-head QK-Norm weight *and +bias* tensors are part of the layer block. + +Per-head QK Norm prevents norm_fuse; Cat is used for QKV fusion +(not FusedQKV) because the per-head norms are handled as separate +TensorList items. + +The QK-Norm tensors are loaded verbatim (``ToFp16()`` only), preserving their +``(1, 128)`` shape. That is deliberate: ``qk_norm_rope_fused_fp16`` reads the +weight as a flat ``[head_dim]`` vector shared across heads, which is exactly +what the checkpoint's ``model_parallel_size == 1`` layout means (upstream +expands it with ``repeat_interleave`` at forward time). No reshape is needed. +""" + +from __future__ import annotations + +from flash_rt.executors.weight_loader import Item, LayerBlock, ModelWeightSpec +from flash_rt.executors.torch_weights import Attr, Cat, FusedGateUp, T, TensorList, ToFp16 + + +def _llm_block() -> LayerBlock: + """Chameleon-7B LLM — 32 layers, FP16 backbone (no quantized spec).""" + lp = "model.layers.{i}" + items = [ + # ── Fused QKV (MHA: no interleave, no norm_fuse) ── + Item("qkv_w", + Cat([f"{lp}.self_attn.q_proj.weight", + f"{lp}.self_attn.k_proj.weight", + f"{lp}.self_attn.v_proj.weight"], dim=0), + [T()], + TensorList("_llm_qkv_w")), + # ── O projection ── + Item("o_w", f"{lp}.self_attn.o_proj.weight", + [ToFp16(), T()], + TensorList("_llm_o_w")), + # ── Fused GateUp (no norm_fuse — per-head QK Norm incompatible) ── + Item("gu_w", + FusedGateUp(gate=f"{lp}.mlp.gate_proj.weight", + up=f"{lp}.mlp.up_proj.weight"), + [T()], + TensorList("_llm_gu_w")), + # ── Down projection ── + Item("d_w", f"{lp}.mlp.down_proj.weight", + [ToFp16(), T()], + TensorList("_llm_d_w")), + # ── Layer norms ── + Item("input_ln_w", f"{lp}.input_layernorm.weight", + [ToFp16()], TensorList("_llm_input_ln_w")), + Item("post_ln_w", f"{lp}.post_attention_layernorm.weight", + [ToFp16()], TensorList("_llm_post_ln_w")), + # ── Per-head Q/K Norm ── + Item("q_norm_w", f"{lp}.self_attn.q_norm.weight", + [ToFp16()], TensorList("_llm_q_norm_w")), + Item("q_norm_b", f"{lp}.self_attn.q_norm.bias", + [ToFp16()], TensorList("_llm_q_norm_b")), + Item("k_norm_w", f"{lp}.self_attn.k_norm.weight", + [ToFp16()], TensorList("_llm_k_norm_w")), + Item("k_norm_b", f"{lp}.self_attn.k_norm.bias", + [ToFp16()], TensorList("_llm_k_norm_b")), + ] + return LayerBlock(prefix_fmt="", num_layers=32, items=items, name="llm") + + +def build_spec() -> ModelWeightSpec: + """Chameleon-7B: 32 decoder layers + embedding, final norm, lm_head.""" + return ModelWeightSpec( + framework="torch", + blocks=[_llm_block()], + singletons=[ + Item("embed_w", "model.embed_tokens.weight", + [ToFp16()], Attr("_llm_embed_w")), + Item("norm_w", "model.norm.weight", + [ToFp16()], Attr("_llm_norm_w")), + # Live output projection. + Item("lm_head_w", "lm_head.weight", + [ToFp16()], Attr("_llm_lm_head_w")), + ], + ) + + +__all__ = ["build_spec"] diff --git a/flash_rt/frontends/torch/_chameleon_thor_spec.py b/flash_rt/frontends/torch/_chameleon_thor_spec.py new file mode 100644 index 000000000..832369e8b --- /dev/null +++ b/flash_rt/frontends/torch/_chameleon_thor_spec.py @@ -0,0 +1,92 @@ +"""Declarative weight spec for standalone Chameleon-7B on Thor. + +Standard Chameleon 32-layer backbone layout: +attention_bias=false, mlp_bias=false, per-head Q/K norm, SwiGLU FFN. +Per-head QK Norm prevents norm_fuse, so QKV is fused with ``Cat`` and the +gate/up pair with ``FusedGateUp``. +""" + +from __future__ import annotations + +from flash_rt.executors.weight_loader import Item, LayerBlock, ModelWeightSpec +from flash_rt.executors.torch_weights import ( + Attr, + Cat, + FusedGateUp, + Quant, + T, + TensorList, + ToFp16, +) + + +def _llm_block(*, use_fp8: bool = True) -> LayerBlock: + """Chameleon-7B LLM — 32 layers, MHA (num_kv_heads=32, no interleave). + + 4 quantized GEMMs per layer (qkv, o, gu, d) → 32 × 4 = 128 scales + appended to ``target._llm_w_scales``. + """ + qkv_tx = [T(), Quant()] if use_fp8 else [T()] + o_tx = [ToFp16(), T(), Quant()] if use_fp8 else [ToFp16(), T()] + gu_tx = [T(), Quant()] if use_fp8 else [T()] + d_tx = [ToFp16(), T(), Quant()] if use_fp8 else [ToFp16(), T()] + scale_into = "_llm_w_scales" if use_fp8 else None + + lp = "model.layers.{i}" + items = [ + # ── Fused QKV (no bias; per-head QK norm as separate items) ── + Item("qkv_w", + Cat([f"{lp}.self_attn.q_proj.weight", + f"{lp}.self_attn.k_proj.weight", + f"{lp}.self_attn.v_proj.weight"], dim=0), + qkv_tx, + TensorList("_llm_qkv_w"), scale_into=scale_into), + # ── O projection (no bias) ── + Item("o_w", f"{lp}.self_attn.o_proj.weight", + o_tx, + TensorList("_llm_o_w"), scale_into=scale_into), + # ── Fused GateUp ── + Item("gu_w", + FusedGateUp(gate=f"{lp}.mlp.gate_proj.weight", + up=f"{lp}.mlp.up_proj.weight"), + gu_tx, + TensorList("_llm_gu_w"), scale_into=scale_into), + # ── Down projection ── + Item("d_w", f"{lp}.mlp.down_proj.weight", + d_tx, + TensorList("_llm_d_w"), scale_into=scale_into), + # ── Layer norms ── + Item("input_ln_w", f"{lp}.input_layernorm.weight", + [ToFp16()], TensorList("_llm_input_ln_w")), + Item("post_ln_w", f"{lp}.post_attention_layernorm.weight", + [ToFp16()], TensorList("_llm_post_ln_w")), + # ── Per-head Q/K Norm ── + Item("q_norm_w", f"{lp}.self_attn.q_norm.weight", + [ToFp16()], TensorList("_llm_q_norm_w")), + Item("q_norm_b", f"{lp}.self_attn.q_norm.bias", + [ToFp16()], TensorList("_llm_q_norm_b")), + Item("k_norm_w", f"{lp}.self_attn.k_norm.weight", + [ToFp16()], TensorList("_llm_k_norm_w")), + Item("k_norm_b", f"{lp}.self_attn.k_norm.bias", + [ToFp16()], TensorList("_llm_k_norm_b")), + ] + return LayerBlock(prefix_fmt="", num_layers=32, items=items, name="llm") + + +def build_spec(*, use_fp8: bool = True) -> ModelWeightSpec: + """Build the standalone Chameleon-7B Thor weight spec.""" + return ModelWeightSpec( + framework="torch", + blocks=[_llm_block(use_fp8=use_fp8)], + singletons=[ + Item("embed_w", "model.embed_tokens.weight", + [ToFp16()], Attr("_llm_embed_w")), + Item("norm_w", "model.norm.weight", + [ToFp16()], Attr("_llm_norm_w")), + Item("lm_head_w", "lm_head.weight", + [ToFp16()], Attr("_llm_lm_head_w")), + ], + ) + + +__all__ = ["build_spec"] diff --git a/flash_rt/frontends/torch/chameleon_rtx_sm87.py b/flash_rt/frontends/torch/chameleon_rtx_sm87.py new file mode 100644 index 000000000..cf5f582a6 --- /dev/null +++ b/flash_rt/frontends/torch/chameleon_rtx_sm87.py @@ -0,0 +1,721 @@ +"""FlashRT — upstream Chameleon-7B VLM frontend for Jetson AGX Orin (SM87). + +Image+text -> text. Constructed directly (**not** via ``flash_rt.load_model``): +this is a chat-style VLM exposing ``set_prompt()`` + ``generate()``, whereas +``VLAModel.predict`` unconditionally reads ``result['actions']``. Same precedent +as ``qwen3_vl`` — see the redirect in ``flash_rt/api.py``. + +Production precision policy (all defaults) +------------------------------------------ +* **Q/K/V/O, FFN gate/up** — INT8 W8A8 **+ Hadamard rotation** (QuaRot at 8 + bits): weights rotated offline, activations rotated inside the fused norm + kernel; per-row dynamic activation scales. +* **FFN down** — INT8 W8A8 per-row dynamic. Not rotated: K=11008 is not a power + of two and its input is the un-rotated BF16 SiLU output. +* **lm_head** (65536x4096) — INT8 W8A8. 268 MB/token = 3.7 % of the decode + budget; FP16 would double that for no measured argmax benefit. +* **residual / QK-LayerNorm / RoPE / attention / KV cache** — FP16, with + ``ffn_down_clamp`` applied on the last ``ffn_down_clamp_last_n`` layers. +* **attention** — FA2 fp16 causal; ``split_kv_bias=4`` on decode, because FA2's + own heuristic returns ``num_splits=1`` at Chameleon's 32 Q heads. +* **VQ-GAN encoder** — FP16 convs + **fp32** codebook distance/argmin. + +Measured (ISL=1032, OSL=16, warm): greedy output **bit-identical to the HF bf16 +reference for 16/16 tokens**, worst per-layer cosine 0.9986, last-row logit +cosine 0.999968, **21.07 tok/s** decode, 273.8 ms prefill, 7.6 GB resident. + +Two non-obvious *correctness* requirements +------------------------------------------ +1. **The Hadamard rotation is not an optimization.** Chameleon's + massive-activation channels (measured: L31 row-0 channel d632 at 2.4e4 against + a row median of 1e4) pin the per-row INT8 amax and round that row's other + ~4090 channels to zero. Plain per-row INT8 reproduces only 8/16 reference + tokens; rotating fixes it. It is free because it preserves per-row scales, so + the stock ``cutlass_int8_rowwise_*`` GEMMs are reused unchanged. +2. **``ffn_down_clamp`` is not a tuning knob.** L31's down output reaches 2.6e5, + past FP16's 65504; without the clamp the residual stores ``inf`` and the final + RMSNorm poisons that row. See ``models/chameleon/pipeline_thor.py``. + +Dim policy: backbone dims (32 layers / 4096 / 32 heads / 11008 / vocab 65536) and +the special token ids are **hard-asserted** against ``config.json``, so this +frontend is Chameleon-7B-specific by construction rather than by convention. +Note ``config.json``'s ``bos_token_id`` is stale (says 1, which is ````); +``tokenizer.json`` gives `` = 0`` and that is what the processor emits. + +Preprocessing specifics +----------------------- +Chameleon needs PIL LANCZOS / 512 / ``u8*0.0078-1.0`` -> ``[-1, +0.989]`` +normalization and a bare 1024-token raster per image (no grid/newline +layout); the quantizers live in ``_chameleon_quant``. + +See ``docs/chameleon7b_rtx_sm87.md`` for the roofline, the measured lever menu, +and the dead-ends. +""" + +from __future__ import annotations + +import json +import logging +import os +import time +from pathlib import Path +from typing import List, Optional, Sequence + +import torch + +import flash_rt.flash_rt_kernels as fvk +from flash_rt.frontends.torch import _chameleon_quant as cq +from flash_rt.frontends.torch._chameleon_rtx_sm87_spec import build_spec +from flash_rt.hardware.rtx.attn_backend_chameleon import ( + ChameleonAttnBackend, make_chameleon_attention_spec) +from flash_rt.models.chameleon.pipeline_rtx import chameleon_forward + +logger = logging.getLogger(__name__) + +_FP16 = torch.float16 +_BF16 = torch.bfloat16 + + +class ChameleonTorchFrontendRtxSm87: + """Chameleon-7B image+text → text on Orin SM87. + + Typical use:: + + f = ChameleonTorchFrontendRtxSm87("/path/to/Chameleon_7B_mGPT") + f.set_prompt("Describe this image.", images=[pil_img]) + print(f.generate(max_new_tokens=32)) + """ + + # Token ids for this checkpoint. Hardcoded so an accelerated deployment does + # not depend on the training code being importable, but asserted against + # config.json at load so a different checkpoint fails loudly (principle #10). + BOS_ID = 0 + EOS_ID = 2 + IMG_PLACEHOLDER_ID = 8711 # + BOI_ID = 8197 # + EOI_ID = 8196 # + SEP_ID = 8710 # + IMG_ID_OFFSET = 4 # token_id = codebook_index + 4 + N_IMG_CODES = 8192 + IMG_TOKENS_PER_VIEW = 1024 # 512/16 = 32 -> 32x32 raster + + # Chameleon-7B backbone dims — hard-asserted against config.json. + D = 4096 + L = 32 + H = 32 + HKV = 32 + HD = 128 + DFF = 11008 + VOCAB = 65536 + + def __init__(self, checkpoint_dir: str, *, + max_seq: int = 2048, + use_int4: bool = False, + use_int4_down: bool = False, + use_hadamard: bool = True, + split_kv_bias: int = 4, + ffn_down_clamp: Optional[float] = None, + vq_argmin_fp32: bool = True, + free_fp16_weights: bool = True, + probe_layers: Optional[Sequence[int]] = None, + **_ignored) -> None: + cc = torch.cuda.get_device_capability(0) + if cc != (8, 7) and os.environ.get("FLASHRT_CHAMELEON_SM87_FORCE") != "1": + raise RuntimeError( + f"ChameleonTorchFrontendRtxSm87 targets SM87 (Orin); found SM{cc[0]}{cc[1]}. " + "The INT8/INT4 path needs ENABLE_SM80_INT8_CUTLASS=ON. Set " + "FLASHRT_CHAMELEON_SM87_FORCE=1 to override.") + + self.checkpoint_dir = Path(checkpoint_dir) + self.max_seq = int(max_seq) + self.use_int4_down = bool(use_int4_down) + self.use_int4 = bool(use_int4) or self.use_int4_down + # W8A8+QuaRot: Hadamard-rotate the six K=4096 projections and quantize + # at 8 bits. Conditions the massive-activation channels (which plain + # per-row INT8 cannot handle) without INT4's noise, and keeps per-row + # scales so the stock CUTLASS GEMM is reused. Ignored on the INT4 tier, + # which already rotates. + self.use_hadamard = bool(use_hadamard) and not self.use_int4 + self.vq_argmin_fp32 = bool(vq_argmin_fp32) + # Required for correctness: L31's down output reaches ~2.6e5, far past + # FP16's 65504. See chameleon_forward's docstring for the measured + # per-layer table and docs/chameleon_acceleration_methodology.md §2.1. + self.ffn_down_clamp = float( + os.environ.get("FLASHRT_CHAMELEON_DOWN_CLAMP", "60000.0") + if ffn_down_clamp is None else ffn_down_clamp) + self.ffn_down_clamp_last_n = int( + os.environ.get("FLASHRT_CHAMELEON_DOWN_CLAMP_LAST_N", "4")) + self._probe_layers = list(probe_layers) if probe_layers else [] + + self._validate_config() + t0 = time.perf_counter() + self._load_weights(free_fp16_weights=free_fp16_weights) + self._build_rope_tables() + self._allocate_buffers() + self._build_attn_backend(split_kv_bias=split_kv_bias) + self._load_processor() + self._load_vqgan() + + # Prompt state. + self.input_ids: Optional[torch.Tensor] = None + self.S = 0 + self._prompt_ready = False + self._timing: dict = {} + + logger.info( + "ChameleonTorchFrontendRtxSm87 ready in %.1fs — tier=%s max_seq=%d " + "gpu_mem=%.2f GB", + time.perf_counter() - t0, self.precision_tier, self.max_seq, + torch.cuda.memory_allocated() / 2 ** 30) + + # ================================================================== + # Load + # ================================================================== + + @property + def precision_tier(self) -> str: + if self.use_int4_down: + return "int4+down" + if self.use_int4: + return "int4" + return "int8+hadamard" if self.use_hadamard else "int8" + + def _validate_config(self) -> None: + cfg = json.loads((self.checkpoint_dir / "config.json").read_text()) + expect = { + "hidden_size": self.D, "num_hidden_layers": self.L, + "num_attention_heads": self.H, "num_key_value_heads": self.HKV, + "intermediate_size": self.DFF, "vocab_size": self.VOCAB, + } + for k, want in expect.items(): + got = cfg.get(k) + if got != want: + raise RuntimeError( + f"config.json {k}={got}, this frontend hardcodes {want}. " + "Chameleon-7B only.") + if cfg.get("attention_bias") or cfg.get("mlp_bias"): + raise RuntimeError( + "This frontend assumes attention_bias=false and mlp_bias=false " + "(upstream Chameleon). A biased checkpoint needs the *_bias " + "GEMM entries wired back in.") + if cfg.get("model_parallel_size", 1) != 1: + raise RuntimeError( + f"model_parallel_size={cfg.get('model_parallel_size')}: the " + "QK-Norm params would differ per head group, but " + "qk_norm_rope_fused_fp16 broadcasts one [head_dim] vector " + "across all heads. Only the 7B (mp=1) layout is supported.") + if self.max_seq > cfg.get("max_position_embeddings", 4096): + raise ValueError( + f"max_seq={self.max_seq} exceeds max_position_embeddings=" + f"{cfg['max_position_embeddings']}") + + # Token-id contract. config.json's bos_token_id is stale (says 1, which + # is ); tokenizer.json is authoritative and gives = 0. + vm = cfg.get("vocabulary_map") or {} + for tok, want in (("", self.IMG_PLACEHOLDER_ID), + ("", self.BOI_ID), + ("", self.EOI_ID), + ("", self.SEP_ID)): + if vm.get(tok) != want: + raise RuntimeError( + f"token id mismatch: {tok} is {vm.get(tok)} in config.json, " + f"expected {want}") + img_ids = sorted(v for k, v in vm.items() if k.startswith("IMGIMG")) + if (len(img_ids) != self.N_IMG_CODES + or img_ids[0] != self.IMG_ID_OFFSET + or img_ids[-1] != self.IMG_ID_OFFSET + self.N_IMG_CODES - 1): + raise RuntimeError( + f"expected {self.N_IMG_CODES} contiguous IMGIMG ids starting at " + f"{self.IMG_ID_OFFSET}; got {len(img_ids)} spanning " + f"[{img_ids[0]}, {img_ids[-1]}]") + self._config = cfg + + def _load_weights(self, *, free_fp16_weights: bool) -> None: + from flash_rt.executors.torch_weights import ( + MultiSafetensorsSource, WeightLoader) + + shards = sorted(self.checkpoint_dir.glob("model-*-of-*.safetensors")) + if not shards: + shards = sorted(self.checkpoint_dir.glob("*.safetensors")) + if not shards: + raise FileNotFoundError(f"no safetensors in {self.checkpoint_dir}") + + src = MultiSafetensorsSource([str(p) for p in shards], device="cuda") + WeightLoader(source=src, target=self, spec=build_spec()).run() + del src + + proj = cq.split_fused_projections( + self._llm_qkv_w, self._llm_gu_w, self._llm_o_w, self._llm_d_w, + D=self.D, Dff=self.DFF) + # The fused tensors are dead once split; release before quantizing so + # the peak is (fp16 split + int8) rather than (fused + split + int8). + self._llm_qkv_w = [] + self._llm_gu_w = [] + torch.cuda.empty_cache() + + self.qw = cq.quantize_int8_all(proj, num_layers=self.L) + if self.use_int4: + cq.quantize_int4_quarot(self.qw, proj, num_layers=self.L, D=self.D, + include_down=self.use_int4_down) + elif self.use_hadamard: + cq.quantize_int8_hadamard(self.qw, proj, num_layers=self.L, D=self.D) + + # lm_head: INT8 in both tiers (see pipeline docstring for the ROI). + self._lm_head_q, self._lm_head_s = cq.quantize_per_row_int8( + self._llm_lm_head_w.t().contiguous()) # [V, D] -> [K=D, N=V] + + if free_fp16_weights: + for key in cq.PROJECTIONS: + proj[key] = [] + self._llm_o_w = [] + self._llm_d_w = [] + self._llm_lm_head_w = None + torch.cuda.empty_cache() + + logger.info("weights: %.2f GB quantized (%s) + %.2f GB lm_head int8 + " + "%.2f GB embed fp16", + self.qw.bytes() / 2 ** 30, self.precision_tier, + self._lm_head_q.numel() / 2 ** 30, + self._llm_embed_w.numel() * 2 / 2 ** 30) + + def _build_rope_tables(self) -> None: + """cos/sin as ``[max_seq, HD]`` fp16, ``cat([f, f], -1)`` tiled. + + The kernel reads ``cos[s*HD + d]`` for d in ``[0, HD)``, so the + half-frequencies must be duplicated across the two halves; a + ``[max_seq, HD/2]`` table indexes out of stride and corrupts RoPE. Being + row-major with stride exactly HD is also what lets a decode step select + position ``pos`` by pointer arithmetic alone. + """ + theta = float(self._config.get("rope_theta", 10000.0)) + inv = 1.0 / (theta ** (torch.arange(0, self.HD, 2, dtype=torch.float32, + device="cuda") / self.HD)) + pos = torch.arange(self.max_seq, device="cuda", dtype=torch.float32) + f = pos[:, None] * inv[None, :] + full = torch.cat([f, f], dim=-1) + self._rope_cos = torch.cos(full).to(_FP16).contiguous() + self._rope_sin = torch.sin(full).to(_FP16).contiguous() + + def _allocate_buffers(self) -> None: + MS, D, Dff, V = self.max_seq, self.D, self.DFF, self.VOCAB + dev = "cuda" + z = lambda *sh, dt: torch.zeros(*sh, dtype=dt, device=dev) # noqa: E731 + + self._x = z(MS, D, dt=_FP16) # residual stream + self._xn = z(MS, D, dt=_FP16) # final-norm output + self._o_proj_out = z(MS, D, dt=_FP16) + self._int8_act_d = z(MS, D, dt=torch.int8) + self._int8_act_ff = z(MS, Dff, dt=torch.int8) + self._bf16_gate_ff = z(MS, Dff, dt=_BF16) + self._bf16_xn_ff = z(MS, Dff, dt=_BF16) + self._int4_act_d = z(MS, D // 2, dt=torch.uint8) if self.use_int4 else None + self._int4_act_ff = (z(MS, Dff // 2, dt=torch.uint8) + if self.use_int4_down else None) + + # Dynamic per-row activation scales — one shared [MS] vector per quant + # site, reused by every layer (decode never uses static calibration: + # that was fitted at prefill M=Se and does not describe one decode row). + self._act_scale = {k: z(MS, dt=torch.float32) + for k in ("qkv", "o", "gu", "down")} + self._lm_act = z(MS, D, dt=torch.int8) + self._lm_act_scale = z(MS, dt=torch.float32) + self._logits = z(1, V, dt=_BF16) + self._logits_all: Optional[torch.Tensor] = None + self._bf16_min = torch.finfo(_BF16).min + + self._probe_bufs = [z(MS, D, dt=_FP16) for _ in self._probe_layers] + self._probe_final = z(MS, D, dt=_FP16) if self._probe_layers else None + + self._tok_dev = z(1, dt=torch.long) + + def _build_attn_backend(self, *, split_kv_bias: int) -> None: + spec = make_chameleon_attention_spec( + num_layers=self.L, num_q_heads=self.H, num_kv_heads=self.HKV, + head_dim=self.HD, max_seq=self.max_seq) + self.attn = ChameleonAttnBackend(spec, max_seq=self.max_seq, + split_kv_bias=split_kv_bias) + + def _load_processor(self) -> None: + """Use the HF ChameleonProcessor verbatim. + + Reimplementing it is a correctness liability: the pipeline is + blend-RGBA-on-white → PIL **LANCZOS** shortest-edge 512 → center-crop + 512 → ``float32(float64(u8) * 0.0078) - 1.0``, giving ``[-1, +0.989]`` + (note: *not* ``[-1, 1]``). ``ChameleonImageProcessorFast`` silently + substitutes BICUBIC for LANCZOS, so the slow/PIL path is required — it + is what ``preprocessor_config.json`` selects by default, and it runs + once per image outside any graph. + """ + from transformers import AutoProcessor + self.processor = AutoProcessor.from_pretrained(str(self.checkpoint_dir)) + if int(getattr(self.processor, "image_seq_length", 0)) != self.IMG_TOKENS_PER_VIEW: + raise RuntimeError( + f"processor image_seq_length=" + f"{self.processor.image_seq_length}, expected " + f"{self.IMG_TOKENS_PER_VIEW}") + + def _load_vqgan(self) -> None: + """Load the HF ``ChameleonVQVAE`` encoder from the checkpoint shards.""" + from safetensors.torch import load_file + from transformers import ChameleonConfig, ChameleonVQVAE + + index = self.checkpoint_dir / "model.safetensors.index.json" + prefix = "model.vqmodel." + sd = {} + if index.exists(): + wmap = json.loads(index.read_text())["weight_map"] + keys = [k for k in wmap if k.startswith(prefix)] + for shard in sorted({wmap[k] for k in keys}): + full = load_file(str(self.checkpoint_dir / shard)) + for k in keys: + if k in full: + sd[k[len(prefix):]] = full[k] + del full + if not sd: + raise FileNotFoundError( + f"no {prefix}* weights found under {self.checkpoint_dir}") + + cfg = ChameleonConfig.from_pretrained(str(self.checkpoint_dir)) + vq = ChameleonVQVAE._from_config(cfg.vq_config) + missing, unexpected = vq.load_state_dict(sd, strict=False) + if unexpected: + raise RuntimeError(f"unexpected vqmodel keys: {unexpected[:4]}") + if missing: + raise RuntimeError(f"missing vqmodel keys: {missing[:4]}") + self.vqgan = vq.eval().to(device="cuda", dtype=_FP16) + # fp32 codebook for the distance/argmin (see _vq_encode). + self._vq_codebook_f32 = ( + self.vqgan.quantize.embedding.weight.detach().float().contiguous()) + logger.info("VQ-GAN encoder loaded (%d tensors, fp16 convs, " + "argmin in %s)", len(sd), + "fp32" if self.vq_argmin_fp32 else "fp16") + + # ================================================================== + # Image tokenization + # ================================================================== + + @torch.no_grad() + def _vq_encode(self, pixel_values: torch.Tensor) -> torch.Tensor: + """``[N, 3, 512, 512]`` fp32 in ``[-1, 0.989]`` → ``[N, 1024]`` token ids. + + The convs run fp16 but the codebook distance and argmin run **fp32**: + ``|z|^2 + |e|^2 - 2 z.e`` is cancellation-prone, and in fp16 it flips + codebook indices (measured 98.14 % → 99.02 % index match vs an all-fp32 + reference for <0.1 ms; see docs §2.4). Tokens are a row-major raster of + the 32x32 latent grid, and the id is simply ``code + 4``. + """ + px = pixel_values.to(device="cuda", dtype=_FP16) + h = self.vqgan.quant_conv(self.vqgan.encoder(px)) + if not self.vq_argmin_fp32: + _, _, idx = self.vqgan.quantize(h) + return idx.view(px.shape[0], -1) + self.IMG_ID_OFFSET + e = self._vq_codebook_f32 # [n_emb, dim] + z = h.permute(0, 2, 3, 1).contiguous().view(-1, e.shape[1]).float() + d = z.pow(2).sum(1, keepdim=True) + e.pow(2).sum(1) - 2.0 * (z @ e.t()) + return d.argmin(1).view(px.shape[0], -1) + self.IMG_ID_OFFSET + + # ================================================================== + # Prompt + # ================================================================== + + def set_prompt(self, text: str, images=None, *, + input_ids: Optional[Sequence[int]] = None) -> None: + """Tokenize, VQ-encode the images, and seed the residual stream. + + Args: + text: prompt containing one ```` per image, e.g. + ``"Describe this image."``. The processor expands each + placeholder to ```` + 1024x ```` + + ````, prepends BOS and appends the sep token. + images: list of PIL images (or a single image). + input_ids: bypass tokenization + VQ entirely and use these ids + verbatim. Used by the precision harness to feed the reference's + exact ids, which isolates LLM error from VQ-GAN index drift. + """ + t0 = time.perf_counter() + if input_ids is not None: + ids = torch.as_tensor(list(input_ids), dtype=torch.long, device="cuda") + n_img = 0 + else: + if images is not None and not isinstance(images, (list, tuple)): + images = [images] + n_text_img = text.count("") + if images and n_text_img != len(images): + raise ValueError( + f"text has {n_text_img} '' placeholders but " + f"{len(images)} images were given") + enc = self.processor(text=text, images=images if images else None, + return_tensors="pt") + ids = enc["input_ids"][0].to("cuda") + n_img = 0 + if images: + img_ids = self._vq_encode(enc["pixel_values"]) # [N, 1024] + n_img = img_ids.shape[0] + # Upstream substitutes at the id level: masked_scatter over + # input_ids == . Order is row-major per image, images in + # order, which matches the placeholder order in the string. + slot = ids == self.IMG_PLACEHOLDER_ID + want = n_img * self.IMG_TOKENS_PER_VIEW + if int(slot.sum()) != want: + raise RuntimeError( + f"{int(slot.sum())} placeholders vs {want} " + "VQ tokens") + ids = ids.masked_scatter(slot, img_ids.reshape(-1).to(ids.dtype)) + + S = int(ids.numel()) + if S > self.max_seq: + raise ValueError( + f"prompt is {S} tokens but max_seq={self.max_seq}. Note " + f"S = 1 + n_img*1026 + n_text + 1.") + if S + 1 > self.max_seq: + raise ValueError( + f"prompt {S} tokens leaves no room to decode within " + f"max_seq={self.max_seq}") + + # NOTE: S is used exactly, never padded. Padding would put junk rows in + # the KV cache that decode would then attend to (a prefill path without a KV cache would pad Se to even). + self.input_ids = ids + self.S = S + self.attn.reset_cache() + # Seed the residual stream. No sqrt(D) scaling — Chameleon feeds raw + # embeddings, and image tokens are ordinary vocab rows (no projector). + torch.index_select(self._llm_embed_w, 0, ids, out=self._x[:S]) + self._prompt_ready = True + self._timing = {"prompt_ms": (time.perf_counter() - t0) * 1e3, + "S": S, "n_images": n_img} + + # ================================================================== + # Forward + # ================================================================== + + def _dims(self) -> dict: + return {"D": self.D, "Dff": self.DFF, "L": self.L, "H": self.H, + "Hd": self.HD, "vocab": self.VOCAB} + + def _bufs(self, *, logits_ptr: int) -> dict: + return { + "x": self._x.data_ptr(), + "xn": self._xn.data_ptr(), + "o_proj_out": self._o_proj_out.data_ptr(), + "int8_act_d": self._int8_act_d.data_ptr(), + "int8_act_ff": self._int8_act_ff.data_ptr(), + "int4_act_d": self._int4_act_d.data_ptr() if self.use_int4 else 0, + "int4_act_ff": (self._int4_act_ff.data_ptr() + if self.use_int4_down else 0), + "bf16_gate_ff": self._bf16_gate_ff.data_ptr(), + "bf16_xn_ff": self._bf16_xn_ff.data_ptr(), + "logits": logits_ptr, + "lm_act": self._lm_act.data_ptr(), + "lm_act_scale": self._lm_act_scale.data_ptr(), + } + + def _weights(self) -> dict: + w = { + "rope_cos": self._rope_cos.data_ptr(), + "rope_sin": self._rope_sin.data_ptr(), + "final_norm_w": self._llm_norm_w.data_ptr(), + "lm_head_w": self._lm_head_q.data_ptr(), + "lm_head_w_scale": self._lm_head_s.data_ptr(), + "input_ln_w": [t.data_ptr() for t in self._llm_input_ln_w], + "post_ln_w": [t.data_ptr() for t in self._llm_post_ln_w], + "q_norm_w": [t.data_ptr() for t in self._llm_q_norm_w], + "q_norm_b": [t.data_ptr() for t in self._llm_q_norm_b], + "k_norm_w": [t.data_ptr() for t in self._llm_k_norm_w], + "k_norm_b": [t.data_ptr() for t in self._llm_k_norm_b], + } + for key in cq.PROJECTIONS: + w[f"{key}_w"] = self.qw.ptr[key] + w[f"{key}_w_scale"] = self.qw.scale_ptr[key] + return w + + def _scales_dev(self) -> dict: + return {f"act_{k}": [v.data_ptr()] * self.L + for k, v in self._act_scale.items()} + + def _probe(self) -> Optional[dict]: + if not self._probe_layers: + return None + return {"layers": self._probe_layers, + "bufs": [b.data_ptr() for b in self._probe_bufs], + "final_buf": self._probe_final.data_ptr()} + + def _require_prompt(self) -> None: + if not self._prompt_ready: + raise RuntimeError("call set_prompt() before prefill()/generate()") + + @torch.no_grad() + def prefill(self, *, logits_all: bool = False) -> torch.Tensor: + """Run the prompt through the 32 layers, filling the KV cache. + + Returns the masked BF16 logits: ``[1, vocab]`` for the last position, or + ``[S, vocab]`` when ``logits_all`` (teacher-forced comparison). + """ + self._require_prompt() + t0 = time.perf_counter() + if logits_all: + if self._logits_all is None or self._logits_all.shape[0] < self.S: + self._logits_all = torch.zeros(self.max_seq, self.VOCAB, + dtype=_BF16, device="cuda") + out = self._logits_all[:self.S] + logits_ptr = self._logits_all.data_ptr() + else: + out = self._logits + logits_ptr = self._logits.data_ptr() + + chameleon_forward( + fvk, self._bufs(logits_ptr=logits_ptr), self._weights(), + self._dims(), self._scales_dev(), + attn=self.attn, S=self.S, pos=None, stream=0, + use_int4=self.use_int4, use_int4_down=self.use_int4_down, + use_hadamard=self.use_hadamard, + ffn_down_clamp_value=self.ffn_down_clamp, + ffn_down_clamp_last_n=self.ffn_down_clamp_last_n, + logits_all=logits_all, probe=self._probe()) + self._mask_image_logits(out) + torch.cuda.synchronize() + self._timing["prefill_ms"] = (time.perf_counter() - t0) * 1e3 + return out + + @torch.no_grad() + def decode_step(self, token_id, *, pos: int, stream: int = 0) -> torch.Tensor: + """One decode step: embed ``token_id``, attend keys ``[0, pos]``. + + ``pos`` is the absolute KV position this token occupies, i.e. ``S`` for + the first generated token. + + ``stream`` must be the capture stream when this body is being recorded + into a CUDA graph. Launching the kernels on the legacy default stream + while another stream is capturing leaves them **out** of the graph + without raising — the replay then only re-runs the torch ops and the + logits never change, which looks exactly like a frozen/stale graph. + """ + self._require_prompt() + if isinstance(token_id, torch.Tensor): + self._tok_dev.copy_(token_id.reshape(1).to(torch.long)) + else: + self._tok_dev.fill_(int(token_id)) + torch.index_select(self._llm_embed_w, 0, self._tok_dev, out=self._x[:1]) + chameleon_forward( + fvk, self._bufs(logits_ptr=self._logits.data_ptr()), self._weights(), + self._dims(), self._scales_dev(), + attn=self.attn, S=1, pos=int(pos), stream=int(stream), + use_int4=self.use_int4, use_int4_down=self.use_int4_down, + use_hadamard=self.use_hadamard, + ffn_down_clamp_value=self.ffn_down_clamp, + ffn_down_clamp_last_n=self.ffn_down_clamp_last_n) + self._mask_image_logits(self._logits) + return self._logits + + def _mask_image_logits(self, logits: torch.Tensor) -> None: + """Suppress the 8192 image-codebook ids, as upstream does every forward. + + Upstream applies this inside ``forward`` at *all* positions, so it is + not something a ``LogitsProcessor`` could be used for and it cannot be + disabled through the generation config. + """ + lo = self.IMG_ID_OFFSET + logits[:, lo:lo + self.N_IMG_CODES].fill_(self._bf16_min) + + # ================================================================== + # Generate + # ================================================================== + + @torch.no_grad() + def generate(self, text: Optional[str] = None, images=None, *, + max_new_tokens: int = 32, eos_token_id: Optional[int] = None, + return_ids: bool = False, skip_special_tokens: bool = True): + """Greedy decode. Returns the decoded string (or the raw id list). + + Greedy only: the checkpoint's generation config is already + ``do_sample=False``, and argmax over BF16 logits is order-preserving, so + it matches an fp32 argmax except on exact BF16 ties. + """ + if text is not None: + self.set_prompt(text, images) + self._require_prompt() + if max_new_tokens < 0: + raise ValueError( + f"max_new_tokens must be >= 0, got {max_new_tokens}") + if max_new_tokens == 0: + return [] if return_ids else "" + eos = self.EOS_ID if eos_token_id is None else int(eos_token_id) + budget = min(max_new_tokens, self.max_seq - self.S) + if budget < max_new_tokens: + logger.warning("max_new_tokens clipped %d -> %d by max_seq=%d", + max_new_tokens, budget, self.max_seq) + + logits = self.prefill() + tok = int(torch.argmax(logits[0]).item()) + out: List[int] = [tok] + + t0 = time.perf_counter() + steps = 0 + for i in range(budget - 1): + if tok == eos: + break + logits = self.decode_step(tok, pos=self.S + i) + tok = int(torch.argmax(logits[0]).item()) + out.append(tok) + steps += 1 + torch.cuda.synchronize() + dt = time.perf_counter() - t0 + self._timing["decode_steps"] = steps + self._timing["decode_ms_per_token"] = (dt / steps * 1e3) if steps else 0.0 + self._timing["decode_tok_s"] = (steps / dt) if dt > 0 else 0.0 + + if out and out[-1] == eos: + out = out[:-1] + if return_ids: + return out + return self.processor.tokenizer.decode( + out, skip_special_tokens=skip_special_tokens) + + # ================================================================== + # Introspection + # ================================================================== + + def snapshot_probe(self) -> dict: + """Per-layer post-residual hidden states captured during the last forward.""" + if not self._probe_layers: + return {} + S = self.S + out = {f"layer_{li}": b[:S].clone() + for li, b in zip(self._probe_layers, self._probe_bufs)} + out["final_norm"] = self._probe_final[:S].clone() + return out + + def reset(self) -> None: + self.attn.reset_cache() + self._prompt_ready = False + self.input_ids = None + self.S = 0 + + @property + def timing(self) -> dict: + return dict(self._timing) + + def precision_spec(self) -> dict: + return { + "tier": self.precision_tier, + "llm_gemms": dict(self.qw.precision), + "lm_head": "int8", + "residual": "fp16", + "attention": "fp16 FA2 causal (split_kv_bias=" + f"{self.attn.split_kv_bias})", + "kv_cache": "fp16", + "ffn_down_clamp": f"{self.ffn_down_clamp:.0f} on last {self.ffn_down_clamp_last_n} layers", + "vqgan": f"fp16 convs, argmin " + f"{'fp32' if self.vq_argmin_fp32 else 'fp16'}", + } + + def get_model_info(self) -> dict: + return { + "model": "chameleon-7b", "arch": "rtx_sm87", + "layers": self.L, "hidden": self.D, "ffn": self.DFF, + "heads": f"{self.H}Q/{self.HKV}KV", "head_dim": self.HD, + "vocab": self.VOCAB, "max_seq": self.max_seq, + "precision": self.precision_spec(), + } + + +__all__ = ["ChameleonTorchFrontendRtxSm87"] diff --git a/flash_rt/frontends/torch/chameleon_thor.py b/flash_rt/frontends/torch/chameleon_thor.py new file mode 100644 index 000000000..994dc4bc0 --- /dev/null +++ b/flash_rt/frontends/torch/chameleon_thor.py @@ -0,0 +1,911 @@ +"""Standalone Chameleon-7B frontend for Jetson Thor (SM110). + +Direct-use VLM/LLM frontend: text + real images -> Chameleon prefill logits +and incremental KV-cache greedy generation. It uses the Chameleon +Thor dynamic-FP8 pipeline and causal FMHA attention backend. +""" + +from __future__ import annotations + +import ctypes +import json +import logging +import math +import os +import pathlib +from typing import List, Optional + +import numpy as np +import PIL +from PIL import Image as _PILImage +import torch + +import flash_rt.flash_rt_kernels as fvk +try: + import flash_rt.flash_rt_fp4 as fvk_fp4 +except ModuleNotFoundError as exc: + if exc.name != "flash_rt.flash_rt_fp4": + raise + fvk_fp4 = None +from flash_rt.hardware.thor.attn_backend_chameleon import ( + ThorChameleonAttnBackend, + make_chameleon_attention_spec, +) +from flash_rt.models.chameleon.pipeline_thor import ( + chameleon_decode_step, + chameleon_forward, + chameleon_forward_calibrate, + chameleon_forward_fp16, +) + +logger = logging.getLogger(__name__) + +fp16 = torch.float16 +fp8 = torch.float8_e4m3fn +_cudart = ctypes.CDLL("libcudart.so") + +D_LLM = 4096 +NH_LLM = 32 +HD_LLM = 128 +L_LLM = 32 +DFF_LLM = 11008 +VOCAB_SIZE = 65536 +ROPE_THETA = 10000.0 + +# Chameleon image/text special ids from the shipped vocabulary_map. +PAD_ID = 1 +EOS_ID = 2 +IMG_START_ID = 8197 # +IMG_END_ID = 8196 # +NEWLINE_ID = 8803 +GRID_TOK_BASE = 8804 +PATCH_SIZE = 32 + + +class ChameleonTorchFrontendThor: + """Standalone Chameleon-7B Thor prefill frontend.""" + + #: Required CUDA capability (Jetson Thor SM110) and the documented + #: dev override that skips the probe. Mirrors the Orin frontend's gate. + _REQUIRED_CAPABILITY = (11, 0) + _FORCE_ARCH_ENV = "FLASHRT_CHAMELEON_THOR_FORCE" + + def _require_arch(self) -> None: + if os.environ.get(self._FORCE_ARCH_ENV) == "1": + return # explicit documented dev override: skip the probe + if not torch.cuda.is_available(): + raise RuntimeError( + "ChameleonTorchFrontendThor requires a Jetson Thor SM110 CUDA " + "device; CUDA is not available.") + cc = torch.cuda.get_device_capability(0) + if cc != self._REQUIRED_CAPABILITY: + raise RuntimeError( + f"ChameleonTorchFrontendThor targets SM110 (Thor); found " + f"SM{cc[0]}{cc[1]}. The FP8 path needs the Thor kernel set. " + f"Set {self._FORCE_ARCH_ENV}=1 to override for development.") + + def __init__( + self, + checkpoint_dir: str, + *, + use_fp8: bool = True, + use_cuda_graph: bool = True, + max_seq: int = 4096, + target_size: int = 512, + tokenizer_path: Optional[str] = None, + vqgan_path: Optional[str] = None, + use_trt_vqgan: bool = False, + trt_vqgan_engine_dir: Optional[str] = None, + use_autotune: bool = True, + ffn_clamp_layers: Optional[List[int]] = None, + fp4_ffn_layers: Optional[List[int]] = None, + use_fa4_attn: Optional[bool] = None, + ) -> None: + self._require_arch() + self.checkpoint_dir = pathlib.Path(checkpoint_dir).expanduser().resolve() + if not self.checkpoint_dir.exists(): + raise FileNotFoundError(f"checkpoint_dir not found: {self.checkpoint_dir}") + self._use_fp8 = bool(use_fp8) + self._use_cuda_graph = bool(use_cuda_graph) + self._max_pos = int(max_seq) + if self._max_pos < 16: + raise ValueError(f"max_seq must be at least 16, got {max_seq}") + self.target_size = int(target_size) + self.tokenizer_path = tokenizer_path + self.vqgan_path = vqgan_path + self._use_trt_vqgan = bool(use_trt_vqgan) + self._trt_vqgan_engine_dir = trt_vqgan_engine_dir + self._trt_vqgan_backend = None + self._vqgan_backend = "eager" + self._trt_stream = torch.cuda.Stream() if self._use_trt_vqgan else None + self._infer_graph = None + self._captured_Se = None + self._last_input_ids: Optional[list[int]] = None + self.Se: Optional[int] = None + self._real_len: int = 0 + self._use_autotune = bool(use_autotune) + self._autotuned_se: set = set() + self._ffn_clamp_layers = self._resolve_ffn_clamp_layers(ffn_clamp_layers) + self._fp4_ffn_layers = self._resolve_fp4_ffn_layers(fp4_ffn_layers) + if use_fa4_attn is None: + use_fa4_attn = os.environ.get("FLASHRT_CHAMELEON_FA4_ATTN", "0") in ("1", "true", "on") + self._use_fa4_attn = bool(use_fa4_attn) + self._stream = torch.cuda.current_stream().cuda_stream + + with open(self.checkpoint_dir / "config.json") as f: + self.config_json = json.load(f) + self._validate_config() + self._build_image_token_mask() + with open(self.checkpoint_dir / "model.safetensors.index.json") as f: + self.weight_index = json.load(f) + + self._load_weights() + self._build_rope_tables() + self._load_tokenizer() + self._load_vqgan() + self._allocate_buffers() + + from flash_rt.core.context import FvkContext + self._ctx = FvkContext() + self._gemm = self._ctx.gemm + self._build_attention_backend() + if self._use_fa4_attn: + self._attn.set_fa4_attn(self._bufs["xn"], self._kv_cache) + self._logits_buf.zero_() + + logger.info( + "ChameleonTorchFrontendThor init: max_seq=%d target_size=%d fp8=%s graph=%s", + self._Se_max, self.target_size, self._use_fp8, self._use_cuda_graph, + ) + + def _resolve_ffn_clamp_layers(self, layers): + spec = os.environ.get("FLASHRT_CHAMELEON_FFN_CLAMP_LAYERS") \ + if layers is None else layers + if spec is None: + return frozenset({31}) + if isinstance(spec, str): + text = spec.strip().lower() + if text == "all": + return None + if text in ("", "none", "off", "false"): + return frozenset() + out = set() + for chunk in text.split(","): + chunk = chunk.strip() + if not chunk: + continue + if "-" in chunk: + a, b = chunk.split("-", 1) + out.update(range(int(a), int(b) + 1)) + else: + out.add(int(chunk)) + return frozenset(i for i in out if 0 <= i < L_LLM) + return frozenset(int(i) for i in spec if 0 <= int(i) < L_LLM) + + def _resolve_fp4_ffn_layers(self, layers): + spec = os.environ.get("FLASHRT_CHAMELEON_FP4_LAYERS") \ + if layers is None else layers + if spec is None: + return frozenset() + parsed = self._parse_layer_spec(spec) + return frozenset() if parsed is None else parsed + + def _parse_layer_spec(self, spec): + if isinstance(spec, str): + text = spec.strip().lower() + if text == "all": + return frozenset(range(L_LLM)) + if text in ("", "none", "off", "false"): + return frozenset() + out = set() + for chunk in text.split(","): + chunk = chunk.strip() + if not chunk: + continue + if "-" in chunk: + a, b = chunk.split("-", 1) + out.update(range(int(a), int(b) + 1)) + else: + out.add(int(chunk)) + return frozenset(i for i in out if 0 <= i < L_LLM) + return frozenset(int(i) for i in spec if 0 <= int(i) < L_LLM) + + def _validate_config(self) -> None: + c = self.config_json + for k, want in (("hidden_size", D_LLM), + ("num_attention_heads", NH_LLM), + ("num_hidden_layers", L_LLM), + ("intermediate_size", DFF_LLM), + ("vocab_size", VOCAB_SIZE)): + if int(c[k]) != want: + raise ValueError(f"config {k}={c[k]}, expected {want}") + if bool(c.get("attention_bias", False)): + raise ValueError("standard Chameleon Thor path expects attention_bias=false") + if bool(c.get("mlp_bias", False)): + raise ValueError("standard Chameleon Thor path expects mlp_bias=false") + + def _build_image_token_mask(self) -> None: + """Image-codebook vocab ids to suppress for text generation. + + Mirrors HF ``ChameleonForConditionalGeneration``'s + ``mask_image_logits``: without this, greedy decode on a text + prompt can emit VQGAN codebook ids (garbage BPE decode) because + those ids are heavily represented in training and the raw + logit distribution favors them for many contexts. + """ + vocab_map = self.config_json.get("vocabulary_map", {}) + image_tokens = sorted( + v for k, v in vocab_map.items() if k.startswith("IMGIMG")) + self._mask_image_logits = bool(self.config_json.get("mask_image_logits", False)) + if image_tokens: + self._image_token_ids = torch.tensor( + image_tokens, dtype=torch.long, device="cuda") + else: + self._image_token_ids = None + self._mask_image_logits = False + + def _load_weights(self) -> None: + from flash_rt.executors.torch_weights import MultiSafetensorsSource, WeightLoader + from flash_rt.frontends.torch._chameleon_thor_spec import build_spec + + shard_paths = sorted(self.checkpoint_dir.glob("model-*-of-*.safetensors")) + if not shard_paths: + raise FileNotFoundError(f"No safetensors shards in {self.checkpoint_dir}") + src = MultiSafetensorsSource([str(p) for p in shard_paths], device="cuda") + WeightLoader(source=src, target=self, spec=build_spec(use_fp8=self._use_fp8)).run() + + self._split_fused_llm_weights() + self._lm_head_w_t = self._llm_lm_head_w.t().contiguous() + if self._use_fp8: + self._setup_fp8_weight_scales() + self._init_fp4_weight_lists() + if self._use_fp8 and self._fp4_ffn_layers: + self._quantize_fp4_weights() + + def _split_fused_llm_weights(self) -> None: + self._q_w, self._k_w, self._v_w = [], [], [] + self._gate_w, self._up_w = [], [] + for li in range(L_LLM): + qkv = self._llm_qkv_w[li] + self._q_w.append(qkv[:, :D_LLM].contiguous()) + self._k_w.append(qkv[:, D_LLM:2 * D_LLM].contiguous()) + self._v_w.append(qkv[:, 2 * D_LLM:].contiguous()) + + gu = self._llm_gu_w[li] + self._gate_w.append(gu[:, :DFF_LLM].contiguous()) + self._up_w.append(gu[:, DFF_LLM:].contiguous()) + + self._llm_q_norm_w[li] = self._llm_q_norm_w[li].reshape(-1).contiguous() + self._llm_q_norm_b[li] = self._llm_q_norm_b[li].reshape(-1).contiguous() + self._llm_k_norm_w[li] = self._llm_k_norm_w[li].reshape(-1).contiguous() + self._llm_k_norm_b[li] = self._llm_k_norm_b[li].reshape(-1).contiguous() + + def _setup_fp8_weight_scales(self) -> None: + self._llm_w_dev = torch.tensor(self._llm_w_scales, dtype=torch.float32, device="cuda") + base = self._llm_w_dev.data_ptr() + self._d_w_qkv_ptrs = [(base + (li * 4 + 0) * 4) for li in range(L_LLM)] + self._d_w_o_ptrs = [(base + (li * 4 + 1) * 4) for li in range(L_LLM)] + self._d_w_gu_ptrs = [(base + (li * 4 + 2) * 4) for li in range(L_LLM)] + self._d_w_d_ptrs = [(base + (li * 4 + 3) * 4) for li in range(L_LLM)] + + def _init_fp4_weight_lists(self) -> None: + self._gu_w_fp4 = [None] * L_LLM + self._gu_sfb = [None] * L_LLM + self._d_w_fp4 = [None] * L_LLM + self._d_sfb = [None] * L_LLM + + def _quantize_fp4_weights(self) -> None: + if fvk_fp4 is None: + logger.warning("flash_rt_fp4 unavailable; FP4 FFN disabled") + self._fp4_ffn_layers = frozenset() + return + try: + if not fvk_fp4.has_nvfp4(): + logger.warning("NVFP4 unavailable on this device; FP4 FFN disabled") + self._fp4_ffn_layers = frozenset() + return + except AttributeError: + pass + + from flash_rt.executors.torch_weights import MultiSafetensorsSource + + shard_paths = sorted(self.checkpoint_dir.glob("model-*-of-*.safetensors")) + src = MultiSafetensorsSource([str(p) for p in shard_paths], device="cuda") + + def _quant(w_fp16: torch.Tensor): + N, K = w_fp16.shape + packed = torch.empty(N, K // 2, dtype=torch.uint8, device="cuda") + sfb_bytes = fvk_fp4.sfa_size_bytes(N, K, True) + sfb = torch.empty(sfb_bytes, dtype=torch.uint8, device="cuda") + rc = fvk_fp4.quantize_fp4_dynamic_sfa_fp16( + w_fp16.data_ptr(), packed.data_ptr(), sfb.data_ptr(), + N, K, True, 0) + if rc != 0: + raise RuntimeError(f"quantize_fp4_dynamic_sfa_fp16 failed rc={rc}") + return packed, sfb + + packed_layers = set() + for li in range(L_LLM): + if li not in self._fp4_ffn_layers: + continue + gate_key = f"model.layers.{li}.mlp.gate_proj.weight" + up_key = f"model.layers.{li}.mlp.up_proj.weight" + down_key = f"model.layers.{li}.mlp.down_proj.weight" + gate_w = src.get(gate_key).to(fp16).contiguous() + up_w = src.get(up_key).to(fp16).contiguous() + down_w = src.get(down_key).to(fp16).contiguous() + gu_w = torch.cat([gate_w, up_w], dim=0).contiguous() + self._gu_w_fp4[li], self._gu_sfb[li] = _quant(gu_w) + self._d_w_fp4[li], self._d_sfb[li] = _quant(down_w) + packed_layers.add(li) + del gate_w, up_w, down_w, gu_w + torch.cuda.synchronize() + self._fp4_ffn_layers = frozenset(packed_layers) + logger.info("FP4 FFN weights quantized for standalone Chameleon layers: %s", + sorted(self._fp4_ffn_layers)) + + def _build_rope_tables(self) -> None: + inv_freq = 1.0 / (ROPE_THETA ** ( + torch.arange(0, HD_LLM, 2, dtype=torch.float32) / HD_LLM)) + pos = torch.arange(self._max_pos, dtype=torch.float32) + freqs = torch.outer(pos, inv_freq) + emb = torch.cat([freqs, freqs], dim=-1) + self._rope_cos = emb.cos().to(fp16).cuda().contiguous() + self._rope_sin = emb.sin().to(fp16).cuda().contiguous() + + def _load_tokenizer(self) -> None: + from transformers import AutoTokenizer + + candidates: list[pathlib.Path] = [] + if self.tokenizer_path: + candidates.append(pathlib.Path(self.tokenizer_path)) + env = os.environ.get("FLASHRT_CHAMELEON_TOKENIZER_DIR") + if env: + candidates.append(pathlib.Path(env)) + candidates.append(self.checkpoint_dir) + + for p in candidates: + if p.exists() and (p / "tokenizer.json").exists(): + self.tokenizer = AutoTokenizer.from_pretrained(str(p), use_fast=True) + self.tokenizer_dir = p + logger.info("Tokenizer loaded from %s", p) + return + raise FileNotFoundError("Chameleon tokenizer not found") + + def _load_vqgan(self) -> None: + from flash_rt.models.chameleon.vqvae_hf import load_chameleon_vqvae + + checkpoint = pathlib.Path(self.vqgan_path).expanduser() \ + if self.vqgan_path else self.checkpoint_dir + self._vqgan_model, self._vqgan_translation = load_chameleon_vqvae( + checkpoint, device="cuda", dtype=torch.float32) + logger.info("Transformers Chameleon VQ-VAE loaded from %s", checkpoint) + + @property + def vqgan_backend(self) -> str: + return self._vqgan_backend + + @property + def fa4_attn_active(self) -> bool: + return bool(getattr(self._attn, "_fa4_enabled", False)) + + def _ensure_trt_vqgan_loaded(self) -> bool: + """Lazy-instantiate optional TensorRT VQGAN acceleration. + + Generic standalone Chameleon defaults to eager Chameleon VQGAN + tokenization. TensorRT is an explicit opt-in framework acceleration + path for deployments that have compatible engines under the FlashRT + engine cache (or a caller-provided engine directory). If unavailable, + the frontend falls back to eager tokenization. + """ + if not self._use_trt_vqgan: + return False + if self._trt_vqgan_backend is not None: + return self._trt_vqgan_backend.is_available() + try: + from flash_rt.hardware.thor.vqgan_trt_backend import VQGANTRTBackend + except ImportError as e: + logger.warning("TRT VQGAN import failed: %s", e) + self._use_trt_vqgan = False + self._vqgan_backend = "eager" + return False + if torch.backends.cudnn.allow_tf32: + logger.info("TRT VQGAN: disabling cuDNN TF32 globally.") + torch.backends.cudnn.allow_tf32 = False + torch.backends.cuda.matmul.allow_tf32 = False + engine_dir = pathlib.Path(self._trt_vqgan_engine_dir) \ + if self._trt_vqgan_engine_dir else None + self._trt_vqgan_backend = VQGANTRTBackend(engine_dir=engine_dir) + available = self._trt_vqgan_backend.is_available() + if not available: + logger.warning( + "TRT VQGAN backend unavailable at %s; falling back to " + "eager PyTorch encode.", + engine_dir or VQGANTRTBackend.ENGINE_DIR) + self._use_trt_vqgan = False + self._vqgan_backend = "eager" + else: + self._vqgan_backend = "trt" + return available + + def _preprocess_image_for_trt(self, pil_image, out_hw): + """PIL uint8 -> CUDA [1,3,H,W] float32 [-1,1] (matches the HF reference preprocessing).""" + H_out, W_out = out_hw + if pil_image.size != (W_out, H_out): + pil_image = pil_image.resize((W_out, H_out), resample=PIL.Image.BICUBIC) + np_img = np.array(pil_image.convert("RGB")) / 255.0 + np_img = np_img * 2.0 - 1.0 + t = torch.from_numpy(np_img).permute(2, 0, 1).unsqueeze(0) + return t.to(dtype=torch.float32, device="cuda").contiguous() + + def _vqgan_encode(self, image) -> list[int]: + if isinstance(image, np.ndarray): + image = _PILImage.fromarray(image.astype(np.uint8)).convert("RGB") + elif not isinstance(image, PIL.Image.Image): + raise TypeError(f"image must be PIL.Image or np.ndarray, got {type(image)!r}") + + # ── TRT fast path ── + if self._ensure_trt_vqgan_loaded(): + H_eng = W_eng = int(self.target_size) + if self._trt_vqgan_backend.supports_resolution(H_eng, W_eng): + with torch.cuda.stream(self._trt_stream): + img = self._preprocess_image_for_trt(image, (H_eng, W_eng)) + indices = self._trt_vqgan_backend.encode(img) + if indices is not None: + latent_ids = indices.view(-1) + global_ids = self._vqgan_translation.convert_img2bpe( + latent_ids).view(-1) + h_lat = H_eng // 16 + w_lat = W_eng // 16 + h_grids = H_eng // PATCH_SIZE + w_grids = W_eng // PATCH_SIZE + grid = global_ids.view(h_lat, w_lat) + newline_col = torch.full( + (h_lat, 1), NEWLINE_ID, dtype=grid.dtype, device=grid.device) + with_nl = torch.cat([grid, newline_col], dim=1).flatten().tolist() + return [IMG_START_ID, GRID_TOK_BASE + h_grids, GRID_TOK_BASE + w_grids, + *with_nl, IMG_END_ID] + + # ── Eager PyTorch fallback ── + # data_lerobot's aspect-preserving center crop when importable; + # otherwise a self-contained re-implementation of the same + # selection (max-coverage crop from the aspect-varied size list), + # so the eager path never depends on external packages. + try: + from data_lerobot.item_processor import ( # type: ignore + var_center_crop, generate_crop_size_list) + except ImportError: + def generate_crop_size_list(num_patches, patch_size, max_ratio=4.0): + crop_size_list = [] + wp, hp = num_patches, 1 + while wp > 0: + if max(wp, hp) / min(wp, hp) <= max_ratio: + crop_size_list.append((wp * patch_size, hp * patch_size)) + if (hp + 1) * wp <= num_patches: + hp += 1 + else: + wp -= 1 + return crop_size_list + + def var_center_crop(image, crop_size_list=None): + crop_size_list = crop_size_list or [(image.size[0], image.size[1])] + w, h = image.size + rem = [min(cw / w, ch / h) / max(cw / w, ch / h) + for cw, ch in crop_size_list] + best = sorted(zip(rem, crop_size_list), reverse=True)[0][1] + left = max(0, (w - best[0]) // 2) + top = max(0, (h - best[1]) // 2) + return image.crop((left, top, left + best[0], top + best[1])) + + crop_size_list = generate_crop_size_list( + (self.target_size // PATCH_SIZE) ** 2, PATCH_SIZE) + cropped = var_center_crop(image, crop_size_list=crop_size_list) + from flash_rt.models.chameleon.vqvae_hf import ( + encode_vqvae_tokens, + preprocess_vqvae_image, + ) + + pixels = preprocess_vqvae_image( + cropped, device="cuda", dtype=torch.float32) + with torch.no_grad(): + latent_ids = encode_vqvae_tokens(self._vqgan_model, pixels) + global_ids = self._vqgan_translation.convert_img2bpe(latent_ids).view(-1) + + w_grids = cropped.size[0] // PATCH_SIZE + h_grids = cropped.size[1] // PATCH_SIZE + w_lat = cropped.size[0] // 16 + h_lat = cropped.size[1] // 16 + grid = global_ids.view(h_lat, w_lat) + newline_col = torch.full((h_lat, 1), NEWLINE_ID, dtype=grid.dtype, device=grid.device) + with_nl = torch.cat([grid, newline_col], dim=1).flatten().tolist() + return [IMG_START_ID, GRID_TOK_BASE + h_grids, GRID_TOK_BASE + w_grids, + *with_nl, IMG_END_ID] + + def _allocate_buffers(self) -> None: + # Capacity is floored to a multiple of 16 so the pad-to-16 in + # set_prompt can never overshoot the allocated buffers/KV cache when + # max_seq itself is not a multiple of 16. + Se = (self._max_pos // 16) * 16 + D = D_LLM + Dff = DFF_LLM + self._Se_max = Se + self._bufs: dict[str, torch.Tensor] = {} + + def _alloc(name, shape, dtype=fp16): + self._bufs[name] = torch.zeros(shape, dtype=dtype, device="cuda") + + _alloc("x", (Se, D)) + _alloc("xn", (Se, D)) + _alloc("xn_fp8", (Se, D), fp8) + _alloc("o_proj_out", (Se, D)) + _alloc("hidden_all", (Se, D)) + _alloc("gate_out", (Se, Dff)) + _alloc("up_out", (Se, Dff)) + _alloc("gu_fp8", (Se, Dff), fp8) + _alloc("zero_bias_d", (D,)) + _alloc("zero_bias_dff", (Dff,)) + _alloc("act_fp4", (Se * D // 2,), torch.uint8) + _alloc("act_sfa", (Se * D // 16 * 2,), torch.uint8) + _alloc("ffn_act_fp4", (Se * Dff // 2,), torch.uint8) + _alloc("ffn_act_sfa", (Se * Dff // 16 * 2,), torch.uint8) + _alloc("gu_merged", (Se, 2 * Dff)) + _alloc("dyn_act_scales", (L_LLM * 4,), torch.float32) + _alloc("last_logits", (VOCAB_SIZE,)) + + self._llm_calib_scales = torch.ones(L_LLM * 4, dtype=torch.float32, device="cuda") + self._kv_cache = torch.zeros(L_LLM, 2, Se, D, dtype=fp16, device="cuda") + self._kv_layer_stride = 2 * Se * D * 2 + logits_sz = max(NH_LLM * Se * Se, 4) + self._logits_buf = torch.zeros(logits_sz, dtype=fp16, device="cuda") + + def _build_attention_backend(self) -> None: + spec = make_chameleon_attention_spec(seq_max=self._Se_max) + kv_base = self._kv_cache.data_ptr() + se_d_bytes = self._Se_max * D_LLM * 2 + chameleon_slots = { + "Q_O": self._bufs["xn"].data_ptr(), + "Kc": kv_base, + "Vc": kv_base + se_d_bytes, + "logits": self._logits_buf.data_ptr(), + "layer_stride": self._kv_layer_stride, + "scale": 1.0 / math.sqrt(HD_LLM), + } + self._attn = ThorChameleonAttnBackend( + spec, self._ctx, chameleon_slots=chameleon_slots) + + def _build_llm_weights(self) -> dict: + w = { + "input_ln_w": [x.data_ptr() for x in self._llm_input_ln_w], + "post_ln_w": [x.data_ptr() for x in self._llm_post_ln_w], + "q_w": [x.data_ptr() for x in self._q_w], + "k_w": [x.data_ptr() for x in self._k_w], + "v_w": [x.data_ptr() for x in self._v_w], + "o_w": [x.data_ptr() for x in self._llm_o_w], + "gate_w": [x.data_ptr() for x in self._gate_w], + "up_w": [x.data_ptr() for x in self._up_w], + "d_w": [x.data_ptr() for x in self._llm_d_w], + "q_norm_w": [x.data_ptr() for x in self._llm_q_norm_w], + "q_norm_b": [x.data_ptr() for x in self._llm_q_norm_b], + "k_norm_w": [x.data_ptr() for x in self._llm_k_norm_w], + "k_norm_b": [x.data_ptr() for x in self._llm_k_norm_b], + "o_b": [self._bufs["zero_bias_d"].data_ptr()] * L_LLM, + "final_norm_w": self._llm_norm_w.data_ptr(), + "rope_cos": self._rope_cos.data_ptr(), + "rope_sin": self._rope_sin.data_ptr(), + } + if self._use_fp8: + w.update({ + "w_scales_flat": self._llm_w_dev.data_ptr(), + "d_w_qkv": self._d_w_qkv_ptrs, + "d_w_o": self._d_w_o_ptrs, + "d_w_gu": self._d_w_gu_ptrs, + "d_w_d": self._d_w_d_ptrs, + "alpha_host": [1.0] * (L_LLM * 4), + "gu_w_fp4": [ + x.data_ptr() if x is not None else 0 for x in self._gu_w_fp4], + "gu_sfb": [ + x.data_ptr() if x is not None else 0 for x in self._gu_sfb], + "d_w_fp4": [ + x.data_ptr() if x is not None else 0 for x in self._d_w_fp4], + "d_sfb": [ + x.data_ptr() if x is not None else 0 for x in self._d_sfb], + }) + return w + + def _build_llm_bufs(self) -> dict: + return {k: self._bufs[k].data_ptr() for k in ( + "x", "xn", "xn_fp8", "o_proj_out", "hidden_all", + "gate_out", "up_out", "gu_fp8", "zero_bias_d", "zero_bias_dff", + "act_fp4", "act_sfa", "ffn_act_fp4", "ffn_act_sfa", "gu_merged", + "dyn_act_scales", + )} + + def _build_llm_scales_dev(self) -> dict: + calib_base = self._llm_calib_scales.data_ptr() + dyn_base = self._bufs["dyn_act_scales"].data_ptr() + return { + "act_qkv": [(calib_base + (li * 4 + 0) * 4) for li in range(L_LLM)], + "act_o": [(calib_base + (li * 4 + 1) * 4) for li in range(L_LLM)], + "act_gu": [(calib_base + (li * 4 + 2) * 4) for li in range(L_LLM)], + "act_down": [(calib_base + (li * 4 + 3) * 4) for li in range(L_LLM)], + "dyn_act_qkv": [(dyn_base + (li * 4 + 0) * 4) for li in range(L_LLM)], + "dyn_act_o": [(dyn_base + (li * 4 + 1) * 4) for li in range(L_LLM)], + "dyn_act_gu": [(dyn_base + (li * 4 + 2) * 4) for li in range(L_LLM)], + "dyn_act_down": [(dyn_base + (li * 4 + 3) * 4) for li in range(L_LLM)], + } + + def encode_prompt(self, text: str, images: Optional[list] = None) -> list[int]: + images = images or [] + chunks = text.split("") + if len(chunks) - 1 not in (0, len(images)): + raise ValueError("number of placeholders must be 0 or match images") + ids: list[int] = [] + bos = getattr(self.tokenizer, "bos_token_id", None) + if bos is not None: + ids.append(int(bos)) + for i, chunk in enumerate(chunks): + if chunk: + ids.extend(self.tokenizer.encode(chunk, add_special_tokens=False)) + if i < len(chunks) - 1: + ids.extend(self._vqgan_encode(images[i])) + if len(chunks) == 1: + for image in images: + ids.extend(self._vqgan_encode(image)) + return ids + + def _embed_ids(self, input_ids: list[int]) -> None: + Se = len(input_ids) + ids_t = torch.tensor(input_ids, dtype=torch.long, device="cuda") + emb = torch.nn.functional.embedding(ids_t, self._llm_embed_w) + self._bufs["x"].zero_() + _cudart.cudaMemcpyAsync( + ctypes.c_void_p(self._bufs["x"].data_ptr()), + ctypes.c_void_p(emb.data_ptr()), + Se * D_LLM * 2, + 3, + ctypes.c_void_p(self._stream), + ) + + def set_prompt(self, text: str, images: Optional[list] = None) -> list[int]: + input_ids = self.encode_prompt(text, images) + self._real_len = len(input_ids) + rem = len(input_ids) % 16 + padded_len = len(input_ids) + ((16 - rem) if rem else 0) + if padded_len > self._Se_max: + raise ValueError( + f"padded sequence length {padded_len} exceeds max_seq=" + f"{self._Se_max} (prompt has {len(input_ids)} tokens)") + if rem: + input_ids.extend([PAD_ID] * (16 - rem)) + self.Se = len(input_ids) + self._last_input_ids = input_ids + if self._use_autotune: + self._autotune_gemms(self.Se) + self._embed_ids(input_ids) + if self._use_cuda_graph: + self._capture_graph(self.Se) + return input_ids + + def _autotune_gemms(self, Se: int, num_algos: int = 16) -> None: + """Per-shape cuBLASLt algo autotune (motus/hyvla pattern). + + Chameleon's per-layer GEMMs collapse to 3 distinct (M,N,K) shapes + at a given Se (q/k/v/o share one shape, gate/up share another, + down is the third) plus the M=1 lm_head projection. Tuning each + shape once mutates the GemmRunner's internal algo cache (keyed on + (M,N,K)); every subsequent real call with that shape — including + inside a captured CUDA graph — picks up the tuned algo for free. + Dummy buffers are only used for timing, not correctness. + """ + if Se in self._autotuned_se: + return + D, Dff = D_LLM, DFF_LLM + dev = "cuda" + if self._use_fp8 and hasattr(self._gemm, "autotune_fp8_nn_dev_fp16"): + shapes = [(Se, D, D), (Se, Dff, D), (Se, D, Dff)] + shapes += [(1, D, D), (1, Dff, D), (1, D, Dff)] # decode + for (M, N, K) in dict.fromkeys(shapes): + A = torch.empty(M, K, dtype=torch.uint8, device=dev) + B = torch.empty(K, N, dtype=torch.uint8, device=dev) + Dbuf = torch.empty(M, N, dtype=fp16, device=dev) + sa = torch.ones(1, dtype=torch.float32, device=dev) + sb = torch.ones(1, dtype=torch.float32, device=dev) + self._gemm.autotune_fp8_nn_dev_fp16( + A.data_ptr(), B.data_ptr(), Dbuf.data_ptr(), + M, N, K, sa.data_ptr(), sb.data_ptr(), num_algos) + elif hasattr(self._gemm, "autotune_fp16_nn"): + shapes = [(Se, D, D), (Se, Dff, D), (Se, D, Dff)] + shapes += [(1, D, D), (1, Dff, D), (1, D, Dff)] # decode + for (M, N, K) in dict.fromkeys(shapes): + A = torch.empty(M, K, dtype=fp16, device=dev) + B = torch.empty(K, N, dtype=fp16, device=dev) + Dbuf = torch.empty(M, N, dtype=fp16, device=dev) + self._gemm.autotune_fp16_nn( + A.data_ptr(), B.data_ptr(), Dbuf.data_ptr(), M, N, K, num_algos) + if hasattr(self._gemm, "autotune_fp16_nn"): + A = torch.empty(1, D, dtype=fp16, device=dev) + B = torch.empty(D, VOCAB_SIZE, dtype=fp16, device=dev) + Dbuf = torch.empty(1, VOCAB_SIZE, dtype=fp16, device=dev) + self._gemm.autotune_fp16_nn( + A.data_ptr(), B.data_ptr(), Dbuf.data_ptr(), 1, VOCAB_SIZE, D, num_algos) + torch.cuda.synchronize() + self._autotuned_se.add(Se) + + def _capture_graph(self, Se: int) -> None: + if self._infer_graph is not None and self._captured_Se == Se: + return + capture_stream = torch.cuda.Stream() + prev_stream_id = self._stream + with torch.cuda.stream(capture_stream): + self._stream = capture_stream.cuda_stream + self._run_backbone(Se) + torch.cuda.synchronize() + # The warmup forward mutates x into the final residual stream; the + # capture pass (and every later replay) must start from clean + # embeddings — see _replay_backbone. + self._embed_ids(self._last_input_ids) + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph, stream=capture_stream): + self._run_backbone(Se) + self._stream = prev_stream_id + self._infer_graph = graph + self._captured_Se = Se + + def _run_backbone(self, Se: int) -> None: + dims = {"Se": Se, "D": D_LLM, "Dff": DFF_LLM, + "L": L_LLM, "H": NH_LLM, "Hd": HD_LLM} + if self._use_fp8: + chameleon_forward( + self._gemm, fvk, self._build_llm_bufs(), self._build_llm_weights(), + dims, self._build_llm_scales_dev(), attn=self._attn, + stream=self._stream, dynamic_fp8_layers=frozenset(range(L_LLM)), + fp4_ffn_layers=self._fp4_ffn_layers, ffn_down_clamp_value=60000.0, + ffn_clamp_layers=self._ffn_clamp_layers, + ) + else: + chameleon_forward_fp16( + self._gemm, fvk, self._build_llm_bufs(), self._build_llm_weights(), + dims, attn=self._attn, stream=self._stream, + ) + + def _project_last(self, last_idx: Optional[int] = None) -> None: + if last_idx is None: + last_idx = self._real_len - 1 + last_hidden_ptr = self._bufs["hidden_all"].data_ptr() + last_idx * D_LLM * 2 + self._gemm.fp16_nn( + last_hidden_ptr, self._lm_head_w_t.data_ptr(), + self._bufs["last_logits"].data_ptr(), 1, VOCAB_SIZE, D_LLM, + int(self._stream), + ) + if self._mask_image_logits and self._image_token_ids is not None: + self._bufs["last_logits"].index_fill_(0, self._image_token_ids, -65504.0) + + def _run_forward(self, Se: int) -> None: + self._run_backbone(Se) + self._project_last() + + def _replay_backbone(self) -> None: + """Replay the captured prefill graph over clean embeddings. + + Every backbone run mutates ``x`` in place into the final residual + stream, so a replay must re-embed first or it recomputes over the + previous run's residuals. + """ + self._embed_ids(self._last_input_ids) + self._infer_graph.replay() + + def prefill(self, text: str, images: Optional[list] = None) -> dict: + input_ids = self.set_prompt(text, images) + if self._use_cuda_graph and self._infer_graph is not None: + self._replay_backbone() + self._project_last() + else: + self._run_forward(self.Se) + torch.cuda.synchronize() + return { + "input_ids": input_ids, + "Se": self.Se, + "vqgan_backend": self._vqgan_backend, + "fa4_attn": self.fa4_attn_active, + "logits": self._bufs["last_logits"].detach().float().cpu(), + "hidden": self._bufs["hidden_all"][:self.Se].detach().float().cpu(), + } + + def generate_greedy(self, text: str, images: Optional[list] = None, + max_new_tokens: int = 16, + eos_token_id: Optional[int] = None) -> dict: + """Greedy generation with incremental KV-cache decode. + + One prefill over the prompt, then single-token decode steps + (``chameleon_decode_step``) — O(n) instead of the historical O(n^2) + full recompute. Decode runs eagerly (no CUDA graph) because the + position is a host scalar baked into RoPE offsets and cache rows. + """ + if not self._use_fp8: + raise NotImplementedError( + "incremental decode requires the dynamic-FP8 path " + "(use_fp8=True)") + if self._fp4_ffn_layers: + raise NotImplementedError( + "incremental decode does not model the NVFP4 FFN tiers " + "(fp4_ffn_layers); use the prefill-only path") + + self.set_prompt(text, images) + generated = list(self._last_input_ids[:self._real_len]) + eos = EOS_ID if eos_token_id is None else int(eos_token_id) + budget = int(max_new_tokens) + if budget < 0: + raise ValueError( + f"max_new_tokens must be >= 0, got {max_new_tokens}") + if budget == 0: + return {"input_ids": generated, + "text": self.tokenizer.decode(generated)} + + # Prefill (graph or eager) + first token from the last prompt row. + if self._use_cuda_graph and self._infer_graph is not None: + self._replay_backbone() + else: + self._run_backbone(self.Se) + self._project_last() + torch.cuda.synchronize() + next_id = int(torch.argmax(self._bufs["last_logits"]).item()) + generated.append(next_id) + + # Decode-step plumbing: all pointers are stable across steps, so + # build the dicts once; only `pos` changes per token. + dims = {"Se": 1, "D": D_LLM, "Dff": DFF_LLM, + "L": L_LLM, "H": NH_LLM, "Hd": HD_LLM} + bufs = self._build_llm_bufs() + weights = self._build_llm_weights() + scales_dev = self._build_llm_scales_dev() + pos = self._real_len + for _ in range(budget - 1): + if next_id == eos or pos >= self._Se_max: + break + # Decode state: single-token embedding in residual row 0. + self._bufs["x"][:1].copy_(self._llm_embed_w[next_id]) + chameleon_decode_step( + self._gemm, fvk, bufs, weights, dims, scales_dev, + attn=self._attn, pos=pos, stream=int(self._stream), + ffn_down_clamp_value=60000.0, + ffn_clamp_layers=self._ffn_clamp_layers, + ) + self._project_last(last_idx=0) + torch.cuda.synchronize() + next_id = int(torch.argmax(self._bufs["last_logits"]).item()) + generated.append(next_id) + pos += 1 + return {"input_ids": generated, "text": self.tokenizer.decode(generated)} + + def _generate_greedy_recompute(self, text: str, + images: Optional[list] = None, + max_new_tokens: int = 16) -> dict: + """Legacy O(n^2) full-recompute greedy path (oracle for decode). + + Always eager: it shares the KV cache with the incremental path, + and graph capture/replay of every growing-Se forward would + overwrite cache rows for pad-16 filler positions. + """ + ids = self.encode_prompt(text, images) + generated = list(ids) + for _ in range(int(max_new_tokens)): + if len(generated) >= self._Se_max: + break + self._real_len = len(generated) + padded = list(generated) + rem = len(padded) % 16 + if rem: + padded.extend([PAD_ID] * (16 - rem)) + self.Se = len(padded) + self._last_input_ids = padded + if self._use_autotune: + self._autotune_gemms(self.Se) + self._embed_ids(padded) + self._run_forward(self.Se) + torch.cuda.synchronize() + next_id = int(torch.argmax(self._bufs["last_logits"]).item()) + generated.append(next_id) + return {"input_ids": generated, "text": self.tokenizer.decode(generated)} + + +__all__ = ["ChameleonTorchFrontendThor"] diff --git a/flash_rt/hardware/__init__.py b/flash_rt/hardware/__init__.py index 5d231ead6..c3bd13722 100644 --- a/flash_rt/hardware/__init__.py +++ b/flash_rt/hardware/__init__.py @@ -159,6 +159,24 @@ def detect_arch() -> str: ("flash_rt.frontends.torch.qwen3_vl_rtx_bf16", "Qwen3VlTorchFrontendRtxBF16"), + # ── Chameleon-7B ── + # Direct frontend (set_prompt/generate), not the VLA predict() surface. + # Orin SM87: INT8/INT4+QuaRot-Hadamard path; compute in + # flash_rt/models/chameleon/pipeline_rtx.py. Registered for resolver / + # direct-construction discovery only; load_model(config="chameleon") + # raises a redirect because this exposes set_prompt() + generate(), + # not the VLA predict() surface. See docs/chameleon7b_rtx_sm87.md. + ("chameleon", "torch", "rtx_sm87"): + ("flash_rt.frontends.torch.chameleon_rtx_sm87", + "ChameleonTorchFrontendRtxSm87"), + # Thor SM110: dynamic-FP8 backbone (optional NVFP4 FFN), attention via + # the dedicated Chameleon Thor backend (FA4 -> CUTLASS causal FMHA -> + # cuBLAS fallback). Compute in flash_rt/models/chameleon/pipeline_thor.py. + # See docs/chameleon_usage.md. + ("chameleon", "torch", "thor"): + ("flash_rt.frontends.torch.chameleon_thor", + "ChameleonTorchFrontendThor"), + # Cosmos3-Edge official Thor baseline. ("cosmos3_edge", "torch", "thor"): ("flash_rt.frontends.torch.cosmos3_edge_thor", "Cosmos3EdgeTorchFrontendThor"), @@ -195,6 +213,7 @@ def detect_arch() -> str: # resolving here and crashing later at the first kernel launch. _SM87_ALLOWED = { ("pi05", "torch", "rtx_sm87"), + ("chameleon", "torch", "rtx_sm87"), ("qwen3_vl", "torch", "rtx_sm87"), } diff --git a/flash_rt/hardware/rtx/attn_backend_chameleon.py b/flash_rt/hardware/rtx/attn_backend_chameleon.py new file mode 100644 index 000000000..8aafba63b --- /dev/null +++ b/flash_rt/hardware/rtx/attn_backend_chameleon.py @@ -0,0 +1,237 @@ +"""FlashRT — Chameleon-7B VLM attention backend for Jetson Orin (SM87). + +This backend owns a **real per-layer KV cache** so the LLM can decode +autoregressively (a prefill-only backend could instead share one K/V scratch +across all 32 layers, ``layer_stride = 0``). + +Two load-bearing design points, both measured — see +``docs/chameleon7b_rtx_sm87.md`` §2.1 and §3.1: + +1. **The K/V GEMMs write straight into the cache.** CUTLASS hard-wires the + output row stride to ``N`` (``cutlass_sm80_int8_rowwise_fp16out.cu:169-171``), + and a per-layer slab of ``[max_seq, num_kv_heads*head_dim]`` has exactly that + row stride, so no staging buffer or copy is needed for either prefill + (``M = S`` at the slab base) or decode (``M = 1`` at ``+ pos*row_stride``). + ``qk_norm_rope_fused_fp16`` then does QK-LayerNorm + RoPE in place. + +2. **Split-KV must be forced on with a biased ``num_sms``.** FA2's heuristic + (``fa2_wrapper_causal.cu:41-43,152-158``) returns ``num_splits = 1`` whenever + ``batch*num_q_heads*num_m_blocks >= 0.8 * (num_sms*2)``. Chameleon decode is + ``1*32*1 = 32`` against ``0.8*32 = 25.6``, so at the true SM count split-KV + silently does nothing (measured: bit-identical output, 1.05x). ``num_sms`` is + a pure heuristic knob in this wrapper, so ``split_kv_bias`` multiplies it; + bias 4 measured 204.9 -> 141.8 us (**1.44x**) at fp16-rounding-level delta. + This is why the Qwen3-VL split-KV lever does not transfer as-is — + that model has 16 Q heads and lands under the threshold naturally. + +FA2 ``fwd_fp16_causal`` is **mandatory** for decode: its causal mask is +bottom-right aligned (``fa2_wrapper_causal.cu:126-138``) so ``q=1, kv=N`` +attends all N keys. The cuBLAS ``attention_mha_causal_fp16`` fallback is **top-left** aligned (``softmax.cu:182-191`` masks with +``q = row % S_q``), so at ``S_q = 1`` only column 0 survives — it is silently +wrong rather than merely slow. This backend therefore raises instead of +degrading to it. +""" + +from __future__ import annotations + +import logging + +import torch + +from flash_rt.hardware.backend import AttentionBackendBase, AttentionSpec + +logger = logging.getLogger(__name__) + +SITE = "llm" + +#: FA2's own cap on the split count (``fa2_wrapper_causal.cu`` passes 128). +_MAX_SPLITS = 128 + + +class ChameleonAttnBackend(AttentionBackendBase): + """Chameleon-7B self-attention with a per-layer FP16 KV cache. + + Owns every buffer it needs (KV cache, Q/O staging, softmax LSE, split-KV + accumulators). Pointers are stable for the object's lifetime, so they are + safe to bake into a captured CUDA graph — but the backend must be kept + alive for as long as any such graph exists. + """ + + def __init__(self, spec: AttentionSpec, *, max_seq: int, + split_kv_bias: int = 4) -> None: + super().__init__(spec) + if set(spec.sites.keys()) != {SITE}: + raise ValueError( + f"ChameleonAttnBackend expects exactly the {SITE!r} site, " + f"got {set(spec.sites.keys())}") + + s = spec.site(SITE) + self.num_layers = int(s.num_layers) + self.num_q_heads_ = int(s.num_q_heads) + self.num_kv_heads_ = int(s.num_kv_heads) + self.head_dim_ = int(s.head_dim) + self.max_seq = int(max_seq) + self.split_kv_bias = int(split_kv_bias) + self.scale = float(self.head_dim_) ** -0.5 + + if self.head_dim_ != 128: + raise ValueError( + f"head_dim must be 128 (FA2 fp16 causal aborts otherwise, " + f"fa2_wrapper_causal.cu:235-243); got {self.head_dim_}") + + try: + import flash_rt.flash_rt_fa2 as _fa2 + except ImportError as e: # pragma: no cover + raise RuntimeError( + "Chameleon decode requires flash_rt_fa2 (the cuBLAS MHA " + "fallback is top-left-aligned causal and therefore WRONG at " + "q_seq=1). Rebuild with -DFLASHRT_ENABLE_FA2=ON.") from e + if not hasattr(_fa2, "fwd_fp16_causal"): + raise RuntimeError( + "flash_rt_fa2 lacks fwd_fp16_causal; Chameleon decode cannot " + "fall back to cuBLAS MHA (top-left-aligned mask is wrong at " + "q_seq=1). Rebuild with -DFA2_DTYPES='fp16;bf16' " + "-DFA2_HDIMS='128;256'.") + self._fa2_fwd = _fa2.fwd_fp16_causal + + dev, fp16, fp32 = "cuda", torch.float16, torch.float32 + kv_row = self.num_kv_heads_ * self.head_dim_ # 4096 == GEMM N + q_row = self.num_q_heads_ * self.head_dim_ + + # Per-layer KV cache. The [max_seq, kv_row] slab per layer is exactly a + # legal CUTLASS destination (row stride == N), which is what lets the + # K/V GEMMs write into it directly. + self.K_cache = torch.zeros(self.num_layers, self.max_seq, kv_row, + dtype=fp16, device=dev) + self.V_cache = torch.zeros(self.num_layers, self.max_seq, kv_row, + dtype=fp16, device=dev) + # Q staging (prefill writes S rows, decode row 0). O aliases Q, matching + # the caller convention: the pipeline reads its result back in place. + self.Q_buf = torch.zeros(self.max_seq, q_row, dtype=fp16, device=dev) + + lse_rows = ((self.max_seq + 127) // 128) * 128 + self.lse_buf = torch.zeros(1, self.num_q_heads_, lse_rows, + dtype=fp32, device=dev) + # Split-KV accumulators, sized for the decode case (seqlen_q == 1). + # Empty splits are self-initialised by the kernel, so no pre-fill. + self.lse_accum = torch.zeros(_MAX_SPLITS, 1, self.num_q_heads_, 1, + dtype=fp32, device=dev) + self.o_accum = torch.zeros(_MAX_SPLITS, 1, self.num_q_heads_, 1, + self.head_dim_, dtype=fp32, device=dev) + + self._num_sms = torch.cuda.get_device_properties( + torch.cuda.current_device()).multi_processor_count + + kv_gb = 2 * self.K_cache.numel() * 2 / 2 ** 30 + logger.info( + "ChameleonAttnBackend: L=%d %dQ/%dKV hd=%d max_seq=%d " + "KV=%.2f GB split_kv_bias=%d (num_sms %d->%d)", + self.num_layers, self.num_q_heads_, self.num_kv_heads_, + self.head_dim_, self.max_seq, kv_gb, self.split_kv_bias, + self._num_sms, self._num_sms * self.split_kv_bias) + + # ------------------------------------------------------------------ + # Layout / pointers + # ------------------------------------------------------------------ + + @property + def kv_layer_stride_bytes(self) -> int: + return self.max_seq * self.num_kv_heads_ * self.head_dim_ * 2 + + @property + def kv_row_stride_bytes(self) -> int: + return self.num_kv_heads_ * self.head_dim_ * 2 + + def _check_layer(self, layer_idx: int) -> None: + if not 0 <= layer_idx < self.num_layers: + raise IndexError( + f"layer_idx {layer_idx} out of range [0, {self.num_layers})") + + def get_slot_ptrs(self, site: str, layer_idx: int) -> dict: + """Prefill slots: Q/O staging plus the base of this layer's KV slab.""" + if site != SITE: + raise KeyError(f"unknown site {site!r}") + self._check_layer(layer_idx) + off = layer_idx * self.kv_layer_stride_bytes + q = self.Q_buf.data_ptr() + return {"Q": q, "O": q, + "K": self.K_cache.data_ptr() + off, + "V": self.V_cache.data_ptr() + off} + + def kv_row_ptrs(self, layer_idx: int, pos: int) -> tuple: + """Decode slots: the single KV row for absolute position ``pos``.""" + self._check_layer(layer_idx) + if not 0 <= pos < self.max_seq: + raise IndexError( + f"pos {pos} out of range [0, max_seq={self.max_seq}); raise " + f"max_seq at construction time") + off = layer_idx * self.kv_layer_stride_bytes + pos * self.kv_row_stride_bytes + return (self.K_cache.data_ptr() + off, self.V_cache.data_ptr() + off) + + # ------------------------------------------------------------------ + # Attention + # ------------------------------------------------------------------ + + def _fa2(self, layer_idx: int, q_seq: int, kv_seq: int, *, + use_split_kv: bool, stream: int) -> int: + qr = self.num_q_heads_ * self.head_dim_ + kr = self.num_kv_heads_ * self.head_dim_ + off = layer_idx * self.kv_layer_stride_bytes + q = self.Q_buf.data_ptr() + # num_sms == 0 disables the split-KV heuristic entirely; a biased value + # is the only way to reach num_splits > 1 at 32 Q heads (see module doc). + num_sms = self._num_sms * self.split_kv_bias if use_split_kv else 0 + self._fa2_fwd( + q, self.K_cache.data_ptr() + off, self.V_cache.data_ptr() + off, q, + self.lse_buf.data_ptr(), + self.lse_accum.data_ptr() if use_split_kv else 0, + self.o_accum.data_ptr() if use_split_kv else 0, + batch=1, seqlen_q=q_seq, seqlen_k=kv_seq, + num_heads_q=self.num_q_heads_, num_heads_kv=self.num_kv_heads_, + head_dim=self.head_dim_, + q_strides=(q_seq * qr, qr, self.head_dim_), + k_strides=(kv_seq * kr, kr, self.head_dim_), + v_strides=(kv_seq * kr, kr, self.head_dim_), + o_strides=(q_seq * qr, qr, self.head_dim_), + softmax_scale=self.scale, num_sms=num_sms, stream=stream) + return q + + def run_prefill(self, layer_idx: int, seq_len: int, *, stream: int = 0) -> int: + """Square causal attention over ``seq_len`` tokens. Result lands in Q_buf.""" + self._check_layer(layer_idx) + if seq_len > self.max_seq: + raise ValueError(f"seq_len {seq_len} > max_seq {self.max_seq}") + # Prefill already fills every SM (num_m_blocks = ceil(S/64)), so + # splitting KV would only add a combine pass. + return self._fa2(layer_idx, seq_len, seq_len, + use_split_kv=False, stream=stream) + + def run_decode(self, layer_idx: int, pos: int, *, stream: int = 0) -> int: + """One query row attending keys ``[0, pos]``. Result lands in Q_buf row 0.""" + self._check_layer(layer_idx) + return self._fa2(layer_idx, 1, pos + 1, + use_split_kv=self.split_kv_bias > 1, stream=stream) + + # No generic ``run(site, layer_idx, q_seq, ...)`` on purpose. Prefill and + # decode differ in more than q_seq here (split-KV on/off, and the caller + # must supply the absolute KV position rather than a length), so a shim that + # inferred the mode from ``q_seq == 1`` would be an untested footgun. The + # inherited ``AttentionBackendBase.run`` raises NotImplementedError, which + # is the behaviour we want if a generic caller ever appears. + + def reset_cache(self) -> None: + self.K_cache.zero_() + self.V_cache.zero_() + + +def make_chameleon_attention_spec(*, num_layers: int, num_q_heads: int, + num_kv_heads: int, head_dim: int, + max_seq: int) -> AttentionSpec: + spec = AttentionSpec() + spec.add_site(SITE, num_layers=num_layers, num_q_heads=num_q_heads, + num_kv_heads=num_kv_heads, head_dim=head_dim, + max_q_seq=max_seq, max_kv_seq=max_seq, causal=True) + return spec + + +__all__ = ["ChameleonAttnBackend", "make_chameleon_attention_spec", "SITE"] diff --git a/flash_rt/hardware/thor/attn_backend_chameleon.py b/flash_rt/hardware/thor/attn_backend_chameleon.py new file mode 100644 index 000000000..f14841583 --- /dev/null +++ b/flash_rt/hardware/thor/attn_backend_chameleon.py @@ -0,0 +1,362 @@ +"""FlashRT — Thor SM110 attention backend for standalone Chameleon-7B. + +One site: + +* **chameleon** — 32-layer Chameleon-7B LLM self-attention (NH=32, HD=128, + causal). Q/O aliased into the frontend's xn buffer. Dispatch order: + FA4 (FlashAttention-4, CuTe-DSL, optional fast path) → CUTLASS causal + FMHA (``libfmha_fp16_causal.so``) → cuBLAS decomposed causal MHA + (``attention_mha_causal_fp16``). + +Prefill (``run``) uses the top-left causal alignment (SQ == SK); single-query +incremental decode (``run_decode``) uses the bottom-right alignment via +``fmha_fp16_causal_br`` (or FA4 / plain cuBLAS MHA, both bottom-right +equivalent at SQ == 1). +""" + +from __future__ import annotations + +import ctypes +import logging +import pathlib + +from flash_rt.hardware.backend import AttentionBackendBase, AttentionSpec + +logger = logging.getLogger(__name__) + +# ── CUTLASS causal FMHA (SM100/110) dynamic loading ── +_fmha_causal_fn = None +_fmha_causal_br_fn = None # bottom-right aligned variant (decode, SQ < SK) + + +def _load_fmha_causal_library() -> bool: + """Load libfmha_fp16_causal.so and resolve the fmha_fp16_causal symbol.""" + global _fmha_causal_fn, _fmha_causal_br_fn + if _fmha_causal_fn is not None: + return True + search_paths = [ + pathlib.Path(__file__).parent.parent.parent / "libfmha_fp16_causal.so", + pathlib.Path(__file__).parent.parent.parent.parent / "build" / "libfmha_fp16_causal.so", + ] + argtypes = [ + ctypes.c_void_p, ctypes.c_void_p, ctypes.c_void_p, ctypes.c_void_p, + ctypes.c_int, ctypes.c_int, ctypes.c_int, + ctypes.c_int, ctypes.c_int, ctypes.c_int, + ctypes.c_void_p, + ] + for p in search_paths: + if p.exists(): + try: + lib = ctypes.CDLL(str(p)) + fn = lib.fmha_fp16_causal + fn.restype = ctypes.c_int + fn.argtypes = argtypes + _fmha_causal_fn = fn + # Bottom-right aligned causal (decode). Older .so builds may + # lack it; run_decode then falls through to the cuBLAS tier. + try: + fn_br = lib.fmha_fp16_causal_br + fn_br.restype = ctypes.c_int + fn_br.argtypes = argtypes + _fmha_causal_br_fn = fn_br + except AttributeError: + logger.warning( + "libfmha_fp16_causal.so lacks fmha_fp16_causal_br " + "(rebuild for CUTLASS decode path)") + logger.info("CUTLASS causal FMHA loaded from %s", p) + return True + except OSError as e: + logger.warning("Failed to load causal FMHA from %s: %s", p, e) + logger.warning("CUTLASS causal FMHA not found — will fall back to cuBLAS MHA") + return False + + +class ThorChameleonAttnBackend(AttentionBackendBase): + """Standalone Chameleon-7B attention backend on Thor (SM110). + + Single ``chameleon`` site: 32-layer causal MHA, Q/O aliased, per-layer + KV cache with ``layer_stride`` bytes between consecutive layers. + """ + + def __init__(self, spec: AttentionSpec, ctx, *, chameleon_slots: dict) -> None: + super().__init__(spec) + + expected_sites = {"chameleon"} + got = set(spec.sites.keys()) + if got != expected_sites: + raise ValueError( + f"ThorChameleonAttnBackend expects sites {expected_sites}, " + f"got {got}") + + self._ctx_cpp = ctx.cpp if hasattr(ctx, "cpp") else ctx + self._slots = {"chameleon": dict(chameleon_slots)} + self._require_keys("chameleon", + ("Q_O", "Kc", "Vc", "logits", "layer_stride", "scale")) + + self._per_layer_kv = {} + s = self._slots["chameleon"] + nL = spec.site("chameleon").num_layers + stride = int(s["layer_stride"]) + Kc = int(s["Kc"]) + Vc = int(s["Vc"]) + self._per_layer_kv["chameleon"] = [ + (Kc + l * stride, Vc + l * stride) for l in range(nL) + ] + + self._fvk = None + self._has_causal_fmha = _load_fmha_causal_library() + + # FA4 (FlashAttention-4, CuTe-DSL) fast path state. Populated by the + # frontend with torch-tensor references to the live Q (xn) and KV + # cache buffers, so run() can slice views (metadata-only, capture + # safe) and dispatch to FA4 instead of the CUTLASS FMHA kernel. + self._fa4_enabled = False + self._fa4_q_tensor = None + self._fa4_kv_cache = None + + def set_fa4_attn(self, q_tensor, kv_cache) -> None: + """Enable the FA4 (FlashAttention-4, CuTe-DSL) fast path. + + q_tensor must alias the chameleon Q_O slot buffer (the frontend's xn + buffer, which Q GEMMs write into and O is read from). kv_cache must + alias the per-layer K/V buffer with layers at [li, 0, :Se] (K) and + [li, 1, :Se] (V). Both are torch tensors so run() can build + capture-safe views. + """ + from flash_rt.hardware.thor import fa4_backend + if not fa4_backend.is_available(): + logger.warning("FA4 unavailable (%s); keeping CUTLASS FMHA", + fa4_backend.status()) + return + if q_tensor is None or kv_cache is None: + raise ValueError("FA4 requires the Q (xn) and KV cache tensor refs") + self._fa4_q_tensor = q_tensor + self._fa4_kv_cache = kv_cache + self._fa4_enabled = True + logger.info("FA4 causal FMHA enabled for chameleon site") + + def disable_fa4_attn(self) -> None: + """Revert chameleon site to the CUTLASS FMHA path.""" + self._fa4_enabled = False + + def _require_keys(self, site, keys): + slot = self._slots[site] + for k in keys: + if k not in slot: + raise ValueError(f"{site}_slots missing required key {k!r}") + if k in ("Q_O", "Kc", "Vc", "logits"): + if int(slot[k]) == 0: + raise ValueError( + f"{site}_slots[{k!r}] is a null device pointer") + + def _fvk_mod(self): + if self._fvk is None: + import flash_rt.flash_rt_kernels as fvk + self._fvk = fvk + return self._fvk + + def get_slot_ptrs(self, site, layer_idx): + """Return {Q, K, V, O} device pointer ints for (site, layer_idx).""" + if site not in self._slots: + raise KeyError(f"unknown site {site!r}") + nL = self._spec.site(site).num_layers + if not (0 <= layer_idx < nL): + raise IndexError( + f"layer_idx {layer_idx} out of range for site {site!r}") + K_ptr, V_ptr = self._per_layer_kv[site][layer_idx] + q_o = int(self._slots[site]["Q_O"]) + return {"Q": q_o, "K": K_ptr, "V": V_ptr, "O": q_o} + + def kv_row_ptrs(self, site, layer_idx, pos): + """Return (K_row, V_row) device pointer ints for token position pos. + + Rows are fp16 [NH*HD] within the layer's cache segment, so decode + can GEMM the new token's K/V directly into the cache. + """ + if site != "chameleon": + raise KeyError(f"unknown site {site!r}") + spec = self._spec.site(site) + if not (0 <= layer_idx < spec.num_layers): + raise IndexError( + f"layer_idx {layer_idx} out of range for site {site!r}") + if not (0 <= pos < spec.max_kv_seq): + raise IndexError(f"pos {pos} out of range for site {site!r}") + row_bytes = spec.num_kv_heads * spec.head_dim * 2 # fp16 + K_ptr, V_ptr = self._per_layer_kv[site][layer_idx] + return K_ptr + pos * row_bytes, V_ptr + pos * row_bytes + + def run(self, site, layer_idx, q_seq, *, kv_seq=None, stream=0, + state_nk=None, cross_attn=False): + """Dispatch causal attention for (site, layer_idx).""" + if site != "chameleon": + raise KeyError(f"unknown site {site!r}") + fvk = self._fvk_mod() + site_spec = self._spec.site(site) + nL = site_spec.num_layers + if not (0 <= layer_idx < nL): + raise IndexError( + f"layer_idx {layer_idx} out of range for site {site!r}") + if kv_seq is None: + kv_seq = q_seq + if q_seq == 1 and kv_seq > q_seq: + raise ValueError( + "run() got a decode shape (q_seq=1 < kv_seq); the top-left " + "causal mask misaligns here — call run_decode() instead") + + s = self._slots[site] + K_ptr, V_ptr = self._per_layer_kv[site][layer_idx] + NH = site_spec.num_q_heads + HD = site_spec.head_dim + + # ── FA4 (FlashAttention-4, CuTe-DSL) fast path ── + if self._fa4_enabled and self._fa4_q_tensor is not None: + try: + from flash_rt.hardware.thor import fa4_backend + fa4 = fa4_backend.fa4_func() + if fa4 is not None: + import torch + qv = self._fa4_q_tensor[:q_seq].view(1, q_seq, NH, HD) + kv_cache = self._fa4_kv_cache + k_view = kv_cache[layer_idx, 0, :kv_seq].view(1, kv_seq, NH, HD) + v_view = kv_cache[layer_idx, 1, :kv_seq].view(1, kv_seq, NH, HD) + with torch.no_grad(): + o = fa4(qv.contiguous(), k_view.contiguous(), + v_view.contiguous(), causal=True, pack_gqa=True) + if isinstance(o, tuple): + o = o[0] + o_flat = o.reshape(q_seq, NH * HD) + # Q_O slot aliases xn (self._fa4_q_tensor): overwrite + # Q with O in place via a metadata-only view. + self._fa4_q_tensor[:q_seq].view(q_seq, NH * HD).copy_(o_flat) + return int(s["Q_O"]) + except Exception as e: # pragma: no cover - defensive fallback + logger.warning("FA4 failed at layer %d (%s); falling back to CUTLASS", + layer_idx, e) + self._fa4_enabled = False + + # ── CUTLASS Causal FMHA (SM100/110) — preferred ── + if self._has_causal_fmha and _fmha_causal_fn is not None: + ret = _fmha_causal_fn( + ctypes.c_void_p(int(s["Q_O"])), + ctypes.c_void_p(K_ptr), + ctypes.c_void_p(V_ptr), + ctypes.c_void_p(int(s["Q_O"])), + 1, q_seq, kv_seq, NH, NH, HD, + ctypes.c_void_p(stream), + ) + if ret != 0: + logger.warning( + "CUTLASS causal FMHA returned %d, falling back to cuBLAS", + ret) + fvk.attention_mha_causal_fp16( + self._ctx_cpp, + int(s["Q_O"]), K_ptr, V_ptr, + int(s["logits"]), int(s["Q_O"]), + q_seq, kv_seq, NH, HD, + float(s["scale"]), stream, + ) + else: + # ── Fallback: cuBLAS decomposed causal MHA ── + fvk.attention_mha_causal_fp16( + self._ctx_cpp, + int(s["Q_O"]), K_ptr, V_ptr, + int(s["logits"]), int(s["Q_O"]), + q_seq, kv_seq, NH, HD, + float(s["scale"]), stream, + ) + return int(s["Q_O"]) + + def run_decode(self, site, layer_idx, kv_len, *, stream=0): + """Dispatch single-query (decode) attention for (site, layer_idx). + + Q is the single row in the Q_O slot; K/V are the first kv_len rows + of the layer's cache segment. The causal mask must be bottom-right + aligned (query at position kv_len-1 attends all kv_len keys). + """ + if site != "chameleon": + raise KeyError(f"unknown site {site!r}") + fvk = self._fvk_mod() + site_spec = self._spec.site(site) + nL = site_spec.num_layers + if not (0 <= layer_idx < nL): + raise IndexError( + f"layer_idx {layer_idx} out of range for site {site!r}") + if not (1 <= kv_len <= site_spec.max_kv_seq): + raise IndexError(f"kv_len {kv_len} out of range for site {site!r}") + + s = self._slots[site] + K_ptr, V_ptr = self._per_layer_kv[site][layer_idx] + NH = site_spec.num_q_heads + HD = site_spec.head_dim + + # ── FA4 fast path: causal is bottom-right (offset_k = sk - sq) ── + if self._fa4_enabled and self._fa4_q_tensor is not None: + try: + from flash_rt.hardware.thor import fa4_backend + fa4 = fa4_backend.fa4_func() + if fa4 is not None: + import torch + qv = self._fa4_q_tensor[:1].view(1, 1, NH, HD) + kv_cache = self._fa4_kv_cache + k_view = kv_cache[layer_idx, 0, :kv_len].view(1, kv_len, NH, HD) + v_view = kv_cache[layer_idx, 1, :kv_len].view(1, kv_len, NH, HD) + with torch.no_grad(): + o = fa4(qv.contiguous(), k_view.contiguous(), + v_view.contiguous(), causal=True, pack_gqa=True) + if isinstance(o, tuple): + o = o[0] + self._fa4_q_tensor[:1].view(1, NH * HD).copy_( + o.reshape(1, NH * HD)) + return int(s["Q_O"]) + except Exception as e: # pragma: no cover - defensive fallback + logger.warning("FA4 decode failed at layer %d (%s); falling back", + layer_idx, e) + self._fa4_enabled = False + + # ── CUTLASS bottom-right causal FMHA ── + if self._has_causal_fmha and _fmha_causal_br_fn is not None: + ret = _fmha_causal_br_fn( + ctypes.c_void_p(int(s["Q_O"])), + ctypes.c_void_p(K_ptr), + ctypes.c_void_p(V_ptr), + ctypes.c_void_p(int(s["Q_O"])), + 1, 1, kv_len, NH, NH, HD, + ctypes.c_void_p(stream), + ) + if ret != 0: + logger.warning( + "CUTLASS causal FMHA (br) returned %d at layer %d, " + "falling back to cuBLAS", ret, layer_idx) + fvk.attention_mha_fp16( + self._ctx_cpp, + int(s["Q_O"]), K_ptr, V_ptr, + int(s["logits"]), int(s["Q_O"]), + 1, kv_len, NH, HD, + float(s["scale"]), stream, + ) + else: + # ── cuBLAS non-causal MHA: with q_seq==1 the bottom-right + # causal mask is the identity, so this is equivalent. ── + fvk.attention_mha_fp16( + self._ctx_cpp, + int(s["Q_O"]), K_ptr, V_ptr, + int(s["logits"]), int(s["Q_O"]), + 1, kv_len, NH, HD, + float(s["scale"]), stream, + ) + return int(s["Q_O"]) + + +def make_chameleon_attention_spec(*, seq_max: int) -> AttentionSpec: + """Build the standalone Chameleon-7B Thor AttentionSpec (1 site).""" + spec = AttentionSpec() + spec.add_site( + "chameleon", + num_layers=32, num_q_heads=32, num_kv_heads=32, head_dim=128, + max_q_seq=int(seq_max), max_kv_seq=int(seq_max), + causal=True, + ) + return spec + + +__all__ = ["ThorChameleonAttnBackend", "make_chameleon_attention_spec"] diff --git a/flash_rt/hardware/thor/vqgan_trt_backend.py b/flash_rt/hardware/thor/vqgan_trt_backend.py new file mode 100644 index 000000000..344610c98 --- /dev/null +++ b/flash_rt/hardware/thor/vqgan_trt_backend.py @@ -0,0 +1,187 @@ +"""TensorRT VQ-GAN encoder backend for Jetson Thor. + +Manages multiple fixed-shape TRT engines (one per resolution). +Lazily loads/deserializes engines on first use of each resolution. +Falls back gracefully if engines are unavailable. +""" + +import json +import logging +from pathlib import Path +from typing import Optional + +import torch + +logger = logging.getLogger(__name__) + +_TRT_AVAILABLE = False +try: + import tensorrt as trt + _TRT_AVAILABLE = True +except ImportError: + pass + + +class VQGANTRTBackend: + ENGINE_DIR = Path.home() / ".flash_rt" / "trt_engines" / "vqgan" + + def __init__(self, engine_dir: Optional[Path] = None): + self._engine_dir = Path(engine_dir) if engine_dir else self.ENGINE_DIR + self._manifest = None + self._engines = {} + self._contexts = {} + self._buffers = {} + self._available = None + self._trt_logger = None + + def is_available(self) -> bool: + if self._available is not None: + return self._available + if not _TRT_AVAILABLE: + logger.warning("tensorrt not importable; TRT VQGAN disabled") + self._available = False + return False + manifest_path = self._engine_dir / "manifest.json" + if not manifest_path.exists(): + logger.warning("No TRT VQGAN manifest at %s", manifest_path) + self._available = False + return False + with open(manifest_path) as f: + self._manifest = json.load(f) + has_engines = len(self._manifest.get("engines", {})) > 0 + if has_engines: + logger.info("TRT VQGAN backend: %d engines at %s (TRT %s)", + len(self._manifest["engines"]), self._engine_dir, + self._manifest.get("trt_version", "?")) + self._available = has_engines + return has_engines + + def supports_resolution(self, height: int, width: int) -> bool: + if not self._available: + return False + key = f"{height}x{width}" + return key in self._manifest.get("engines", {}) + + def encode(self, image_tensor: torch.Tensor) -> Optional[torch.Tensor]: + """Run TRT VQGAN encoder. + + Args: + image_tensor: [1, 3, H, W] float32 on CUDA, range [-1, 1] + + Returns: + [1, H//16, W//16] int64 codebook indices on CUDA, or None on failure. + """ + _, _, h, w = image_tensor.shape + key = f"{h}x{w}" + + if key not in self._contexts: + if not self._load_engine(h, w): + return None + + ctx = self._contexts[key] + inp_buf, out_buf, out_dtype_is_int32 = self._buffers[key] + + inp_buf.copy_(image_tensor) + ctx.set_tensor_address("image", inp_buf.data_ptr()) + ctx.set_tensor_address("indices", out_buf.data_ptr()) + + stream = torch.cuda.current_stream() + ok = ctx.execute_async_v3(stream_handle=stream.cuda_stream) + if not ok: + logger.error("TRT execute_async_v3 failed for resolution %s", key) + return None + + if out_dtype_is_int32: + return out_buf.to(torch.int64) + return out_buf.clone() + + def encode_batch(self, image_tensor: torch.Tensor) -> Optional[torch.Tensor]: + """Run TRT VQGAN encoder on a multi-view batch. + + Args: + image_tensor: [B, 3, H, W] float32 on CUDA, range [-1, 1] + where B == engine batch (from per-engine manifest entry). + + Returns: + [B, H//16, W//16] int64 codebook indices, or None on failure / + batch mismatch. + """ + b, _, h, w = image_tensor.shape + key = f"{h}x{w}" + + if key not in self._contexts: + if not self._load_engine(h, w): + return None + + engine_batch = self._manifest["engines"][key]["input_shape"][0] + if b != engine_batch: + return None # caller should fall back to per-view encode + + ctx = self._contexts[key] + inp_buf, out_buf, out_dtype_is_int32 = self._buffers[key] + + inp_buf.copy_(image_tensor) + ctx.set_tensor_address("image", inp_buf.data_ptr()) + ctx.set_tensor_address("indices", out_buf.data_ptr()) + + stream = torch.cuda.current_stream() + ok = ctx.execute_async_v3(stream_handle=stream.cuda_stream) + if not ok: + logger.error("TRT execute_async_v3 (batch=%d) failed at %s", b, key) + return None + + if out_dtype_is_int32: + return out_buf.to(torch.int64) + return out_buf.clone() + + def _load_engine(self, height: int, width: int) -> bool: + key = f"{height}x{width}" + meta = self._manifest["engines"].get(key) + if meta is None: + return False + + engine_path = self._engine_dir / meta["file"] + if not engine_path.exists(): + logger.error("Engine file missing: %s", engine_path) + return False + + if self._trt_logger is None: + self._trt_logger = trt.Logger(trt.Logger.WARNING) + if hasattr(trt, "init_libnvinfer_plugins"): + trt.init_libnvinfer_plugins(self._trt_logger, "") + + runtime = trt.Runtime(self._trt_logger) + with open(engine_path, "rb") as f: + engine = runtime.deserialize_cuda_engine(f.read()) + + if engine is None: + logger.error("Failed to deserialize TRT engine: %s", engine_path) + return False + + context = engine.create_execution_context() + self._engines[key] = engine + self._contexts[key] = context + + self._allocate_buffers(key, height, width) + logger.info("TRT VQGAN engine loaded: %s (%s)", key, engine_path.name) + return True + + def _allocate_buffers(self, key: str, height: int, width: int): + # Per-engine batch (each res may have been built independently with + # a different batch); the top-level manifest["batch"] is only the + # last-built entry and unreliable when engines coexist. + batch = self._manifest["engines"][key]["input_shape"][0] + h_lat, w_lat = height // 16, width // 16 + device = torch.device("cuda") + + inp_buf = torch.empty(batch, 3, height, width, device=device, dtype=torch.float32) + + engine = self._engines[key] + out_dtype_trt = engine.get_tensor_dtype("indices") + out_dtype_is_int32 = (out_dtype_trt == trt.DataType.INT32) + if out_dtype_is_int32: + out_buf = torch.empty(batch, h_lat, w_lat, device=device, dtype=torch.int32) + else: + out_buf = torch.empty(batch, h_lat, w_lat, device=device, dtype=torch.int64) + + self._buffers[key] = (inp_buf, out_buf, out_dtype_is_int32) diff --git a/flash_rt/models/chameleon/__init__.py b/flash_rt/models/chameleon/__init__.py new file mode 100644 index 000000000..82f4de55b --- /dev/null +++ b/flash_rt/models/chameleon/__init__.py @@ -0,0 +1,18 @@ +"""FlashRT Chameleon-7B (VLM/LLM) model namespace. + +Chameleon-7B is a 32-layer early-fusion multimodal LLM (MHA 32x128 with +per-head QK LayerNorm + RoPE, attention_bias=false, mlp_bias=false, +SwiGLU FFN, 8192-token VQ-GAN image vocabulary). + +Two hardware paths: + +- Thor SM110 (``pipeline_thor.py``): dynamic per-tensor FP8 backbone with + fused norm/activation+quantize kernels, optional NVFP4 FFN layers, + causal CUTLASS SM100 FMHA (``libfmha_fp16_causal.so`` / + ``libfmha_fp8_causal.so``) and CUDA-graph decode. +- Orin SM87 (``pipeline_rtx.py``): INT8 W8A8 + INT4 W4A4 QuaRot-Hadamard + rotated weights, SM80 CUTLASS rowwise GEMMs, FA2 fp16 causal attention. + +Image tokenization uses the Apache-2.0 Transformers ``ChameleonVQVAE`` +implementation and loads only ``model.vqmodel.*`` checkpoint tensors. +""" diff --git a/flash_rt/models/chameleon/pipeline_rtx.py b/flash_rt/models/chameleon/pipeline_rtx.py new file mode 100644 index 000000000..6e11cf392 --- /dev/null +++ b/flash_rt/models/chameleon/pipeline_rtx.py @@ -0,0 +1,305 @@ +"""FlashRT — Chameleon-7B VLM forward for Jetson AGX Orin (SM87). + +Standalone text-generating forward on the Chameleon-7B INT8/QuaRot-INT4 +kernel set (SM80 CUTLASS rowwise GEMMs + Hadamard rotations). Key design +points: + +1. **No attention bias.** Upstream Chameleon has ``attention_bias = false`` and + ``mlp_bias = false``, so the no-bias GEMM entries are used unconditionally. +2. **A real KV cache.** The K and V GEMMs write straight into the attention + backend's per-layer slab — legal because CUTLASS hard-wires its output row + stride to ``N``, which equals the cache's row stride. No staging, no copy. +3. **One code path for prefill and decode.** ``pos is None`` means prefill + (``S`` rows at the slab base, RoPE from position 0); ``pos`` set means a + single decode row written at ``pos``, with the cos/sin pointers advanced by + ``pos`` rows. ``qk_norm_rope_fused_fp16`` needs no modification for this: + it derives position as ``row / num_heads``, which is 0 at ``S = 1``, so the + position lives entirely in the table pointer. +4. **lm_head tail instead of an action head.** Final RMSNorm, then per-row INT8 + quantization of the wanted row(s), then an INT8 GEMM to BF16 logits. The + ``mask_image_logits`` step and the argmax live in the frontend (they are + torch ops on the logits view, which are graph-safe — verified). + +lm_head stays INT8 in both precision tiers: it is 268 MB/token = 1.74 ms = +3.7 % of the INT8 decode budget, and dropping to 15 INT4 levels over a +65536-row output is not worth 0.8 ms. + +Raw-pointer interface only (int pointers + Python primitives) for CUDA-Graph +safety: no torch ops, no allocation, no sync inside the forward. +""" + +from __future__ import annotations + +import ctypes + +_CUDART = None + + +def _gpu_copy(dst_ptr: int, src_ptr: int, nbytes: int, stream: int) -> None: + """Async D2D copy — the in-graph layer-probe mechanism.""" + global _CUDART + if _CUDART is None: + _CUDART = ctypes.CDLL("libcudart.so") + _CUDART.cudaMemcpyAsync( + ctypes.c_void_p(dst_ptr), ctypes.c_void_p(src_ptr), + ctypes.c_size_t(nbytes), 3, ctypes.c_void_p(stream)) + + +def _check(status, name: str, shape) -> None: + if status != 0: + raise RuntimeError(f"{name} failed: status={status} shape={shape}") + + +def chameleon_forward( + fvk, bufs, weights, dims, scales_dev, + *, attn, S: int, pos=None, stream: int = 0, + use_int4: bool = False, use_int4_down: bool = False, + use_hadamard: bool = False, + ffn_down_clamp_value: float = 60000.0, + ffn_down_clamp_last_n: int = 4, + logits_all: bool = False, probe=None, +) -> None: + """Run the 32-layer Chameleon-7B decoder and the lm_head. + + Args: + fvk: ``flash_rt.flash_rt_kernels`` module. + bufs / weights / dims / scales_dev: int-pointer dicts from the frontend. + attn: ``ChameleonAttnBackend`` (owns the KV cache). + S: rows to process. ``1`` on the decode path. + pos: ``None`` for prefill; otherwise the absolute KV position of the + single decode row. RoPE and the KV write both key off this. + ffn_down_clamp_value: symmetric clamp applied to the down-projection + output before the residual add; ``<= 0`` disables it. + **This is required for correctness, not a tuning knob.** Measured on + this checkpoint (ISL=1032), the reference's magnitudes are tiny + through L30 and then explode at L31 only: + + layer L28 L29 L30 L31 + residual 1616 1720 2032 266240 + down_out 1120 1032 1880 264192 + + 2.6e5 is far beyond FP16's 65504, so the store becomes ``inf`` and + the final RMSNorm then poisons that row's logits. Because the + pre-L31 residual is only ~2032, clamping the down output at 60000 + keeps the residual at ~62000 < 65504 — so no BF16 residual stream is + needed. See ``docs/chameleon7b_rtx_sm87.md``. Unlike the + Thor path we do **not** clamp the down *input*: ours is BF16 + (``cutlass_int8_silu_gated_bf16out``), whose range absorbs the + 151552 without issue. + logits_all: compute logits for **all** S rows (teacher-forced + precision comparison). Requires an ``[S, vocab]`` logits buffer. + probe: ``{"layers": [...], "bufs": [...], "final_buf": ptr}`` to snapshot + the post-residual hidden state; ``None`` disables it at zero cost. + """ + decode = pos is not None + if decode and S != 1: + raise ValueError(f"decode path requires S == 1, got {S}") + + D = int(dims["D"]) + Dff = int(dims["Dff"]) + L = int(dims["L"]) + H = int(dims["H"]) + Hd = int(dims["Hd"]) + V = int(dims["vocab"]) + + x_ptr = int(bufs["x"]) + xn_ptr = int(bufs["xn"]) + int8_act_d_ptr = int(bufs["int8_act_d"]) + int8_act_ff_ptr = int(bufs["int8_act_ff"]) + int4_act_d_ptr = int(bufs.get("int4_act_d", 0)) + int4_act_ff_ptr = int(bufs.get("int4_act_ff", 0)) + bf16_gate_ptr = int(bufs["bf16_gate_ff"]) + bf16_xn_ff_ptr = int(bufs["bf16_xn_ff"]) + o_proj_out_ptr = int(bufs["o_proj_out"]) + logits_ptr = int(bufs["logits"]) + lm_act_ptr = int(bufs["lm_act"]) + lm_scale_ptr = int(bufs["lm_act_scale"]) + + # RoPE position enters purely through the table pointer (see module doc). + rope_off = (pos if decode else 0) * Hd * 2 + cos_ptr = int(weights["rope_cos"]) + rope_off + sin_ptr = int(weights["rope_sin"]) + rope_off + + probe_map = None + probe_final = None + if probe is not None: + layers = probe.get("layers") or [] + pbufs = probe.get("bufs") or [] + if len(layers) != len(pbufs): + raise ValueError("probe['layers'] and probe['bufs'] length mismatch") + probe_map = {int(li): int(p) for li, p in zip(layers, pbufs)} + probe_final = int(probe.get("final_buf") or 0) + + # ── entry: fused RMSNorm + quantize (layer 0's input_layernorm) ── + if use_int4: + fvk.rms_norm_fht_int4_fp16( + x_ptr, int(weights["input_ln_w"][0]), + int4_act_d_ptr, int(scales_dev["act_qkv"][0]), + S, D, 1e-5, int(stream)) + elif use_hadamard: + fvk.rms_norm_fht_int8_fp16( + x_ptr, int(weights["input_ln_w"][0]), + int8_act_d_ptr, int(scales_dev["act_qkv"][0]), + S, D, 1e-5, int(stream)) + else: + fvk.rms_norm_int8_rowwise_fp16( + x_ptr, int(weights["input_ln_w"][0]), + int8_act_d_ptr, int(scales_dev["act_qkv"][0]), + S, D, 1e-5, int(stream)) + + act_d_ptr = int4_act_d_ptr if use_int4 else int8_act_d_ptr + + for li in range(L): + if decode: + K_ptr, V_ptr = attn.kv_row_ptrs(li, pos) + Q_ptr = int(attn.get_slot_ptrs("llm", li)["Q"]) + else: + slots = attn.get_slot_ptrs("llm", li) + Q_ptr, K_ptr, V_ptr = int(slots["Q"]), int(slots["K"]), int(slots["V"]) + + a_qkv = int(scales_dev["act_qkv"][li]) + a_o = int(scales_dev["act_o"][li]) + a_gu = int(scales_dev["act_gu"][li]) + a_d = int(scales_dev["act_down"][li]) + + # ── Q/K/V: K and V land directly in the KV cache ── + qkv_gemm = (fvk.cutlass_int4_rowwise_fp16out if use_int4 + else fvk.cutlass_int8_rowwise_fp16out) + for name, out_ptr in (("q_w", Q_ptr), ("k_w", K_ptr), ("v_w", V_ptr)): + _check(qkv_gemm(act_d_ptr, int(weights[name][li]), a_qkv, + int(weights[name + "_scale"][li]), out_ptr, + S, D, D, int(stream)), + f"{'int4' if use_int4 else 'int8'} {name}", (S, D, D)) + + # ── fused per-head QK LayerNorm(+bias) + rotate-half RoPE, in place ── + fvk.qk_norm_rope_fused_fp16( + Q_ptr, K_ptr, + int(weights["q_norm_w"][li]), int(weights["q_norm_b"][li]), + int(weights["k_norm_w"][li]), int(weights["k_norm_b"][li]), + cos_ptr, sin_ptr, + S, H, Hd, 1e-5, int(stream)) + + # ── causal MHA (result written back into the Q slot) ── + if decode: + attn.run_decode(li, pos, stream=int(stream)) + else: + attn.run_prefill(li, S, stream=int(stream)) + + # ── O projection ── + if use_int4: + fvk.fht_int4_quant_fp16(Q_ptr, int4_act_d_ptr, a_o, S, D, int(stream)) + elif use_hadamard: + fvk.fht_int8_quant_fp16(Q_ptr, int8_act_d_ptr, a_o, S, D, int(stream)) + else: + fvk.quantize_int8_rowwise_fp16(Q_ptr, int8_act_d_ptr, a_o, + S, D, int(stream)) + o_gemm = (fvk.cutlass_int4_rowwise_fp16out if use_int4 + else fvk.cutlass_int8_rowwise_fp16out) + _check(o_gemm(act_d_ptr, int(weights["o_w"][li]), a_o, + int(weights["o_w_scale"][li]), o_proj_out_ptr, + S, D, D, int(stream)), "o_proj", (S, D, D)) + + # ── residual_1 + post-attention RMSNorm + quantize ── + if use_int4: + fvk.residual_add_rms_norm_fht_int4_fp16( + x_ptr, o_proj_out_ptr, int(weights["post_ln_w"][li]), + int4_act_d_ptr, a_gu, S, D, 1e-5, int(stream)) + elif use_hadamard: + fvk.residual_add_rms_norm_fht_int8_fp16( + x_ptr, o_proj_out_ptr, int(weights["post_ln_w"][li]), + int8_act_d_ptr, a_gu, S, D, 1e-5, int(stream)) + else: + fvk.residual_add_rms_norm_int8_rowwise_fp16( + x_ptr, o_proj_out_ptr, int(weights["post_ln_w"][li]), + int8_act_d_ptr, a_gu, S, D, 1e-5, int(stream)) + + # ── FFN: gate -> BF16, up with fused SiLU(gate)* -> BF16 ── + if use_int4: + _check(fvk.cutlass_int4_rowwise_bf16out( + int4_act_d_ptr, int(weights["gate_w"][li]), a_gu, + int(weights["gate_w_scale"][li]), bf16_gate_ptr, + S, Dff, D, int(stream)), "int4 gate", (S, Dff, D)) + _check(fvk.cutlass_int4_silu_gated_bf16out( + int4_act_d_ptr, int(weights["up_w"][li]), a_gu, + int(weights["up_w_scale"][li]), bf16_gate_ptr, bf16_xn_ff_ptr, + S, Dff, D, int(stream)), "int4 up+silu", (S, Dff, D)) + else: + _check(fvk.cutlass_int8_rowwise_bf16out( + int8_act_d_ptr, int(weights["gate_w"][li]), a_gu, + int(weights["gate_w_scale"][li]), bf16_gate_ptr, + S, Dff, D, int(stream)), "int8 gate", (S, Dff, D)) + _check(fvk.cutlass_int8_silu_gated_bf16out( + int8_act_d_ptr, int(weights["up_w"][li]), a_gu, + int(weights["up_w_scale"][li]), bf16_gate_ptr, bf16_xn_ff_ptr, + S, Dff, D, int(stream)), "int8 up+silu", (S, Dff, D)) + + # ── down projection ── + if use_int4_down: + fvk.fht128_int4_quant_bf16(bf16_xn_ff_ptr, int4_act_ff_ptr, a_d, + S, Dff, int(stream)) + _check(fvk.cutlass_int4_rowwise_fp16out( + int4_act_ff_ptr, int(weights["d_w"][li]), a_d, + int(weights["d_w_scale"][li]), o_proj_out_ptr, + S, D, Dff, int(stream)), "int4 down", (S, D, Dff)) + else: + fvk.quantize_int8_rowwise(bf16_xn_ff_ptr, int8_act_ff_ptr, a_d, + S, Dff, int(stream)) + _check(fvk.cutlass_int8_rowwise_fp16out( + int8_act_ff_ptr, int(weights["d_w"][li]), a_d, + int(weights["d_w_scale"][li]), o_proj_out_ptr, + S, D, Dff, int(stream)), "int8 down", (S, D, Dff)) + + # Guard the FP16 residual against L31's massive down output (see the + # measured table in the docstring). Restricted to the last + # ``ffn_down_clamp_last_n`` layers: the magnitude grows monotonically + # with depth and L28 measures 1616, i.e. 37x below the clamp, so the + # earlier layers cannot reach it. Clamping all 32 layers instead costs + # 6.2 ms of a 281 ms prefill (2.2 %) for no effect. The Gate-1 harness + # reports per-layer clamp saturation, so a checkpoint that violates the + # assumption is detectable — raise this to L if that ever happens. + if ffn_down_clamp_value > 0.0 and li >= L - ffn_down_clamp_last_n: + fvk.clamp_inplace_fp16(o_proj_out_ptr, float(ffn_down_clamp_value), + S * D, int(stream)) + + # ── residual_2 (+ next layer's input_layernorm + quantize) ── + if li < L - 1: + if use_int4: + fvk.residual_add_rms_norm_fht_int4_fp16( + x_ptr, o_proj_out_ptr, int(weights["input_ln_w"][li + 1]), + int4_act_d_ptr, int(scales_dev["act_qkv"][li + 1]), + S, D, 1e-5, int(stream)) + elif use_hadamard: + fvk.residual_add_rms_norm_fht_int8_fp16( + x_ptr, o_proj_out_ptr, int(weights["input_ln_w"][li + 1]), + int8_act_d_ptr, int(scales_dev["act_qkv"][li + 1]), + S, D, 1e-5, int(stream)) + else: + fvk.residual_add_rms_norm_int8_rowwise_fp16( + x_ptr, o_proj_out_ptr, int(weights["input_ln_w"][li + 1]), + int8_act_d_ptr, int(scales_dev["act_qkv"][li + 1]), + S, D, 1e-5, int(stream)) + else: + fvk.residual_add_fp16(x_ptr, o_proj_out_ptr, S * D, int(stream)) + + if probe_map is not None and li in probe_map: + _gpu_copy(probe_map[li], x_ptr, S * D * 2, stream) + + # ── final RMSNorm ── + fvk.rms_norm_fp16(x_ptr, int(weights["final_norm_w"]), xn_ptr, + S, D, 1e-5, int(stream)) + if probe_final: + _gpu_copy(probe_final, xn_ptr, S * D * 2, stream) + + # ── lm_head: INT8 W8A8 -> BF16 logits ── + # Next-token prediction reads the last row; logits_all is for the + # teacher-forced precision comparison. + rows, row_off = (S, 0) if logits_all else (1, S - 1) + fvk.quantize_int8_rowwise_fp16(xn_ptr + row_off * D * 2, lm_act_ptr, + lm_scale_ptr, rows, D, int(stream)) + _check(fvk.cutlass_int8_rowwise_bf16out( + lm_act_ptr, int(weights["lm_head_w"]), lm_scale_ptr, + int(weights["lm_head_w_scale"]), logits_ptr, + rows, V, D, int(stream)), "lm_head", (rows, V, D)) + + +__all__ = ["chameleon_forward"] diff --git a/flash_rt/models/chameleon/pipeline_thor.py b/flash_rt/models/chameleon/pipeline_thor.py new file mode 100644 index 000000000..c76d0d387 --- /dev/null +++ b/flash_rt/models/chameleon/pipeline_thor.py @@ -0,0 +1,1126 @@ +"""FlashRT — standalone Chameleon-7B Thor SM110 pipeline forward functions. + +Chameleon-7B LLM (32-layer MHA, attention_bias=false, mlp_bias=false, +per-head QK LayerNorm + RoPE). Used by the standalone Chameleon Thor +frontend. + +All functions use raw-pointer interface (int pointers + Python primitives) +for CUDA Graph compatibility. No dynamic allocation, no torch ops, no sync. + +Functions: + chameleon_forward — Chameleon-7B LLM inference (dynamic FP8 default) + chameleon_forward_fp16 — pure-FP16 reference path + chameleon_forward_calibrate — FP8 static-scale calibration +""" + +from __future__ import annotations + +import math + +import flash_rt.flash_rt_kernels as _fvk +import torch + +from flash_rt.hardware.thor.shared_primitives import ( + _measure_scale_gpu, + _gpu_copy, + _gpu_sync, + _gpu_zero, +) + +try: + import flash_rt.flash_rt_fp4 as _fvk_fp4 +except Exception: + _fvk_fp4 = None + + +def _parse_fp4_layer_policy() -> frozenset: + """FP4 FFN layer policy from FLASHRT_CHAMELEON_FP4_LAYERS env var. + + Values: + - unset / "": default = L0-L2 FP4, L3-L31 FP8 (safe default). + - "0-7": FP4 for L0..L7 inclusive. + - "0-14,20-31": FP4 for L0..L14 + L20..L31 (skip outlier L15-L19). + This is the SM120-sweep-validated aggressive setting (13ms savings on + RTX 5090; expect similar Thor gains). + - comma-separated list: "0,1,2,5" → FP4 on those specific layers. + + Returns the frozenset of FP8 layer indices (complement of FP4 set). + """ + import os as _os + val = _os.environ.get("FLASHRT_CHAMELEON_FP4_LAYERS", "").strip() + if not val: + return frozenset(range(3, 32)) # default: L0-L2 FP4 + + fp4_layers = set() + for chunk in val.split(","): + chunk = chunk.strip() + if "-" in chunk: + a, b = chunk.split("-") + fp4_layers.update(range(int(a), int(b) + 1)) + else: + fp4_layers.add(int(chunk)) + fp4_layers = {li for li in fp4_layers if 0 <= li < 32} + return frozenset(li for li in range(32) if li not in fp4_layers) + + +# ── FP4 GEMM variant and layer policy ── +FP4_VARIANT = 9 +_FFN_FP8_LAYERS = _parse_fp4_layer_policy() + + +# ══════════════════════════════════════════════════════════════════ +# Chameleon-7B LLM (32 layers, mixed FP4/FP8 GEMMs) +# ══════════════════════════════════════════════════════════════════ + +def chameleon_forward( + gemm, fvk, bufs, weights, dims, scales_dev, + *, attn, stream: int = 0, + alpha_host=None, awq_v_proj=None, + ffn_down_clamp_value: float = 10000.0, + ffn_clamp_layers=None, + dynamic_fp8_layers: frozenset = frozenset(), + fp4_ffn_layers: frozenset = frozenset(), + probe=None, +) -> None: + """Chameleon-7B LLM forward pass (32 layers). + + Production precision: dynamic per-tensor FP8 on all layers, with the FFN of + ``fp4_ffn_layers`` optionally run in NVFP4 W4A16 (decoupled from attention). + + Output: hidden_all = RMSNorm(x) written to bufs['hidden_all'] as + full [Se, D] FP16 tensor (consumed by action_head_forward). + + ``ffn_down_clamp_value``: Clamp down_out (o_proj_out) to ±V after + the FFN before residual_2. Chameleon-7B L31 down_proj's FP32 + accumulator × alpha can push output beyond ±65504 producing inf. + Applied after both FP4 and FP8 FFN paths. Set to <= 0 to disable. + + ``ffn_clamp_layers``: Optional set/frozenset of layer indices where the + FFN clamps are applied. ``None`` preserves the historical behavior and + clamps every layer; callers can pass e.g. ``frozenset({31})`` to clamp + only the deep outlier layer and avoid redundant elementwise passes. + + ``dynamic_fp8_layers``: frozenset of layer indices to run with full FP8 + GEMMs but RUNTIME (per-forward) per-tensor activation scaling instead of + the static calibrated scale. Uses quantize_fp8_device_fp16 (GPU amax) + + fp8_nn_dev_fp16 (device-scale GEMM). Restores 512-res precision (static + scale is wrong for long sequences) at full FP8 speed. CUDA-Graph safe. + This is the default for all 32 layers; see docs/chameleon_thor_sm110.md. + + ``fp4_ffn_layers``: frozenset of layer indices whose FFN runs in NVFP4 + W4A16 while attention stays on the dynamic-FP8 path (decoupled). Requires + those layers' FP4 weights to be packed. 512-res default = L0-7 (~1.10x E2E, + cos >= 0.99 on 3/4 variants); empty at 256-res (FP4 breaks precision). See + docs/chameleon_thor_sm110.md (§6 FP4). + + Optional ``probe`` dict for layer-wise precision debugging:: + + { + 'layers': tuple[int, ...], # layer indices to snapshot post-residual-2 + 'bufs': tuple[int, ...], # device pointers, one per layer index, + # each must hold >= Se*D fp16 elements + 'final_buf': int, # device pointer for post-final-RMSNorm + # snapshot. Pass 0 to skip. + } + """ + Se = int(dims['Se']) + D = int(dims['D']) + Dff = int(dims['Dff']) + L = int(dims['L']) + H = int(dims['H']) + Hd = int(dims['Hd']) + + if alpha_host is None: + alpha_host = weights.get('alpha_host') + if awq_v_proj is None: + awq_v_proj = weights.get('awq_v_proj') + + x_ptr = int(bufs['x']) + xn_ptr = int(bufs['xn']) + xn_fp8_ptr = int(bufs['xn_fp8']) + o_proj_out_ptr = int(bufs['o_proj_out']) + hidden_all_ptr = int(bufs['hidden_all']) + + # FP4 FFN buffers + act_fp4_ptr = int(bufs['act_fp4']) + act_sfa_ptr = int(bufs['act_sfa']) + ffn_act_fp4_ptr = int(bufs['ffn_act_fp4']) + ffn_act_sfa_ptr = int(bufs['ffn_act_sfa']) + gu_merged_ptr = int(bufs['gu_merged']) + + cos_ptr = int(weights['rope_cos']) + sin_ptr = int(weights['rope_sin']) + + # Layer-0 pre-attention RMSNorm + FP8 quantize + fvk.rms_norm_fp16( + x_ptr, int(weights['input_ln_w'][0]), xn_ptr, + Se, D, 1e-5, int(stream), + ) + fvk.quantize_fp8_static_fp16( + xn_ptr, xn_fp8_ptr, int(scales_dev["act_qkv"][0]), Se * D, int(stream), + ) + + for li in range(L): + clamp_this_layer = ( + ffn_down_clamp_value > 0.0 + and (ffn_clamp_layers is None or li in ffn_clamp_layers) + ) + slots = attn.get_slot_ptrs("chameleon", li) + Q_ptr = int(slots["Q"]) + K_ptr = int(slots["K"]) + V_ptr = int(slots["V"]) + O_ptr = int(slots["O"]) + + d_act_qkv = int(scales_dev["act_qkv"][li]) + d_act_o = int(scales_dev["act_o"][li]) + d_act_gu = int(scales_dev["act_gu"][li]) + d_act_d = int(scales_dev["act_down"][li]) + + d_w_qkv = int(weights['d_w_qkv'][li]) + d_w_o = int(weights['d_w_o'][li]) + d_w_gu = int(weights['d_w_gu'][li]) + d_w_d = int(weights['d_w_d'][li]) + + q_w_ptr = int(weights['q_w'][li]) + k_w_ptr = int(weights['k_w'][li]) + v_w_ptr = int(weights['v_w'][li]) + + + # ═══ Dynamic per-tensor FP8 branch ═══ + # Runtime amax scaling (quantize_fp8_device_fp16 -> fp8_nn_dev_fp16) + # instead of static-calibrated scale. The static per-tensor scale is + # wrong for long-sequence (512-res) variants where deep-layer + # activations shift; recomputing amax each forward restores precision + # at full FP8 speed (no FP16 GEMM). CUDA-Graph safe: the device scale + # pointer is dereferenced at replay, so it adapts per input. + if li in dynamic_fp8_layers: + # Self-sufficient entry: re-derive FP16 xn from residual stream. + # Fused RMSNorm + dynamic FP8 quantize: amax is measured inside + # the norm's own output-write pass, skipping the separate + # absmax_kernel read of xn (one fewer full pass over Se*D). + dyn_qkv = int(scales_dev['dyn_act_qkv'][li]) + dyn_o = int(scales_dev['dyn_act_o'][li]) + dyn_gu = int(scales_dev['dyn_act_gu'][li]) + dyn_d = int(scales_dev['dyn_act_down'][li]) + + fvk.rms_norm_quantize_dynamic_fp8_fp16( + x_ptr, int(weights['input_ln_w'][li]), xn_ptr, xn_fp8_ptr, + dyn_qkv, Se, D, 1e-5, int(stream)) + + # QKV: device-scale GEMM (xn_fp8 already produced above). + # All three GEMMs read xn_fp8; Q writes into xn (aliased Q_O) + # AFTER the quantize, so the clobber is harmless. + gemm.fp8_nn_dev_fp16(xn_fp8_ptr, v_w_ptr, V_ptr, + Se, D, D, dyn_qkv, d_w_qkv, int(stream)) + gemm.fp8_nn_dev_fp16(xn_fp8_ptr, k_w_ptr, K_ptr, + Se, D, D, dyn_qkv, d_w_qkv, int(stream)) + gemm.fp8_nn_dev_fp16(xn_fp8_ptr, q_w_ptr, Q_ptr, + Se, D, D, dyn_qkv, d_w_qkv, int(stream)) + + fvk.qk_norm_rope_fused_fp16( + Q_ptr, K_ptr, + int(weights['q_norm_w'][li]), int(weights['q_norm_b'][li]), + int(weights['k_norm_w'][li]), int(weights['k_norm_b'][li]), + cos_ptr, sin_ptr, + Se, H, Hd, 1e-5, int(stream)) + attn.run("chameleon", li, q_seq=Se, kv_seq=Se, stream=int(stream)) + + # O projection: dynamic. + fvk.quantize_fp8_device_fp16( + O_ptr, xn_fp8_ptr, dyn_o, Se * D, int(stream)) + gemm.fp8_nn_dev_fp16(xn_fp8_ptr, int(weights['o_w'][li]), + o_proj_out_ptr, Se, D, D, dyn_o, d_w_o, + int(stream)) + + # Residual 1. + fvk.residual_add_fp16(x_ptr, o_proj_out_ptr, Se * D, int(stream)) + + # ── FFN: decoupled NVFP4 (opt-in) or dynamic FP8 ── + # DECOUPLED mode keeps attention on dynamic FP8 (above) but runs the + # FFN in NVFP4 W4A16 for FP4-eligible layers — isolating FFN + # precision from attention precision. Falls back to dynamic FP8 FFN. + _use_fp4_ffn = ( + _fvk_fp4 is not None + and li in fp4_ffn_layers + and int(weights['gu_w_fp4'][li]) != 0 + ) + if _use_fp4_ffn: + # Post-attn RMSNorm -> FP16 xn (no FP8 quantize needed here; + # the FP4 path quantizes xn_ptr directly below). + fvk.rms_norm_fp16(x_ptr, int(weights['post_ln_w'][li]), xn_ptr, + Se, D, 1e-5, int(stream)) + # NOTE: intended for SHALLOW layers only. Unlike the FP8 branch + # below, there is no intermediate clamp on the SwiGLU output + # (gate_geglu fuses silu*mul + FP4 quantize). A deep layer whose + # gu exceeds fp16 max (~65504, e.g. L31) could overflow to + # inf/nan *before* the FP4 quantize; the final o_proj_out clamp + # cannot recover that. The shipped 512 default (L0-7) is safe. + # gate+up merged NVFP4 GEMM (dynamic per-block SFA). + _fvk_fp4.quantize_fp4_dynamic_sfa_fp16( + xn_ptr, act_fp4_ptr, act_sfa_ptr, Se, D, False, int(stream)) + _fvk_fp4.cutlass_fp4_gemm_variant( + FP4_VARIANT, act_fp4_ptr, act_sfa_ptr, + int(weights['gu_w_fp4'][li]), int(weights['gu_sfb'][li]), + gu_merged_ptr, Se, 2 * Dff, D, 1.0, 0.0, int(stream)) + # fused SwiGLU + quantize down-proj input to NVFP4. + _fvk_fp4.gate_geglu_fp4_sfa_v2_fp16( + gu_merged_ptr, ffn_act_fp4_ptr, ffn_act_sfa_ptr, + Se, Dff, int(stream)) + _fvk_fp4.cutlass_fp4_gemm_variant( + FP4_VARIANT, ffn_act_fp4_ptr, ffn_act_sfa_ptr, + int(weights['d_w_fp4'][li]), int(weights['d_sfb'][li]), + o_proj_out_ptr, Se, D, Dff, 1.0, 0.0, int(stream)) + else: + # Post-attn residual add + RMSNorm + dynamic FP8 quantize, + # fused into one elementwise kernel: residual is fp16-rounded + # (same as residual_add_fp16), ssq is over the rounded values + # (same as rms_norm reading the fp16 residual), the amax for + # dyn_gu is folded into the xn write pass, and the residual is + # register-cached so xn is never re-read from global. + fvk.residual_add_rms_norm_quantize_dynamic_fp8_fp16( + x_ptr, o_proj_out_ptr, int(weights['post_ln_w'][li]), + xn_ptr, xn_fp8_ptr, dyn_gu, Se, D, 1e-5, int(stream)) + gemm.fp8_nn_dev_fp16(xn_fp8_ptr, int(weights['gate_w'][li]), + int(bufs['gate_out']), Se, Dff, D, + dyn_gu, d_w_gu, int(stream)) + gemm.fp8_nn_dev_fp16(xn_fp8_ptr, int(weights['up_w'][li]), + int(bufs['up_out']), Se, Dff, D, + dyn_gu, d_w_gu, int(stream)) + + # Down input gu = silu(gate)*up in FP16, fused with its own + # dynamic FP8 quantize on layers with no outlier clamp (amax + # measured inside the SwiGLU write pass, skipping a separate + # absmax_kernel read of the Se*Dff intermediate). Layers that + # need the outlier clamp (default: L31 only) keep the + # unfused gate_geglu -> clamp -> quantize sequence since the + # clamp must run BEFORE the amax/scale is computed. + if clamp_this_layer: + fvk.gate_geglu_fp16(int(bufs['gate_out']), int(bufs['up_out']), + int(bufs['gate_out']), Se * Dff, int(stream)) + fvk.clamp_inplace_fp16(int(bufs['gate_out']), + float(ffn_down_clamp_value), + Se * Dff, int(stream)) + fvk.quantize_fp8_device_fp16( + int(bufs['gate_out']), int(bufs['gu_fp8']), dyn_d, + Se * Dff, int(stream)) + else: + fvk.gate_geglu_quantize_dynamic_fp8_fp16( + int(bufs['gate_out']), int(bufs['up_out']), + int(bufs['gate_out']), int(bufs['gu_fp8']), dyn_d, + Se * Dff, int(stream)) + gemm.fp8_nn_dev_fp16(int(bufs['gu_fp8']), int(weights['d_w'][li]), + o_proj_out_ptr, Se, D, Dff, dyn_d, d_w_d, + int(stream)) + + if clamp_this_layer: + fvk.clamp_inplace_fp16(o_proj_out_ptr, + float(ffn_down_clamp_value), + Se * D, int(stream)) + + # Residual 2 + next-layer prep. + if li < L - 1: + fvk.residual_add_rms_norm_fp16( + x_ptr, o_proj_out_ptr, int(weights['input_ln_w'][li + 1]), + xn_ptr, Se, D, 1e-5, int(stream)) + # A following STATIC FP8 layer consumes xn_fp8; produce it. + if (li + 1) not in dynamic_fp8_layers: + fvk.quantize_fp8_static_fp16( + xn_ptr, xn_fp8_ptr, + int(scales_dev["act_qkv"][li + 1]), Se * D, int(stream)) + else: + fvk.residual_add_fp16(x_ptr, o_proj_out_ptr, Se * D, + int(stream)) + + if probe is not None: + _pl = probe.get('layers') or () + if li in _pl: + _gpu_copy(int(probe['bufs'][_pl.index(li)]), + x_ptr, Se * D * 2, stream) + continue + # ═══ End dynamic per-tensor FP8 branch ═══ + + # ── QKV FP8 GEMMs ── + # Chameleon has no attention bias; pass zero-buffer for fp8_nn_bias epilogue + if alpha_host is not None: + alpha_qkv = float(alpha_host[li * 4 + 0]) + zero_bias_ptr = int(bufs['zero_bias_d']) + + # V projection: AWQ path OR shared-scale path. + # IMPORTANT ORDER: The chameleon attention backend aliases + # Q_ptr to xn_ptr (bufs['xn']). If Q is written first, xn is + # clobbered, and the AWQ V path reads Q's output instead of + # the RMSNormed xn. So run AWQ V-quantize + V-GEMM BEFORE Q + # (V's destination is the KV cache, separate buffer, safe). + if awq_v_proj is not None: + fvk.awq_quant_fp8_static_fp16( + xn_ptr, + int(awq_v_proj['inv_s_ptrs'][li]), + int(awq_v_proj['xn_v_fp8']), + int(awq_v_proj['act_scale_ptrs'][li]), + Se, D, int(stream), + ) + v_alpha = float(awq_v_proj['alpha_host'][li]) + gemm.fp8_nn_bias( + int(awq_v_proj['xn_v_fp8']), + int(awq_v_proj['w_ptrs'][li]), + V_ptr, zero_bias_ptr, + Se, D, D, v_alpha, int(stream), + ) + else: + gemm.fp8_nn_bias( + xn_fp8_ptr, v_w_ptr, V_ptr, zero_bias_ptr, + Se, D, D, alpha_qkv, int(stream), + ) + + # Q and K: also use AWQ smoothed path (same xn_v_fp8, per-proj weights) + if awq_v_proj is not None and 'w_q_ptrs' in awq_v_proj: + q_alpha = float(awq_v_proj['alpha_q_host'][li]) + k_alpha = float(awq_v_proj['alpha_k_host'][li]) + gemm.fp8_nn_bias( + int(awq_v_proj['xn_v_fp8']), + int(awq_v_proj['w_q_ptrs'][li]), + Q_ptr, zero_bias_ptr, + Se, D, D, q_alpha, int(stream), + ) + gemm.fp8_nn_bias( + int(awq_v_proj['xn_v_fp8']), + int(awq_v_proj['w_k_ptrs'][li]), + K_ptr, zero_bias_ptr, + Se, D, D, k_alpha, int(stream), + ) + else: + gemm.fp8_nn_bias( + xn_fp8_ptr, q_w_ptr, Q_ptr, zero_bias_ptr, + Se, D, D, alpha_qkv, int(stream), + ) + gemm.fp8_nn_bias( + xn_fp8_ptr, k_w_ptr, K_ptr, zero_bias_ptr, + Se, D, D, alpha_qkv, int(stream), + ) + else: + gemm.fp8_nn_dev_fp16( + xn_fp8_ptr, q_w_ptr, Q_ptr, + Se, D, D, d_act_qkv, d_w_qkv, int(stream), + ) + gemm.fp8_nn_dev_fp16( + xn_fp8_ptr, k_w_ptr, K_ptr, + Se, D, D, d_act_qkv, d_w_qkv, int(stream), + ) + gemm.fp8_nn_dev_fp16( + xn_fp8_ptr, v_w_ptr, V_ptr, + Se, D, D, d_act_qkv, d_w_qkv, int(stream), + ) + + # ── Fused per-head QK LayerNorm + RoPE ── + fvk.qk_norm_rope_fused_fp16( + Q_ptr, K_ptr, + int(weights['q_norm_w'][li]), int(weights['q_norm_b'][li]), + int(weights['k_norm_w'][li]), int(weights['k_norm_b'][li]), + cos_ptr, sin_ptr, + Se, H, Hd, 1e-5, int(stream), + ) + + # ── MHA via attention backend ── + attn.run("chameleon", li, q_seq=Se, kv_seq=Se, stream=int(stream)) + + # ── O projection ── + fvk.quantize_fp8_static_fp16( + O_ptr, xn_fp8_ptr, d_act_o, Se * D, int(stream), + ) + + if alpha_host is not None: + alpha_o = float(alpha_host[li * 4 + 1]) + gemm.fp8_nn_bias( + xn_fp8_ptr, int(weights['o_w'][li]), o_proj_out_ptr, + int(bufs['zero_bias_d']), + Se, D, D, alpha_o, int(stream), + ) + else: + gemm.fp8_nn_dev_fp16( + xn_fp8_ptr, int(weights['o_w'][li]), o_proj_out_ptr, + Se, D, D, d_act_o, d_w_o, int(stream), + ) + + # ── Post-attention: residual + RMSNorm (path depends on FFN precision) ── + if li in _FFN_FP8_LAYERS: + fvk.residual_add_rms_norm_fp8_fp16( + x_ptr, o_proj_out_ptr, int(weights['post_ln_w'][li]), + xn_fp8_ptr, Se, D, 1e-5, + d_act_gu, int(stream), + ) + else: + fvk.residual_add_fp16( + x_ptr, o_proj_out_ptr, Se * D, int(stream), + ) + fvk.rms_norm_fp16( + x_ptr, int(weights['post_ln_w'][li]), xn_ptr, + Se, D, 1e-5, int(stream), + ) + + # ── FFN (legacy static path; only for layers excluded from + # dynamic_fp8_layers): NVFP4 if li not in _FFN_FP8_LAYERS, else FP8 ── + gu_w_fp4_li = int(weights['gu_w_fp4'][li]) if li not in _FFN_FP8_LAYERS else 0 + fp4_available = ( + _fvk_fp4 is not None + and li not in _FFN_FP8_LAYERS + and gu_w_fp4_li != 0 + ) + if fp4_available: + # FP4 path + _fvk_fp4.quantize_fp4_dynamic_sfa_fp16( + xn_ptr, act_fp4_ptr, act_sfa_ptr, + Se, D, False, int(stream), + ) + _fvk_fp4.cutlass_fp4_gemm_variant( + FP4_VARIANT, + act_fp4_ptr, act_sfa_ptr, + int(weights['gu_w_fp4'][li]), int(weights['gu_sfb'][li]), + gu_merged_ptr, Se, 2 * Dff, D, + 1.0, 0.0, int(stream), + ) + _fvk_fp4.gate_geglu_fp4_sfa_v2_fp16( + gu_merged_ptr, ffn_act_fp4_ptr, ffn_act_sfa_ptr, + Se, Dff, int(stream), + ) + _fvk_fp4.cutlass_fp4_gemm_variant( + FP4_VARIANT, + ffn_act_fp4_ptr, ffn_act_sfa_ptr, + int(weights['d_w_fp4'][li]), int(weights['d_sfb'][li]), + o_proj_out_ptr, Se, D, Dff, + 1.0, 0.0, int(stream), + ) + else: + # FP8 path (also used as fallback when FP4 is unavailable + # for a layer in the L0-2 range). + gate_w_ptr = int(weights['gate_w'][li]) + up_w_ptr = int(weights['up_w'][li]) + gemm.fp8_nn_dev_fp16( + xn_fp8_ptr, gate_w_ptr, int(bufs['gate_out']), + Se, Dff, D, d_act_gu, d_w_gu, int(stream), + ) + gemm.fp8_nn_dev_fp16( + xn_fp8_ptr, up_w_ptr, int(bufs['up_out']), + Se, Dff, D, d_act_gu, d_w_gu, int(stream), + ) + # Down-proj: 2-tier adaptive dispatch. + # Tier 1 (standard FP8): no AWQ dict → fused silu_mul_split_fp8 + # Tier 2 (AWQ D smooth): li in awq_d_layers → per-K smooth + FP8 + _has_awq = (awq_v_proj is not None + and 'inv_s_D_ptrs' in awq_v_proj) + _in_awq_d = (_has_awq and li in awq_v_proj.get( + 'awq_d_layers', frozenset())) + + if _in_awq_d: + # Tier 2: AWQ D smoothed FP8 path (fused SwiGLU). + fvk.gate_geglu_fp16(int(bufs['gate_out']), int(bufs['up_out']), + int(bufs['gate_out']), Se * Dff, int(stream)) + fvk.awq_quant_fp8_static_fp16( + int(bufs['gate_out']), + int(awq_v_proj['inv_s_D_ptrs'][li]), + int(bufs['gu_fp8']), + int(awq_v_proj['act_scale_D_ptrs'][li]), + Se, Dff, int(stream)) + gemm.fp8_nn_bias( + int(bufs['gu_fp8']), + int(awq_v_proj['w_D_ptrs'][li]), + o_proj_out_ptr, int(bufs['zero_bias_d']), + Se, D, Dff, + float(awq_v_proj['alpha_D_host'][li]), + int(stream)) + else: + # Tier 1: standard per-tensor FP8 (no AWQ) + fvk.silu_mul_split_fp8_fp16( + int(bufs['gate_out']), int(bufs['up_out']), + int(bufs['gu_fp8']), + Se * Dff, d_act_d, int(stream), + ) + gemm.fp8_nn_dev_fp16( + int(bufs['gu_fp8']), int(weights['d_w'][li]), + o_proj_out_ptr, + Se, D, Dff, d_act_d, d_w_d, int(stream), + ) + + # ── Clamp down_out to fp16 range ── + # Chameleon-7B L31 down_proj's FP32 accumulator × alpha can push + # the FP16 output beyond ±65504 producing inf, which propagates + # through the final RMSNorm and destroys action precision. Mirror + # the RTX FP8 path's clamp (pipeline_rtx.py:829-833). + if clamp_this_layer: + fvk.clamp_inplace_fp16( + o_proj_out_ptr, float(ffn_down_clamp_value), + Se * D, int(stream), + ) + + # ── Fused: residual_2 + next-layer input_ln + FP8 quantize ── + # When AWQ V-proj is enabled we need the FP16 xn buffer populated + # for the NEXT layer's V-proj activation smoothing kernel (which + # reads FP16 xn, not FP8). Split the fused path into + # residual_add_rms_norm_fp16 + quantize_fp8_static_fp16 (one + # extra launch per layer) so xn_ptr stays valid. + if li < L - 1: + if awq_v_proj is not None: + fvk.residual_add_rms_norm_fp16( + x_ptr, o_proj_out_ptr, + int(weights['input_ln_w'][li + 1]), + xn_ptr, Se, D, 1e-5, int(stream), + ) + fvk.quantize_fp8_static_fp16( + xn_ptr, xn_fp8_ptr, + int(scales_dev["act_qkv"][li + 1]), + Se * D, int(stream), + ) + else: + fvk.residual_add_rms_norm_fp8_fp16( + x_ptr, o_proj_out_ptr, + int(weights['input_ln_w'][li + 1]), + xn_fp8_ptr, Se, D, 1e-5, + int(scales_dev["act_qkv"][li + 1]), int(stream), + ) + else: + fvk.residual_add_fp16(x_ptr, o_proj_out_ptr, Se * D, int(stream)) + + # ── Optional layer probe: snapshot post-residual-2 hidden state ── + if probe is not None: + probe_layers = probe.get('layers') or () + if li in probe_layers: + idx = probe_layers.index(li) + snap_ptr = int(probe['bufs'][idx]) + if snap_ptr != 0: + _gpu_copy(snap_ptr, x_ptr, Se * D * 2, stream) + + # ── Final RMSNorm → hidden_all [Se, D] FP16 ── + fvk.rms_norm_fp16( + x_ptr, int(weights['final_norm_w']), hidden_all_ptr, + Se, D, 1e-5, int(stream), + ) + + # ── Optional final probe ── + if probe is not None: + final_buf = int(probe.get('final_buf', 0)) + if final_buf != 0: + _gpu_copy(final_buf, hidden_all_ptr, Se * D * 2, stream) + + +def chameleon_decode_step( + gemm, fvk, bufs, weights, dims, scales_dev, + *, attn, pos: int, stream: int = 0, + ffn_down_clamp_value: float = 60000.0, + ffn_clamp_layers=frozenset({31}), +) -> None: + """Single-token (Se=1) incremental decode step at position ``pos``. + + Mirrors chameleon_forward's dynamic per-tensor FP8 branch op-for-op, + with: + * K/V GEMMs writing the KV cache row at ``pos`` + (``attn.kv_row_ptrs``) — history rows were filled by prefill; + * RoPE cos/sin taken at row ``pos`` via byte offsets; + * attention dispatched through ``attn.run_decode`` (bottom-right + causal mask) over kv_len=pos+1 keys; + * final RMSNorm written to ``hidden_all`` row 0. + + The token embedding must already be in ``bufs['x']`` row 0, and the + residual stream in ``bufs['x']`` is updated in place across layers. + CUDA-Graph safe (int pointers only, no allocations); ``pos`` is a host + scalar, so callers run this eagerly, outside graph capture. + """ + D = int(dims['D']) + Dff = int(dims['Dff']) + L = int(dims['L']) + H = int(dims['H']) + Hd = int(dims['Hd']) + + x_ptr = int(bufs['x']) + xn_ptr = int(bufs['xn']) + xn_fp8_ptr = int(bufs['xn_fp8']) + o_proj_out_ptr = int(bufs['o_proj_out']) + hidden_all_ptr = int(bufs['hidden_all']) + + cos_ptr = int(weights['rope_cos']) + pos * Hd * 2 + sin_ptr = int(weights['rope_sin']) + pos * Hd * 2 + + for li in range(L): + clamp_this_layer = ( + ffn_down_clamp_value > 0.0 + and (ffn_clamp_layers is None or li in ffn_clamp_layers) + ) + dyn_qkv = int(scales_dev['dyn_act_qkv'][li]) + dyn_o = int(scales_dev['dyn_act_o'][li]) + dyn_gu = int(scales_dev['dyn_act_gu'][li]) + dyn_d = int(scales_dev['dyn_act_down'][li]) + + d_w_o = int(weights['d_w_o'][li]) + d_w_gu = int(weights['d_w_gu'][li]) + d_w_d = int(weights['d_w_d'][li]) + + K_row_ptr, V_row_ptr = attn.kv_row_ptrs("chameleon", li, pos) + slots = attn.get_slot_ptrs("chameleon", li) + Q_ptr = int(slots["Q"]) + O_ptr = int(slots["O"]) + + # Fused RMSNorm + dynamic FP8 quantize from the residual stream. + fvk.rms_norm_quantize_dynamic_fp8_fp16( + x_ptr, int(weights['input_ln_w'][li]), xn_ptr, xn_fp8_ptr, + dyn_qkv, 1, D, 1e-5, int(stream)) + + # QKV GEMMs (M=1): K/V land in the cache row at pos, Q in Q_O. + gemm.fp8_nn_dev_fp16(xn_fp8_ptr, int(weights['v_w'][li]), V_row_ptr, + 1, D, D, dyn_qkv, int(weights['d_w_qkv'][li]), + int(stream)) + gemm.fp8_nn_dev_fp16(xn_fp8_ptr, int(weights['k_w'][li]), K_row_ptr, + 1, D, D, dyn_qkv, int(weights['d_w_qkv'][li]), + int(stream)) + gemm.fp8_nn_dev_fp16(xn_fp8_ptr, int(weights['q_w'][li]), Q_ptr, + 1, D, D, dyn_qkv, int(weights['d_w_qkv'][li]), + int(stream)) + + fvk.qk_norm_rope_fused_fp16( + Q_ptr, K_row_ptr, + int(weights['q_norm_w'][li]), int(weights['q_norm_b'][li]), + int(weights['k_norm_w'][li]), int(weights['k_norm_b'][li]), + cos_ptr, sin_ptr, + 1, H, Hd, 1e-5, int(stream)) + attn.run_decode("chameleon", li, kv_len=pos + 1, stream=int(stream)) + + # O projection (M=1). + fvk.quantize_fp8_device_fp16( + O_ptr, xn_fp8_ptr, dyn_o, D, int(stream)) + gemm.fp8_nn_dev_fp16(xn_fp8_ptr, int(weights['o_w'][li]), + o_proj_out_ptr, 1, D, D, dyn_o, d_w_o, + int(stream)) + + # Residual 1. + fvk.residual_add_fp16(x_ptr, o_proj_out_ptr, D, int(stream)) + + # FFN: residual add + RMSNorm + dynamic FP8 quantize (fused). + fvk.residual_add_rms_norm_quantize_dynamic_fp8_fp16( + x_ptr, o_proj_out_ptr, int(weights['post_ln_w'][li]), + xn_ptr, xn_fp8_ptr, dyn_gu, 1, D, 1e-5, int(stream)) + gemm.fp8_nn_dev_fp16(xn_fp8_ptr, int(weights['gate_w'][li]), + int(bufs['gate_out']), 1, Dff, D, + dyn_gu, d_w_gu, int(stream)) + gemm.fp8_nn_dev_fp16(xn_fp8_ptr, int(weights['up_w'][li]), + int(bufs['up_out']), 1, Dff, D, + dyn_gu, d_w_gu, int(stream)) + + if clamp_this_layer: + fvk.gate_geglu_fp16(int(bufs['gate_out']), int(bufs['up_out']), + int(bufs['gate_out']), Dff, int(stream)) + fvk.clamp_inplace_fp16(int(bufs['gate_out']), + float(ffn_down_clamp_value), + Dff, int(stream)) + fvk.quantize_fp8_device_fp16( + int(bufs['gate_out']), int(bufs['gu_fp8']), dyn_d, + Dff, int(stream)) + else: + fvk.gate_geglu_quantize_dynamic_fp8_fp16( + int(bufs['gate_out']), int(bufs['up_out']), + int(bufs['gate_out']), int(bufs['gu_fp8']), dyn_d, + Dff, int(stream)) + gemm.fp8_nn_dev_fp16(int(bufs['gu_fp8']), int(weights['d_w'][li]), + o_proj_out_ptr, 1, D, Dff, dyn_d, d_w_d, + int(stream)) + + if clamp_this_layer: + fvk.clamp_inplace_fp16(o_proj_out_ptr, + float(ffn_down_clamp_value), + D, int(stream)) + + # Residual 2 (next layer re-derives xn from the residual stream). + fvk.residual_add_fp16(x_ptr, o_proj_out_ptr, D, int(stream)) + + # Final RMSNorm → hidden_all row 0. + fvk.rms_norm_fp16( + x_ptr, int(weights['final_norm_w']), hidden_all_ptr, + 1, D, 1e-5, int(stream), + ) + + +def chameleon_forward_fp16( + gemm, fvk, bufs, weights, dims, + *, attn, stream: int = 0, + ffn_gate_clamp_value: float = 10000.0, + probe=None, +) -> None: + """Chameleon-7B FP16-only forward (no FP8, no FP4, no AWQ). + + Ported from pipeline_rtx.chameleon_forward. All 32 layers run pure + FP16 GEMMs via ``gemm.fp16_nn`` — same on Thor as RTX. + + Precision-optimal path (cosine target ≥ 0.99 vs HF bf16) at the + cost of ~2× the FP8 path latency in the LLM. Recommended when + downstream ActionHead is sensitive to accumulated FP8 error. + + Weights required (all FP16, KN row-major layout — spec built with + ``use_fp8=False``): + q_w[li], k_w[li], v_w[li], o_w[li] : (D, D) + gate_w[li], up_w[li] : (D, Dff) + d_w[li] / down_w[li] : (Dff, D) + input_ln_w[li], post_ln_w[li] : (D,) + q_norm_w/b[li], k_norm_w/b[li] : (1, HD) or (HD,) + final_norm_w : (D,) + rope_cos, rope_sin : (max_pos, HD) + + Output: hidden_all = RMSNorm(x_post_res_2, final_norm_w) written to + bufs['hidden_all'] as (Se, D) FP16 (consumed by action_head_forward). + """ + Se = int(dims['Se']) + D = int(dims['D']) + Dff = int(dims['Dff']) + L = int(dims['L']) + H = int(dims['H']) + Hd = int(dims['Hd']) + + x_ptr = int(bufs['x']) + xn_ptr = int(bufs['xn']) + o_proj_out_ptr = int(bufs['o_proj_out']) + hidden_all_ptr = int(bufs['hidden_all']) + # Reuse existing gate/up buffers (Se, Dff) fp16 + gate_ptr = int(bufs['gate_out']) + up_ptr = int(bufs['up_out']) + + cos_ptr = int(weights['rope_cos']) + sin_ptr = int(weights['rope_sin']) + + q_w = weights['q_w'] + k_w = weights['k_w'] + v_w = weights['v_w'] + o_w = weights['o_w'] + gate_w = weights['gate_w'] + up_w = weights['up_w'] + down_w = weights['d_w'] # frontend uses 'd_w' key for down projection + + input_ln_w = weights['input_ln_w'] + post_ln_w = weights['post_ln_w'] + q_norm_w = weights['q_norm_w'] + q_norm_b = weights['q_norm_b'] + k_norm_w = weights['k_norm_w'] + k_norm_b = weights['k_norm_b'] + final_norm_w = int(weights['final_norm_w']) + + for li in range(L): + slots = attn.get_slot_ptrs("chameleon", li) + Q_ptr = int(slots["Q"]) + K_ptr = int(slots["K"]) + V_ptr = int(slots["V"]) + + # input_layernorm (RMSNorm eps=1e-5) → xn + fvk.rms_norm_fp16( + x_ptr, int(input_ln_w[li]), xn_ptr, + Se, D, 1e-5, int(stream), + ) + + # Q / K / V GEMMs (no bias in Chameleon). + # NOTE: Q_ptr aliases xn_ptr on Thor (chameleon slots["Q_O"] = + # bufs['xn'].data_ptr()). Because gemm.fp16_nn reads A (xn) and + # writes D (Q) at the SAME fp16 dtype and SAME buffer, cuBLAS + # in-place semantics are undefined → output corruption. + # Route Q through o_proj_out_ptr scratch, then copy into Q slot. + # (V and K write to separate KV cache buffers; safe direct.) + gemm.fp16_nn(xn_ptr, int(v_w[li]), V_ptr, Se, D, D, int(stream)) + gemm.fp16_nn(xn_ptr, int(k_w[li]), K_ptr, Se, D, D, int(stream)) + gemm.fp16_nn(xn_ptr, int(q_w[li]), o_proj_out_ptr, + Se, D, D, int(stream)) + _gpu_copy(Q_ptr, o_proj_out_ptr, Se * D * 2, stream) + + # Per-head QK LayerNorm + RoPE (in-place). + fvk.qk_norm_rope_fused_fp16( + Q_ptr, K_ptr, + int(q_norm_w[li]), int(q_norm_b[li]), + int(k_norm_w[li]), int(k_norm_b[li]), + cos_ptr, sin_ptr, + Se, H, Hd, 1e-5, int(stream), + ) + + # Causal MHA (CUTLASS SM110 FMHA / cuBLAS fallback). + attn.run("chameleon", li, q_seq=Se, kv_seq=Se, stream=int(stream)) + + # O projection. + gemm.fp16_nn(Q_ptr, int(o_w[li]), o_proj_out_ptr, + Se, D, D, int(stream)) + + # Residual 1. + fvk.residual_add_fp16(x_ptr, o_proj_out_ptr, Se * D, int(stream)) + + # post_attention_layernorm. + fvk.rms_norm_fp16( + x_ptr, int(post_ln_w[li]), xn_ptr, + Se, D, 1e-5, int(stream), + ) + + # FFN gate / up (SwiGLU). + gemm.fp16_nn(xn_ptr, int(gate_w[li]), gate_ptr, + Se, Dff, D, int(stream)) + gemm.fp16_nn(xn_ptr, int(up_w[li]), up_ptr, + Se, Dff, D, int(stream)) + fvk.gate_geglu_fp16(gate_ptr, up_ptr, gate_ptr, Se * Dff, int(stream)) + + # L31 gate*up overflow guard (fp16 max ≈ 65504, Chameleon L31 + # gate*up amax observed ≈ 48000). See pipeline_rtx.py:243-256. + if ffn_gate_clamp_value > 0.0: + fvk.clamp_inplace_fp16( + gate_ptr, float(ffn_gate_clamp_value), + Se * Dff, int(stream), + ) + + # down projection. + gemm.fp16_nn(gate_ptr, int(down_w[li]), o_proj_out_ptr, + Se, D, Dff, int(stream)) + + # Residual 2. + fvk.residual_add_fp16(x_ptr, o_proj_out_ptr, Se * D, int(stream)) + + # Optional probe: snapshot post-residual-2 hidden state. + if probe is not None: + probe_layers = probe.get('layers') or () + if li in probe_layers: + idx = probe_layers.index(li) + snap_ptr = int(probe['bufs'][idx]) + if snap_ptr != 0: + _gpu_copy(snap_ptr, x_ptr, Se * D * 2, stream) + + # Final RMSNorm → hidden_all. + fvk.rms_norm_fp16( + x_ptr, final_norm_w, hidden_all_ptr, + Se, D, 1e-5, int(stream), + ) + + if probe is not None: + final_buf = int(probe.get('final_buf', 0)) + if final_buf != 0: + _gpu_copy(final_buf, hidden_all_ptr, Se * D * 2, stream) + + +# ══════════════════════════════════════════════════════════════════ +# Chameleon-7B LLM calibration (FP16 + amax measurement) +# ══════════════════════════════════════════════════════════════════ + +def _d2h_float(d_ptr: int) -> float: + """Read a single float32 from device to host.""" + t = torch.empty(1, dtype=torch.float32, device='cuda') + import ctypes + ctypes.CDLL('libcudart.so').cudaMemcpy( + ctypes.c_void_p(t.data_ptr()), + ctypes.c_void_p(d_ptr), + 4, 2, # cudaMemcpyDeviceToDevice is 3, D2H is 2 + ) + return float(t.item()) + + +def _d2h_floats(d_ptr: int, n: int) -> list: + """Read n float32 values from device to host.""" + t = torch.empty(n, dtype=torch.float32, device='cuda') + import ctypes + ctypes.CDLL('libcudart.so').cudaMemcpy( + ctypes.c_void_p(t.data_ptr()), + ctypes.c_void_p(d_ptr), + n * 4, 2, + ) + return t.cpu().tolist() + + +def chameleon_forward_calibrate( + gemm, fvk_mod, bufs, weights, dims, + calib_scales_ptr, stream: int = 0, + attn_calib_scales_ptr: int = 0, +) -> None: + """Calibrate Chameleon-7B FP8 scales. + + 4 quantization points per layer × 32 layers = 128 scales. + Points: act_qkv, act_o, act_gu, act_down. + """ + Se = int(dims['Se']) + D = int(dims['D']) + Dff = int(dims['Dff']) + L = int(dims['L']) + H = int(dims['H']) + Hd = int(dims['Hd']) + + import numpy as np + + x_ptr = int(bufs['x']) + xn_ptr = int(bufs['xn']) + o_proj_out_ptr = int(bufs['o_proj_out']) + gate_out_ptr = int(bufs['gate_out']) + up_out_ptr = int(bufs['up_out']) + down_out_ptr = int(bufs['down_out']) + Q_ptr_buf = int(bufs['Q']) + K_ptr_buf = int(bufs['K']) + V_ptr_buf = int(bufs['V']) + O_ptr_buf = int(bufs['O']) + + calib_buf = int(bufs['calib_buf']) + d_scale = int(bufs['d_scale']) + fp8_scratch = int(bufs['fp8_scratch']) + norm_scratch = int(bufs['norm_scratch']) + + cos_ptr = int(weights['rope_cos']) + sin_ptr = int(weights['rope_sin']) + + w_scales_dev = int(weights['w_scales_flat']) + ws_host = _d2h_floats(w_scales_dev, L * 4) + + _gpu_zero(calib_buf, L * 4 * 4, stream) + + for li in range(L): + # ── 1. RMSNorm → measure amax (act_qkv scale) ── + fvk_mod.rms_norm_fp16( + x_ptr, int(weights['input_ln_w'][li]), norm_scratch, + Se, D, 1e-5, int(stream), + ) + _measure_scale_gpu(fvk_mod, norm_scratch, Se * D, d_scale, fp8_scratch, stream) + _gpu_sync(stream) + as_qkv = _d2h_float(d_scale) + cs_qkv = calib_buf + (li * 4 + 0) * 4 + _gpu_copy(cs_qkv, d_scale, 4, stream) + + # ── 2. Quantize xn → FP8 ── + fvk_mod.quantize_fp8_static_fp16( + norm_scratch, int(bufs['xn_fp8']), cs_qkv, + Se * D, int(stream), + ) + + # ── 3. Q/K/V FP8 GEMMs (no bias in Chameleon) ── + q_w_ptr = int(weights['q_w'][li]) + k_w_ptr = int(weights['k_w'][li]) + v_w_ptr = int(weights['v_w'][li]) + + alpha_qkv = float(np.float32(as_qkv) * np.float32(ws_host[li * 4 + 0])) + zero_bias_ptr = int(bufs['zero_bias_d']) + + gemm.fp8_nn_bias( + int(bufs['xn_fp8']), q_w_ptr, Q_ptr_buf, zero_bias_ptr, + Se, D, D, alpha_qkv, int(stream), + ) + gemm.fp8_nn_bias( + int(bufs['xn_fp8']), k_w_ptr, K_ptr_buf, zero_bias_ptr, + Se, D, D, alpha_qkv, int(stream), + ) + gemm.fp8_nn_bias( + int(bufs['xn_fp8']), v_w_ptr, V_ptr_buf, zero_bias_ptr, + Se, D, D, alpha_qkv, int(stream), + ) + + # ── 4. Fused QK LayerNorm + RoPE ── + fvk_mod.qk_norm_rope_fused_fp16( + Q_ptr_buf, K_ptr_buf, + int(weights['q_norm_w'][li]), int(weights['q_norm_b'][li]), + int(weights['k_norm_w'][li]), int(weights['k_norm_b'][li]), + cos_ptr, sin_ptr, + Se, H, Hd, 1e-5, int(stream), + ) + + # ── 4b. Optional Q/K/V amax for FP8 attention ── + if attn_calib_scales_ptr: + n_qkv = Se * H * Hd + for i, ptr in enumerate((Q_ptr_buf, K_ptr_buf, V_ptr_buf)): + _measure_scale_gpu( + fvk_mod, ptr, n_qkv, d_scale, fp8_scratch, stream) + _gpu_sync(stream) + _gpu_copy( + attn_calib_scales_ptr + (li * 3 + i) * 4, + d_scale, 4, stream, + ) + + # ── 5. Attention (cuBLAS — no FMHA during calibration) ── + attn_scale = 1.0 / math.sqrt(float(Hd)) + fvk_mod.attention_qkv_fp16( + bufs['ctx'], Q_ptr_buf, K_ptr_buf, V_ptr_buf, + int(bufs['logits']), O_ptr_buf, + Se, Se, H, Hd, attn_scale, int(stream), + ) + + # ── 6. O proj — measure amax → quantize → GEMM ── + _measure_scale_gpu(fvk_mod, O_ptr_buf, Se * D, d_scale, fp8_scratch, stream) + _gpu_sync(stream) + as_o = _d2h_float(d_scale) + cs_o = calib_buf + (li * 4 + 1) * 4 + _gpu_copy(cs_o, d_scale, 4, stream) + fvk_mod.quantize_fp8_static_fp16( + O_ptr_buf, int(bufs['xn_fp8']), cs_o, + Se * D, int(stream), + ) + alpha_o = float(np.float32(as_o) * np.float32(ws_host[li * 4 + 1])) + gemm.fp8_nn_bias( + int(bufs['xn_fp8']), int(weights['o_w'][li]), o_proj_out_ptr, + zero_bias_ptr, + Se, D, D, alpha_o, int(stream), + ) + + # ── 7. Residual 1 ── + fvk_mod.residual_add_fp16(x_ptr, o_proj_out_ptr, Se * D, int(stream)) + + # ── 8. Post-attn RMSNorm → measure amax (act_gu scale) ── + fvk_mod.rms_norm_fp16( + x_ptr, int(weights['post_ln_w'][li]), norm_scratch, + Se, D, 1e-5, int(stream), + ) + _measure_scale_gpu(fvk_mod, norm_scratch, Se * D, d_scale, fp8_scratch, stream) + _gpu_sync(stream) + as_gu = _d2h_float(d_scale) + cs_gu = calib_buf + (li * 4 + 2) * 4 + _gpu_copy(cs_gu, d_scale, 4, stream) + + # ── 9. Quantize → FP8 ── + fvk_mod.quantize_fp8_static_fp16( + norm_scratch, int(bufs['xn_fp8']), cs_gu, + Se * D, int(stream), + ) + + # ── 10. Gate + Up FP8 GEMMs (no bias) ── + gate_w_ptr = int(weights['gate_w'][li]) + up_w_ptr = int(weights['up_w'][li]) + alpha_gu = float(np.float32(as_gu) * np.float32(ws_host[li * 4 + 2])) + zero_bias_dff = int(bufs['zero_bias_dff']) + gemm.fp8_nn_bias( + int(bufs['xn_fp8']), gate_w_ptr, gate_out_ptr, + zero_bias_dff, + Se, Dff, D, alpha_gu, int(stream), + ) + gemm.fp8_nn_bias( + int(bufs['xn_fp8']), up_w_ptr, up_out_ptr, + zero_bias_dff, + Se, Dff, D, alpha_gu, int(stream), + ) + + # ── 11. SiLU(gate)*up → measure amax → FP8 ── + silu_scr = int(bufs['silu_scratch']) + _gpu_copy(silu_scr, gate_out_ptr, Se * Dff * 2, stream) + fvk_mod.gate_geglu_fp16(silu_scr, up_out_ptr, down_out_ptr, + Se * Dff, int(stream)) + _measure_scale_gpu(fvk_mod, down_out_ptr, Se * Dff, d_scale, fp8_scratch, stream) + _gpu_sync(stream) + as_d = _d2h_float(d_scale) + cs_d = calib_buf + (li * 4 + 3) * 4 + _gpu_copy(cs_d, d_scale, 4, stream) + fvk_mod.silu_mul_split_fp8_fp16( + gate_out_ptr, up_out_ptr, int(bufs['gu_fp8']), + Se * Dff, cs_d, int(stream), + ) + + # ── 12. Down FP8 GEMM ── + alpha_d = float(np.float32(as_d) * np.float32(ws_host[li * 4 + 3])) + gemm.fp8_nn_bias( + int(bufs['gu_fp8']), int(weights['d_w'][li]), o_proj_out_ptr, + zero_bias_ptr, + Se, D, Dff, alpha_d, int(stream), + ) + + # ── 13. Residual 2 ── + fvk_mod.residual_add_fp16(x_ptr, o_proj_out_ptr, Se * D, int(stream)) + + # ── Final RMSNorm ── + fvk_mod.rms_norm_fp16( + x_ptr, int(weights['final_norm_w']), int(bufs['xn']), + Se, D, 1e-5, int(stream), + ) + + # ── Copy calibrated scales to output ── + _gpu_copy(calib_scales_ptr, calib_buf, L * 4 * 4, stream) + _gpu_sync(stream) + + +__all__ = [ + "chameleon_forward", + "chameleon_forward_fp16", + "chameleon_forward_calibrate", +] diff --git a/flash_rt/models/chameleon/vqvae_hf.py b/flash_rt/models/chameleon/vqvae_hf.py new file mode 100644 index 000000000..d7eb2850c --- /dev/null +++ b/flash_rt/models/chameleon/vqvae_hf.py @@ -0,0 +1,124 @@ +"""Apache-2.0 Transformers-backed Chameleon VQ-VAE helpers.""" + +from __future__ import annotations + +import hashlib +import json +from pathlib import Path + +import numpy as np +import torch +from PIL import Image + + +_VQ_PREFIXES = ("model.vqmodel.", "vqmodel.") + + +def _checkpoint_index(checkpoint_dir: Path) -> tuple[dict[str, str], Path]: + index_path = checkpoint_dir / "model.safetensors.index.json" + if index_path.is_file(): + payload = json.loads(index_path.read_text(encoding="utf-8")) + weight_map = payload.get("weight_map") + if not isinstance(weight_map, dict): + raise ValueError(f"invalid weight_map in {index_path}") + return weight_map, index_path + + single = checkpoint_dir / "model.safetensors" + if single.is_file(): + from safetensors import safe_open + + with safe_open(single, framework="pt", device="cpu") as handle: + return {key: single.name for key in handle.keys()}, single + raise FileNotFoundError( + f"no model.safetensors or model.safetensors.index.json under " + f"{checkpoint_dir}") + + +def _vq_weight_map(checkpoint_dir: Path) -> tuple[dict[str, str], str, Path]: + weight_map, manifest = _checkpoint_index(checkpoint_dir) + for prefix in _VQ_PREFIXES: + selected = { + key[len(prefix):]: shard + for key, shard in weight_map.items() + if key.startswith(prefix) + } + if selected: + return selected, prefix, manifest + raise FileNotFoundError( + f"no Chameleon VQ-VAE weights ({', '.join(_VQ_PREFIXES)}) under " + f"{checkpoint_dir}") + + +def load_chameleon_vqvae( + checkpoint_dir: str | Path, + *, + device: str | torch.device = "cuda", + dtype: torch.dtype = torch.float16, +): + """Load only the HF Chameleon VQ-VAE from checkpoint safetensors.""" + from safetensors import safe_open + from transformers import ChameleonConfig, ChameleonVQVAE + from transformers.models.chameleon.modeling_chameleon import ( + ChameleonImageVocabularyMapping, + ) + + root = Path(checkpoint_dir).expanduser().resolve() + selected, prefix, _ = _vq_weight_map(root) + state_dict: dict[str, torch.Tensor] = {} + by_shard: dict[str, list[str]] = {} + for key, shard in selected.items(): + by_shard.setdefault(shard, []).append(key) + for shard, keys in by_shard.items(): + with safe_open(root / shard, framework="pt", device="cpu") as handle: + for key in keys: + state_dict[key] = handle.get_tensor(prefix + key) + + config = ChameleonConfig.from_pretrained(str(root)) + model = ChameleonVQVAE(config.vq_config) + incompatible = model.load_state_dict(state_dict, strict=True) + if incompatible.missing_keys or incompatible.unexpected_keys: + raise RuntimeError( + "Chameleon VQ-VAE checkpoint mismatch: " + f"missing={incompatible.missing_keys[:4]}, " + f"unexpected={incompatible.unexpected_keys[:4]}") + model = model.eval().to(device=device, dtype=dtype) + mapping = ChameleonImageVocabularyMapping(config.vocabulary_map) + return model, mapping + + +def preprocess_vqvae_image( + image: Image.Image, + *, + device: str | torch.device = "cuda", + dtype: torch.dtype = torch.float16, +) -> torch.Tensor: + """Convert a PIL image to the Chameleon VQ-VAE ``[-1, 1]`` tensor.""" + rgba = image.convert("RGBA") + background = Image.new("RGBA", rgba.size, (255, 255, 255, 255)) + rgb = Image.alpha_composite(background, rgba).convert("RGB") + values = np.asarray(rgb, dtype=np.float32) / 127.5 - 1.0 + tensor = torch.from_numpy(values).permute(2, 0, 1).unsqueeze(0) + return tensor.to(device=device, dtype=dtype).contiguous() + + +def encode_vqvae_tokens(model, pixel_values: torch.Tensor) -> torch.Tensor: + """Return image-token indices across Transformers 4.43+ output APIs.""" + output = model.encode(pixel_values) + if hasattr(output, "image_tokens"): + return output.image_tokens + return output[2] + + +def vqvae_checkpoint_digest(checkpoint_dir: str | Path) -> str: + """Return a stable digest of the VQ-VAE config and referenced shards.""" + root = Path(checkpoint_dir).expanduser().resolve() + selected, _, manifest = _vq_weight_map(root) + paths = [root / "config.json", manifest] + paths.extend(root / name for name in sorted(set(selected.values()))) + digest = hashlib.sha256() + for path in paths: + digest.update(path.name.encode("utf-8")) + with path.open("rb") as handle: + for chunk in iter(lambda: handle.read(1024 * 1024), b""): + digest.update(chunk) + return digest.hexdigest()[:16] diff --git a/pyproject.toml b/pyproject.toml index 1fd541501..22930a2a6 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -16,6 +16,7 @@ dependencies = [ [project.optional-dependencies] torch = ["torch", "safetensors", "sentencepiece"] +chameleon = ["torch", "safetensors", "transformers>=4.43,<4.56", "Pillow"] jax = ["jax", "jaxlib", "ml_dtypes", "orbax-checkpoint", "flax"] server = ["fastapi", "uvicorn"] eval = ["opencv-python", "tqdm"] @@ -42,7 +43,7 @@ flash-attn = ["flash-attn"] # (flash_rt/hardware/thor/fa4_backend.py) sets CUTE_DSL_ARCH=sm_101a. # Requires an editable / source-tree install (pip install -e). pip install ".[thor-fa4]". thor-fa4 = ["nvidia-cutlass-dsl==4.5.1", "quack-kernels==0.4.1"] -all = ["flash-rt[torch,jax,server,eval,minimax-remover]"] +all = ["flash-rt[torch,jax,server,eval,minimax-remover,chameleon]"] [tool.setuptools.packages.find] include = ["flash_rt*"] diff --git a/scripts/bench_chameleon_thor.py b/scripts/bench_chameleon_thor.py new file mode 100644 index 000000000..8c8e1aed1 --- /dev/null +++ b/scripts/bench_chameleon_thor.py @@ -0,0 +1,359 @@ +#!/usr/bin/env python3 +"""Real-image latency benchmark for standalone Chameleon-7B on Thor. + +Measures HF BF16, FlashRT FP16, and FlashRT dynamic FP8 prefill +latency on the same real-image prompt. Inputs are always real images +(from a user-supplied directory), never synthetic token ids. + +Usage +----- + PYTHONPATH=. python scripts/bench_chameleon_thor.py \\ + --checkpoint /path/to/Chameleon_7B_mGPT \\ + --image-dir /path/to/images \\ + --prompt "Describe the image." \\ + --iters 10 --warmup 2 \\ + --output /tmp/chameleon_thor_bench.json +""" + +from __future__ import annotations + +import argparse +import json +import pathlib +import time +from typing import Dict, List + +import numpy as np + + +def _stats(xs: List[float]) -> Dict[str, float]: + a = np.asarray(xs, dtype=np.float64) + return { + "mean": float(a.mean()), + "p50": float(np.percentile(a, 50)), + "min": float(a.min()), + "max": float(a.max()), + } + + +def _load_real_images(image_dir: pathlib.Path, max_images: int): + from PIL import Image + exts = (".jpg", ".jpeg", ".png", ".bmp") + paths = sorted(p for p in image_dir.iterdir() if p.suffix.lower() in exts) + if not paths: + raise FileNotFoundError(f"No real images found under {image_dir}") + paths = paths[:max_images] + return [Image.open(p).convert("RGB") for p in paths], [str(p) for p in paths] + + +def _pad_ids(input_ids: List[int], pad_id: int = 1) -> tuple[List[int], int]: + real_len = len(input_ids) + padded = list(input_ids) + rem = len(padded) % 16 + if rem: + padded.extend([pad_id] * (16 - rem)) + return padded, real_len + + +def _estimate_prefill_tflops(Se: int) -> float: + D, Dff, L, vocab = 4096, 11008, 32, 65536 + + def gemm_flops(M: int, N: int, K: int) -> int: + return 2 * M * N * K + + per_layer_gemm = ( + gemm_flops(Se, 3 * D, D) + + gemm_flops(Se, D, D) + + gemm_flops(Se, 2 * Dff, D) + + gemm_flops(Se, D, Dff) + ) + per_layer_attn = 4 * Se * Se * D + lm_head = gemm_flops(1, vocab, D) + return (L * (per_layer_gemm + per_layer_attn) + lm_head) / 1e12 + + +def _roofline(Se: int, prefill_ms: float, peak_tflops: float) -> Dict[str, float]: + tflops = _estimate_prefill_tflops(Se) + achieved = tflops / (prefill_ms / 1000.0) if prefill_ms > 0 else 0.0 + floor_ms = tflops / peak_tflops * 1000.0 if peak_tflops > 0 else 0.0 + return { + "estimated_tflops": float(tflops), + "assumed_peak_tflops": float(peak_tflops), + "achieved_tflops": float(achieved), + "efficiency_vs_peak": float(achieved / peak_tflops) if peak_tflops > 0 else 0.0, + "optimistic_compute_floor_ms": float(floor_ms), + "measured_over_floor": float(prefill_ms / floor_ms) if floor_ms > 0 else 0.0, + } + + +def _run_flashrt_prefill_once(fe, prompt: str, images, cached_ids, + *, use_cuda_graph: bool): + import torch + + times = {} + t0 = time.perf_counter() + if cached_ids is None: + ids = fe.encode_prompt(prompt, images) + else: + ids = cached_ids + torch.cuda.synchronize() + t1 = time.perf_counter() + + padded, real_len = _pad_ids(ids) + fe._real_len = real_len + fe.Se = len(padded) + fe._last_input_ids = padded + if fe._use_autotune: + fe._autotune_gemms(fe.Se) + torch.cuda.synchronize() + t2 = time.perf_counter() + + fe._embed_ids(padded) + torch.cuda.synchronize() + t3 = time.perf_counter() + + if use_cuda_graph: + fe._capture_graph(fe.Se) + fe._infer_graph.replay() + else: + fe._run_backbone(fe.Se) + torch.cuda.synchronize() + t4 = time.perf_counter() + + fe._project_last() + torch.cuda.synchronize() + t5 = time.perf_counter() + + times["encode_ms"] = (t1 - t0) * 1000.0 + times["prepare_ms"] = (t2 - t1) * 1000.0 + times["embed_ms"] = (t3 - t2) * 1000.0 + times["backbone_ms"] = (t4 - t3) * 1000.0 + times["lm_head_ms"] = (t5 - t4) * 1000.0 + times["transformer_prefill_ms"] = times["embed_ms"] + times["backbone_ms"] + times["lm_head_ms"] + times["total_ms"] = (t5 - t0) * 1000.0 + return times, fe.Se, real_len, fe.vqgan_backend + + +def _bench_flashrt(checkpoint_dir: pathlib.Path, prompt: str, images, + *, use_fp8: bool, use_cuda_graph: bool, + target_size: int, use_trt_vqgan: bool, + trt_vqgan_engine_dir: str | None, + iters: int, warmup: int, + reuse_input_ids: bool, + generate_greedy: int, + peak_tflops: float) -> Dict: + import torch + from flash_rt.frontends.torch.chameleon_thor import ChameleonTorchFrontendThor + + fe = ChameleonTorchFrontendThor( + str(checkpoint_dir), use_fp8=use_fp8, use_cuda_graph=use_cuda_graph, + target_size=target_size, use_trt_vqgan=use_trt_vqgan, + trt_vqgan_engine_dir=trt_vqgan_engine_dir) + + cached_ids = None + input_build_ms = None + if reuse_input_ids: + t0 = time.perf_counter() + cached_ids = fe.encode_prompt(prompt, images) + torch.cuda.synchronize() + input_build_ms = (time.perf_counter() - t0) * 1000.0 + + for _ in range(warmup): + _run_flashrt_prefill_once( + fe, prompt, images, cached_ids, use_cuda_graph=use_cuda_graph) + + stage_values: Dict[str, List[float]] = {} + Se = real_len = None + backend = fe.vqgan_backend + for _ in range(iters): + times, Se, real_len, backend = _run_flashrt_prefill_once( + fe, prompt, images, cached_ids, use_cuda_graph=use_cuda_graph) + for k, v in times.items(): + stage_values.setdefault(k, []).append(v) + + stage_stats = {k: _stats(v) for k, v in stage_values.items()} + prefill_ms = stage_stats["transformer_prefill_ms"]["p50"] + result = { + "Se": int(Se), + "real_len": int(real_len), + "vqgan_backend": backend, + "fa4_attn": fe.fa4_attn_active, + "reuse_input_ids": bool(reuse_input_ids), + "one_time_input_build_ms": input_build_ms, + "latency_ms": stage_stats["total_ms"], + "stage_breakdown_ms": stage_stats, + "roofline": _roofline(int(Se), prefill_ms, peak_tflops), + } + + if generate_greedy > 0: + for _ in range(max(1, min(warmup, 2))): + fe.generate_greedy(prompt, images, max_new_tokens=generate_greedy) + gen_lat = [] + for _ in range(iters): + t0 = time.perf_counter() + out = fe.generate_greedy(prompt, images, max_new_tokens=generate_greedy) + torch.cuda.synchronize() + gen_lat.append((time.perf_counter() - t0) * 1000.0) + result["generate_greedy"] = { + "max_new_tokens": int(generate_greedy), + "latency_ms": _stats(gen_lat), + "ms_per_token": _stats([x / generate_greedy for x in gen_lat]), + "output_token_count": len(out["input_ids"]), + } + + del fe + torch.cuda.empty_cache() + return result + + +def _bench_hf(checkpoint_dir: pathlib.Path, prompt: str, images, + *, target_size: int, use_trt_vqgan: bool, + trt_vqgan_engine_dir: str | None, + iters: int, warmup: int) -> Dict: + import torch + from transformers import AutoConfig + from safetensors.torch import load_file + + try: + from transformers import ChameleonForConditionalGeneration as _Cls + except (ImportError, ModuleNotFoundError) as e: + print(f"[bench] HF BF16 reference unavailable ({e}); skipping") + return None + + cfg = AutoConfig.from_pretrained(str(checkpoint_dir)) + cfg.rope_scaling = None + if not hasattr(cfg, "rope_theta") or cfg.rope_theta is None: + cfg.rope_theta = 10000.0 + + model = _Cls(cfg) + sd = {} + for shard in sorted(checkpoint_dir.glob("model-*-of-*.safetensors")): + sd.update(load_file(str(shard))) + model.load_state_dict(sd, strict=False, assign=False) + model = model.to(torch.bfloat16).cuda().eval() + + from flash_rt.frontends.torch.chameleon_thor import ChameleonTorchFrontendThor + fe = ChameleonTorchFrontendThor( + str(checkpoint_dir), use_fp8=False, use_cuda_graph=False, + target_size=target_size, use_trt_vqgan=use_trt_vqgan, + trt_vqgan_engine_dir=trt_vqgan_engine_dir) + ids = fe.encode_prompt(prompt, images) + backend = fe.vqgan_backend + del fe + torch.cuda.empty_cache() + + ids_t = torch.tensor([ids], dtype=torch.long, device="cuda") + + def _fwd(): + with torch.no_grad(): + model(input_ids=ids_t, use_cache=False) + torch.cuda.synchronize() + + for _ in range(warmup): + _fwd() + lat = [] + for _ in range(iters): + t0 = time.perf_counter() + _fwd() + lat.append((time.perf_counter() - t0) * 1000.0) + del model + torch.cuda.empty_cache() + return {"Se": len(ids), "vqgan_backend": backend, "latency_ms": _stats(lat)} + + +def main() -> None: + ap = argparse.ArgumentParser(description=__doc__) + ap.add_argument("--checkpoint", required=True) + ap.add_argument("--image-dir", required=True, + help="Directory of real input images") + ap.add_argument("--prompt", default="Describe the image.") + ap.add_argument("--max-images", type=int, default=1) + ap.add_argument("--target-size", type=int, default=512) + ap.add_argument("--use-trt-vqgan", action="store_true", + help="Use TensorRT VQGAN if compatible engines exist " + "(recommended when available; default is eager VQGAN)") + ap.add_argument("--trt-vqgan-engine-dir", default=None) + ap.add_argument("--iters", type=int, default=10) + ap.add_argument("--warmup", type=int, default=2) + ap.add_argument("--no-graph", action="store_true", + help="Disable CUDA Graph capture for FlashRT paths") + ap.add_argument("--reuse-input-ids", action="store_true", + help="Build real-image input ids once and benchmark transformer prefill only") + ap.add_argument("--stage-breakdown", action="store_true", + help="Include per-stage timing in JSON output (currently always collected)") + ap.add_argument("--generate-greedy", type=int, default=0, + help="Also benchmark full-prefix greedy generation for N new tokens") + ap.add_argument("--peak-tflops", type=float, default=240.0, + help="Measured Thor FP8 GEMM plateau used for roofline efficiency") + ap.add_argument("--skip-hf", action="store_true") + ap.add_argument("--output", default="/tmp/chameleon_thor_bench.json") + args = ap.parse_args() + + import torch + + checkpoint_dir = pathlib.Path(args.checkpoint) + image_dir = pathlib.Path(args.image_dir) + images, image_paths = _load_real_images(image_dir, args.max_images) + device_name = torch.cuda.get_device_name(0) + + result: Dict = { + "checkpoint": str(checkpoint_dir), + "image_dir": str(image_dir), + "image_paths": image_paths, + "prompt": args.prompt, + "num_images": len(images), + "target_size": args.target_size, + "use_trt_vqgan": bool(args.use_trt_vqgan), + "trt_vqgan_engine_dir": args.trt_vqgan_engine_dir, + "vqgan_backend_requested": "trt" if args.use_trt_vqgan else "eager", + "device": device_name, + "iters": args.iters, + "warmup": args.warmup, + "graph": not args.no_graph, + "reuse_input_ids": bool(args.reuse_input_ids), + "stage_breakdown": bool(args.stage_breakdown), + "generate_greedy": int(args.generate_greedy), + "peak_tflops": float(args.peak_tflops), + } + + print("[bench] FlashRT FP16...") + result["flashrt_fp16"] = _bench_flashrt( + checkpoint_dir, args.prompt, images, use_fp8=False, + use_cuda_graph=not args.no_graph, target_size=args.target_size, + use_trt_vqgan=args.use_trt_vqgan, + trt_vqgan_engine_dir=args.trt_vqgan_engine_dir, + iters=args.iters, warmup=args.warmup, + reuse_input_ids=args.reuse_input_ids, + generate_greedy=args.generate_greedy, + peak_tflops=args.peak_tflops) + print(f"[bench] FlashRT FP16: {result['flashrt_fp16']}") + + print("[bench] FlashRT dynamic FP8...") + result["flashrt_fp8"] = _bench_flashrt( + checkpoint_dir, args.prompt, images, use_fp8=True, + use_cuda_graph=not args.no_graph, target_size=args.target_size, + use_trt_vqgan=args.use_trt_vqgan, + trt_vqgan_engine_dir=args.trt_vqgan_engine_dir, + iters=args.iters, warmup=args.warmup, + reuse_input_ids=args.reuse_input_ids, + generate_greedy=args.generate_greedy, + peak_tflops=args.peak_tflops) + print(f"[bench] FlashRT FP8: {result['flashrt_fp8']}") + + if not args.skip_hf: + print("[bench] HF BF16 (eager)...") + result["hf_bf16"] = _bench_hf( + checkpoint_dir, args.prompt, images, + target_size=args.target_size, + use_trt_vqgan=args.use_trt_vqgan, + trt_vqgan_engine_dir=args.trt_vqgan_engine_dir, + iters=args.iters, warmup=args.warmup) + print(f"[bench] HF BF16: {result['hf_bf16']}") + + with open(args.output, "w") as f: + json.dump(result, f, indent=2) + print(f"[bench] wrote {args.output}") + + +if __name__ == "__main__": + main() diff --git a/scripts/build_vqgan_trt.py b/scripts/build_vqgan_trt.py new file mode 100644 index 000000000..a4729d54d --- /dev/null +++ b/scripts/build_vqgan_trt.py @@ -0,0 +1,316 @@ +#!/usr/bin/env python3 +"""Build TensorRT FP16 engines for the Chameleon VQ-GAN encoder. + +Exports fixed-shape ONNX per resolution, then compiles TRT engines. +Engines are cached at ~/.flash_rt/trt_engines/vqgan/ (or --output_dir). + +Must be run on the TARGET hardware (engines are not portable across GPUs). + +Usage +----- +python scripts/build_vqgan_trt.py \ + --checkpoint /path/to/Chameleon_7B_mGPT \ + --resolutions 384x512 384x384 384x672 512x512 \ + --verify +""" + +import argparse +import json +import os +import platform +import sys +import time +from pathlib import Path + +import numpy as np +import torch +import torch.nn as nn +import tensorrt as trt + +from flash_rt.models.chameleon.vqvae_hf import ( + load_chameleon_vqvae, + vqvae_checkpoint_digest, +) + + +class VQGANEncoderWrapper(nn.Module): + """Image tensor -> codebook indices. + + Input : x of shape (B, 3, H, W), float32 in [-1, 1] + Output : indices of shape (B, H/16, W/16), int64 + """ + + def __init__(self, vq_model: nn.Module): + super().__init__() + self.encoder = vq_model.encoder + self.quant_conv = vq_model.quant_conv + # Re-host the codebook as a self-contained nn.Embedding so this + # wrapper is a fully standard nn.Module (no external parameter sharing). + n_e, e_dim = vq_model.quantize.embedding.weight.shape + self.codebook = nn.Embedding(n_e, e_dim) + with torch.no_grad(): + self.codebook.weight.copy_(vq_model.quantize.embedding.weight) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + h = self.encoder(x) # (B, 2*z_channels, H/16, W/16) when double_z=True + h = self.quant_conv(h) # (B, e_dim, H/16, W/16) + + b, c, hh, ww = h.shape + # (B, e_dim, H/16, W/16) -> (B, H/16, W/16, e_dim) -> (B*N, e_dim) + z_flat = h.permute(0, 2, 3, 1).contiguous().view(-1, c) + e = self.codebook.weight # (n_e, e_dim) + + # ||z - e||^2 = ||z||^2 + ||e||^2 - 2 z·e + d = ( + (z_flat * z_flat).sum(dim=1, keepdim=True) + + (e * e).sum(dim=1) + - 2.0 * torch.matmul(z_flat, e.t()) + ) # (B*N, n_e) + idx = torch.argmin(d, dim=1) # (B*N,) + idx = idx.view(b, hh, ww).to(torch.int64) + return idx + + +def build_vqmodel(checkpoint: str, device: torch.device) -> nn.Module: + vq_model, _ = load_chameleon_vqvae( + checkpoint, device=device, dtype=torch.float32) + for p in vq_model.parameters(): + p.requires_grad_(False) + return vq_model + + +def parse_resolution(s: str) -> tuple: + h, w = s.lower().split("x") + return int(h), int(w) + + +def export_onnx(vq_model, height, width, batch, opset, output_path, device): + wrapper = VQGANEncoderWrapper(vq_model).to(device).eval() + dummy = torch.randn(batch, 3, height, width, device=device, dtype=torch.float32) + export_kwargs = dict( + input_names=["image"], + output_names=["indices"], + dynamic_axes=None, + opset_version=opset, + do_constant_folding=True, + ) + # PyTorch 2.5+ defaults to dynamo=True which emits TRT-incompatible + # IR-10/opset-18 nodes. Force legacy exporter on those versions. + # On older PyTorch (< 2.5) the kwarg doesn't exist and isn't needed. + _torch_ver = tuple(int(x) for x in torch.__version__.split(".")[:2]) + if _torch_ver >= (2, 5): + export_kwargs["dynamo"] = False + torch.onnx.export(wrapper, dummy, output_path, **export_kwargs) + print(f" [onnx] exported {output_path} (shape=[{batch},3,{height},{width}])") + return wrapper + + +def build_engine(onnx_path, engine_path, workspace_gb, opt_level, fp16=True): + logger = trt.Logger(trt.Logger.WARNING) + builder = trt.Builder(logger) + + # Choose network creation flags. TRT 11+ supports STRONGLY_TYPED which + # makes the engine honor the ONNX's native dtypes verbatim (so FP16 + # weights stay FP16). Older TRT uses EXPLICIT_BATCH + BuilderFlag.FP16 + # which silently demotes to FP32 for some ops on Ada (we observed this + # on TRT 11.1 — the engine ran ~2× slower in FP32 by default). + use_strong = (fp16 and + hasattr(trt.NetworkDefinitionCreationFlag, "STRONGLY_TYPED")) + + if use_strong: + # Strongly typed mode requires the ONNX itself to be FP16. + # Auto-convert from the FP32 ONNX (idempotent: skips if already FP16). + from onnxconverter_common import float16 + import onnx as _onnx_mod + + onnx_path_obj = Path(onnx_path) + fp16_onnx = onnx_path_obj.with_suffix(".fp16.onnx") + if not fp16_onnx.exists(): + print(f" [onnx-fp16] converting {onnx_path_obj.name} → " + f"{fp16_onnx.name}") + mdl = _onnx_mod.load(str(onnx_path_obj)) + mdl16 = float16.convert_float_to_float16(mdl, keep_io_types=True) + _onnx_mod.save(mdl16, str(fp16_onnx)) + else: + print(f" [onnx-fp16] reusing existing {fp16_onnx.name}") + parse_path = str(fp16_onnx) + network = builder.create_network( + 1 << int(trt.NetworkDefinitionCreationFlag.STRONGLY_TYPED)) + elif hasattr(trt.NetworkDefinitionCreationFlag, "EXPLICIT_BATCH"): + # Legacy path (TRT < 10). + parse_path = str(onnx_path) + network = builder.create_network( + 1 << int(trt.NetworkDefinitionCreationFlag.EXPLICIT_BATCH)) + else: + # Very old TRT — default network. + parse_path = str(onnx_path) + network = builder.create_network(0) + + parser = trt.OnnxParser(network, logger) + with open(parse_path, "rb") as f: + if not parser.parse(f.read()): + for i in range(parser.num_errors): + print(f" [trt] parse error: {parser.get_error(i)}") + raise RuntimeError(f"Failed to parse ONNX: {parse_path}") + + config = builder.create_builder_config() + config.set_memory_pool_limit(trt.MemoryPoolType.WORKSPACE, + int(workspace_gb * (1 << 30))) + if fp16 and not use_strong and hasattr(trt.BuilderFlag, "FP16"): + # Old TRT path that needs the explicit FP16 flag. + config.set_flag(trt.BuilderFlag.FP16) + config.builder_optimization_level = opt_level + + t0 = time.time() + serialized = builder.build_serialized_network(network, config) + elapsed = time.time() - t0 + + if serialized is None: + raise RuntimeError(f"TRT engine build failed for {onnx_path}") + + with open(engine_path, "wb") as f: + f.write(serialized) + size_mb = os.path.getsize(engine_path) / (1024 * 1024) + print(f" [trt] built {engine_path} ({size_mb:.1f} MB, {elapsed:.1f}s)") + return engine_path + + +def verify_engine(engine_path, torch_wrapper, height, width, batch, device): + logger = trt.Logger(trt.Logger.WARNING) + runtime = trt.Runtime(logger) + with open(engine_path, "rb") as f: + engine = runtime.deserialize_cuda_engine(f.read()) + if engine is None: + print(" [verify] FAILED: could not deserialize engine") + return False + + context = engine.create_execution_context() + stream = torch.cuda.current_stream() + + h_lat, w_lat = height // 16, width // 16 + inp = torch.randn(batch, 3, height, width, device=device, dtype=torch.float32) + out = torch.zeros(batch, h_lat, w_lat, device=device, dtype=torch.int64) + + # TRT may use int32 for output; detect binding dtype + out_name = "indices" + out_dtype_trt = engine.get_tensor_dtype(out_name) + if out_dtype_trt == trt.DataType.INT32: + out_buf = torch.zeros(batch, h_lat, w_lat, device=device, dtype=torch.int32) + else: + out_buf = out + + context.set_tensor_address("image", inp.data_ptr()) + context.set_tensor_address("indices", out_buf.data_ptr()) + context.execute_async_v3(stream_handle=stream.cuda_stream) + stream.synchronize() + + if out_dtype_trt == trt.DataType.INT32: + trt_indices = out_buf.to(torch.int64) + else: + trt_indices = out_buf + + with torch.no_grad(): + pt_indices = torch_wrapper(inp) + + match = (trt_indices == pt_indices).sum().item() + total = trt_indices.numel() + mismatch_pct = 100.0 * (1.0 - match / total) + ok = mismatch_pct < 0.5 + print(f" [verify] match={match}/{total} ({100*match/total:.2f}%), " + f"mismatch={mismatch_pct:.3f}% {'OK' if ok else 'WARN'}") + return ok + + +def main(): + parser = argparse.ArgumentParser(description="Build TRT engines for VQ-GAN encoder") + parser.add_argument("--checkpoint", type=str, required=True, + help="Path to a Transformers Chameleon checkpoint") + parser.add_argument("--resolutions", nargs="+", default=["384x512", "384x384", "384x672", "512x512"], + help="HxW resolutions to build engines for") + parser.add_argument("--batch", type=int, default=1) + parser.add_argument("--output_dir", type=str, + default=str(Path.home() / ".flash_rt" / "trt_engines" / "vqgan")) + parser.add_argument("--workspace_gb", type=float, default=2.0) + parser.add_argument("--opt_level", type=int, default=5, + help="TRT builder optimization level (0-5)") + parser.add_argument("--opset", type=int, default=17) + parser.add_argument("--keep_onnx", action="store_true", + help="Keep intermediate ONNX files in output_dir") + parser.add_argument("--verify", action="store_true", + help="Run parity check TRT vs PyTorch after build") + args = parser.parse_args() + + device = torch.device("cuda") + output_dir = Path(args.output_dir) + output_dir.mkdir(parents=True, exist_ok=True) + + print(f"TensorRT {trt.__version__}") + print(f"Platform: {platform.machine()}") + print(f"Output dir: {output_dir}") + print(f"Resolutions: {args.resolutions}") + print() + + vq_model = build_vqmodel(args.checkpoint, device) + ckpt_hash = vqvae_checkpoint_digest(args.checkpoint) + + manifest_path = output_dir / "manifest.json" + if manifest_path.exists(): + with open(manifest_path) as f: + manifest = json.load(f) + manifest["build_date"] = time.strftime("%Y-%m-%dT%H:%M:%S") + else: + manifest = { + "trt_version": trt.__version__, + "platform": platform.machine(), + "ckpt_hash": ckpt_hash, + "build_date": time.strftime("%Y-%m-%dT%H:%M:%S"), + "batch": args.batch, + "precision": "fp16", + "engines": {}, + } + + for res_str in args.resolutions: + height, width = parse_resolution(res_str) + assert height % 16 == 0 and width % 16 == 0, f"H/W must be multiples of 16, got {res_str}" + h_lat, w_lat = height // 16, width // 16 + + print(f"── {res_str} ({height}×{width} → {h_lat}×{w_lat} latent) ──") + + engine_name = f"vqgan_encoder_b{args.batch}_{height}x{width}_fp16.engine" + engine_path = output_dir / engine_name + + # Export ONNX + onnx_path = output_dir / f"vqgan_encoder_{height}x{width}.onnx" + wrapper = export_onnx(vq_model, height, width, args.batch, args.opset, + str(onnx_path), device) + + # Build TRT engine + build_engine(str(onnx_path), str(engine_path), args.workspace_gb, args.opt_level) + + # Verify + if args.verify: + verify_engine(str(engine_path), wrapper.to(device), height, width, args.batch, device) + + # Clean ONNX + if not args.keep_onnx: + onnx_path.unlink(missing_ok=True) + + manifest["engines"][res_str] = { + "file": engine_name, + "height": height, + "width": width, + "input_shape": [args.batch, 3, height, width], + "output_shape": [args.batch, h_lat, w_lat], + } + print() + + # Write manifest + manifest_path = output_dir / "manifest.json" + with open(manifest_path, "w") as f: + json.dump(manifest, f, indent=2) + print(f"Manifest written: {manifest_path}") + print("Done.") + + +if __name__ == "__main__": + main() diff --git a/scripts/chameleon_orin_check.py b/scripts/chameleon_orin_check.py new file mode 100644 index 000000000..71246c58f --- /dev/null +++ b/scripts/chameleon_orin_check.py @@ -0,0 +1,393 @@ +#!/usr/bin/env python3 +"""Gate-1 correctness harness for Chameleon-7B on Orin SM87. + +Compares the FlashRT INT8/INT4 frontend against a **stock transformers 4.57.1** +``ChameleonForConditionalGeneration`` reference (bf16, eager attention) on the +*same* token ids, and runs the CUDA-Graph safety gate on the decode body. + +Two non-obvious things this handles: + +* **The checkpoint does not load into stock transformers as-is.** Its + ``ChameleonLayerNorm`` builds ``(num_heads, head_dim) = (32,128)`` weights + while this Lumina-mGPT export stores ``(1,128)`` — the shard is + ``model_parallel_size`` x ``head_dim`` and upstream expands it with + ``repeat_interleave`` at forward time. We expand it at load instead, which is + exactly equivalent for the 7B (mp=1) layout. transformers 4.57 also rejects + ``from_pretrained(..., state_dict=...)``, so the model is built with a naked + constructor + ``load_state_dict``. + +* **VQ-GAN index drift would poison every number.** FlashRT runs the encoder + convs in fp16, so codebook indices can differ from an all-fp32 reference. The + FlashRT side therefore *exports* the ids it computed and the reference is fed + those verbatim, isolating LLM error from tokenizer error. Run with + ``--vq-fp16-argmin`` to measure the drift itself instead. + +Usage: + PYTHONPATH=. python scripts/chameleon_orin_check.py \ + --checkpoint /path/to/Chameleon_7B_mGPT \ + --image FlashRT.png --prompt "Describe this image." --steps 16 + ... --int4 # QuaRot W4A4 tier + ... --text-only # skip the image (fast smoke) +""" + +from __future__ import annotations + +import argparse +import sys + +import torch + +PROBE_LAYERS = [0, 4, 8, 12, 16, 20, 24, 28, 31] + +# Chameleon suppresses the 8192 image-codebook ids at every forward, so they +# carry no information and must be excluded from any similarity metric — +# including them makes cosine NaN (finfo(bf16).min squared overflows fp32). +IMG_LO, IMG_HI = 4, 8196 + + +def cosine(a: torch.Tensor, b: torch.Tensor) -> float: + a = a.detach().float().flatten() + b = b.detach().float().flatten() + return float(a @ b / (a.norm() * b.norm() + 1e-30)) + + +def text_slice(logits: torch.Tensor) -> torch.Tensor: + """Drop the masked image-id band before comparing logits.""" + return torch.cat([logits[..., :IMG_LO], logits[..., IMG_HI:]], dim=-1) + + +# ====================================================================== +# Reference +# ====================================================================== + +def load_reference(ckpt: str): + """Stock transformers Chameleon, bf16, eager, with qk_norm expanded.""" + import json + from pathlib import Path + from safetensors.torch import load_file + from transformers import ChameleonConfig, ChameleonForConditionalGeneration + + ckpt_p = Path(ckpt) + index = json.loads((ckpt_p / "model.safetensors.index.json").read_text()) + sd, n_exp = {}, 0 + for shard in sorted(set(index["weight_map"].values())): + full = load_file(str(ckpt_p / shard)) + for k, t in full.items(): + if ((".q_norm." in k or ".k_norm." in k) + and t.dim() == 2 and t.shape[0] == 1): + t = t.repeat_interleave(32, dim=0) + n_exp += 1 + sd[k] = t + del full + if n_exp != 128: + print(f" [warn] expanded {n_exp} qk_norm tensors, expected 128") + + cfg = ChameleonConfig.from_pretrained(ckpt) + cfg._attn_implementation = "eager" + torch.set_default_dtype(torch.bfloat16) + model = ChameleonForConditionalGeneration(cfg) + torch.set_default_dtype(torch.float32) + missing, unexpected = model.load_state_dict(sd, strict=False) + missing = [m for m in missing if "inv_freq" not in m] + if missing or unexpected: + raise RuntimeError(f"ref load: missing={missing[:4]} unexpected={unexpected[:4]}") + del sd + return model.eval().cuda() + + +@torch.no_grad() +def reference_forward(model, ids: list): + """Teacher-forced forward. Returns (per-layer hidden states, logits). + + Hidden states come from forward hooks on the decoder layers, NOT from + ``output_hidden_states=True``: HF's ``all_hidden_states`` replaces the last + entry with the *post-final-norm* tensor, so comparing it against a + pre-norm probe shows a false cosine collapse. + """ + caught = {} + handles = [] + for li in PROBE_LAYERS: + def hook(_m, _inp, out, li=li): + caught[li] = (out[0] if isinstance(out, tuple) else out)[0].float().cpu() + return None # a non-None hook return REPLACES the output + handles.append(model.model.layers[li].register_forward_hook(hook)) + + def norm_hook(_m, _inp, out): + caught["final_norm"] = out[0].float().cpu() + return None + handles.append(model.model.norm.register_forward_hook(norm_hook)) + try: + t = torch.tensor([ids], device="cuda") + logits = model(input_ids=t).logits[0].float().cpu() + finally: + for h in handles: + h.remove() + return caught, logits + + +# ====================================================================== +# Gates +# ====================================================================== + +def gate_graph_safety(front) -> bool: + """Capture the decode body and prove it is not frozen (stale-value test).""" + print("\n── graph-safety gate (decode body) ──") + pos = front.S + tok_a, tok_b = 16853, 40000 + try: + front.decode_step(tok_a, pos=pos) # warm every M=1 shape first + front.decode_step(tok_a, pos=pos) + torch.cuda.synchronize() + g = torch.cuda.CUDAGraph() + s = torch.cuda.Stream() + with torch.cuda.stream(s): + with torch.cuda.graph(g, stream=s): + # The kernels MUST be launched on the capture stream; on stream 0 + # they are silently not recorded and the replay looks frozen. + front.decode_step(front._tok_dev, pos=pos, + stream=int(s.cuda_stream)) + print(" capture : OK (no code=13)") + except Exception as e: # pragma: no cover + print(f" capture : FAIL — {type(e).__name__}: {e}") + return False + + front._tok_dev.fill_(tok_a); g.replay(); torch.cuda.synchronize() + la = front._logits.clone() + front._tok_dev.fill_(tok_b); g.replay(); torch.cuda.synchronize() + lb = front._logits.clone() + c = cosine(text_slice(la), text_slice(lb)) + frozen = torch.equal(la, lb) + print(f" stale-value : {'FAIL (frozen)' if frozen else 'PASS'} " + f"(cos between two seed tokens = {c:.4f})") + return not frozen + + +def gate_overflow(front) -> bool: + """FP16 residual health. + + The gate is **finiteness**, not an absolute magnitude. Chameleon's L31 + residual legitimately reaches ~2.6e5 in the bf16 reference (the massive + activation), so any absolute threshold below that would fail by + construction. What must not happen is inf/nan, which the + ``ffn_down_clamp`` prevents by capping the down output just under FP16's + 65504. Saturation at the clamp is therefore *expected* at L31 and is + reported for information only. + """ + print("\n── fp16 residual health (finite + clamp saturation) ──") + snaps = front.snapshot_probe() + if not snaps: + print(" (frontend built without probe_layers — skipped)") + return True + clamp = float(getattr(front, "ffn_down_clamp", 0.0) or 0.0) + bad, sat = [], [] + worst, worst_k = 0.0, "" + for k, v in snaps.items(): + if not bool(torch.isfinite(v).all()): + bad.append(k) + m = float(v[torch.isfinite(v)].abs().max()) if v.numel() else 0.0 + if clamp and m >= clamp * 0.98: + sat.append(k) + if m > worst: + worst, worst_k = m, k + print(f" max finite |x| : {worst:.0f} at {worst_k} " + f"(fp16 max 65504, clamp {clamp:.0f})") + print(f" saturating at clamp: {sat if sat else 'none'} " + f"{'(expected at L31)' if sat else ''}") + ok = not bad + print(f" inf/nan : {bad if bad else 'none'} -> " + f"{'PASS' if ok else 'FAIL'}") + return ok + + +def main() -> int: + ap = argparse.ArgumentParser() + ap.add_argument("--checkpoint", required=True) + ap.add_argument("--image", default=None) + ap.add_argument("--prompt", default="Describe this image.") + ap.add_argument("--steps", type=int, default=16, help="greedy tokens to compare") + ap.add_argument("--max-seq", type=int, default=1280) + ap.add_argument("--int4", action="store_true") + ap.add_argument("--int4-down", action="store_true") + ap.add_argument("--text-only", action="store_true") + ap.add_argument("--vq-fp16-argmin", action="store_true", + help="measure VQ index drift instead of avoiding it") + ap.add_argument("--split-kv-bias", type=int, default=4) + ap.add_argument("--skip-ref", action="store_true", + help="run only the FlashRT-side gates") + args = ap.parse_args() + + from flash_rt.frontends.torch.chameleon_rtx_sm87 import ChameleonTorchFrontendRtxSm87 + + print("=" * 68) + print("Chameleon-7B Orin SM87 — Gate 1") + print("=" * 68) + + front = ChameleonTorchFrontendRtxSm87( + args.checkpoint, max_seq=args.max_seq, + use_int4=args.int4, use_int4_down=args.int4_down, + split_kv_bias=args.split_kv_bias, + vq_argmin_fp32=not args.vq_fp16_argmin, + probe_layers=PROBE_LAYERS) + print(f"tier={front.precision_tier} spec={front.precision_spec()}") + + # ---- prompt ---- + images = None + text = args.prompt + if not args.text_only: + if args.image: + from PIL import Image + images = [Image.open(args.image).convert("RGB")] + else: + import numpy as np + from PIL import Image + images = [Image.fromarray(np.random.RandomState(0).randint( + 0, 256, (480, 640, 3), dtype=np.uint8))] + print(" [note] no --image given; using deterministic noise") + text = "" + args.prompt + front.set_prompt(text, images=images) + ids = front.input_ids.tolist() + print(f"\nISL={len(ids)} (1 BOS + n_img*1026 + text + 1 sep) " + f"images={front.timing['n_images']} " + f"prompt_ms={front.timing['prompt_ms']:.0f}") + + # ---- FlashRT teacher-forced logits + probes ---- + lg_frt = front.prefill(logits_all=True).float().cpu() + probes = front.snapshot_probe() + print(f"prefill_ms={front.timing['prefill_ms']:.0f}") + + ok_overflow = gate_overflow(front) + ok_graph = gate_graph_safety(front) + + # ---- greedy text ---- + front.set_prompt(text, images=images) + frt_ids = front.generate(max_new_tokens=args.steps, return_ids=True) + frt_txt = front.processor.tokenizer.decode(frt_ids, skip_special_tokens=True) + tm = front.timing + print(f"\n── generation ──\n FlashRT ids : {frt_ids}" + f"\n FlashRT text: {frt_txt!r}" + f"\n decode : {tm['decode_ms_per_token']:.2f} ms/token " + f"= {tm['decode_tok_s']:.2f} tok/s (ISL={len(ids)}, OSL={args.steps})") + + if args.skip_ref: + print("\n(--skip-ref: reference comparison not run)") + return 0 if (ok_graph and ok_overflow) else 1 + + # ---- reference ---- + print(f"\n── HF reference (bf16 eager) ──") + ref = load_reference(args.checkpoint) + ref_h, ref_lg = reference_forward(ref, ids) + + print("\n layer cosine norm-ratio FlashRT|max| ref|max|") + worst = 1.0 + for li in PROBE_LAYERS: + a = probes[f"layer_{li}"].float().cpu() + b = ref_h[li] + c = cosine(a, b) + worst = min(worst, c) + print(f" L{li:<5d} {c:.6f} {float(a.norm()/b.norm()):.4f} " + f"{float(a.abs().max()):9.1f} {float(b.abs().max()):9.1f}") + c_fn = cosine(probes["final_norm"].float().cpu(), ref_h["final_norm"]) + print(f" final {c_fn:.6f}") + + # logits + argmax over the whole teacher-forced sequence + a_lg = text_slice(lg_frt) + b_lg = text_slice(ref_lg) + c_last = cosine(a_lg[-1], b_lg[-1]) + am_frt = a_lg.argmax(-1) + am_ref = b_lg.argmax(-1) + exact = am_frt == am_ref + + # A BF16 tie is not a precision failure: when the reference's top-1 and + # top-2 are within one BF16 ULP the winner is numerically arbitrary, and + # any engine may legitimately pick either. Classify those separately + # instead of scoring them as errors. + top2 = b_lg.topk(2, dim=-1).values + gap = (top2[:, 0] - top2[:, 1]).abs() + ulp = top2[:, 0].abs() * 2 ** -8 # BF16 has 8 mantissa bits + tied = gap <= ulp + real_bad = (~exact) & (~tied) + n = len(am_ref) + match_exact = float(exact.float().mean()) + match_adj = float((exact | tied).float().mean()) + print(f"\n last-row logit cosine : {c_last:.6f}") + print(f" argmax exact match : {match_exact*100:.2f}% " + f"({int(exact.sum())}/{n})") + print(f" of {int((~exact).sum())} mismatches: {int(((~exact) & tied).sum())} " + f"are BF16 ties (gap <= 1 ulp), {int(real_bad.sum())} are real") + print(f" tie-adjusted match : {match_adj*100:.2f}%") + if int(real_bad.sum()): + g = gap[real_bad] + print(f" real-mismatch ref gap : median={float(g.median()):.4f} " + f"max={float(g.max()):.4f} (logit scale " + f"~{float(top2[:, 0].abs().median()):.1f})") + match = match_adj + + # Split by position class. At an *image* position the model predicts the + # next token while all 8192 image ids are masked out of the logits, so the + # winner is an arbitrary low-confidence text token — averaging over the 1024 + # image positions swamps the handful that actually drive generation. Gate on + # the text positions only. + ids_t = torch.tensor(ids) + is_img = ((ids_t >= IMG_LO) & (ids_t < IMG_HI)) | (ids_t == 8197) | (ids_t == 8196) + for label, sel in (("image", is_img), ("text ", ~is_img)): + k = int(sel.sum()) + if not k: + continue + e = float(exact[sel].float().mean()) + a = float((exact | tied)[sel].float().mean()) + print(f" {label} positions ({k:4d}) : exact {e*100:6.2f}% " + f"tie-adjusted {a*100:6.2f}% " + f"median ref gap {float(gap[sel].median()):.3f}") + text_sel = ~is_img + if int(text_sel.sum()): + match = float((exact | tied)[text_sel].float().mean()) + + # greedy text identity + with torch.no_grad(): + gen = ref.generate(input_ids=torch.tensor([ids], device="cuda"), + max_new_tokens=args.steps, do_sample=False, + num_beams=1) + ref_new = gen[0, len(ids):].tolist() + ref_txt = front.processor.tokenizer.decode(ref_new, skip_special_tokens=True) + n_pref = 0 + for x, y in zip(frt_ids, ref_new): + if x != y: + break + n_pref += 1 + print(f"\n reference ids : {ref_new}") + print(f" reference text: {ref_txt!r}") + print(f" identical prefix: {n_pref}/{min(len(frt_ids), len(ref_new))} tokens") + + # ---- verdict ---- + cos_gate = 0.99 if front.use_int4 else 0.97 + n_text = int(text_sel.sum()) + checks = [ + ("worst layer cosine", worst >= cos_gate, f"{worst:.4f} >= {cos_gate}"), + ("last-row logit cosine", c_last >= 0.999, f"{c_last:.6f} >= 0.999"), + ("greedy text identical", n_pref == len(ref_new), + f"{n_pref}/{len(ref_new)}"), + ("graph safety", ok_graph, ""), + ("fp16 residual finite", ok_overflow, ""), + ] + # Only binding with a meaningful sample: an image-heavy prompt leaves a + # handful of text positions, where one near-tie flip swings the rate by + # >15 points. Greedy text identity above is the metric that actually + # tracks generation quality. + if n_text >= 32: + checks.insert(2, ("argmax match (text)", match >= 0.99, + f"{match*100:.2f}% >= 99% (n={n_text})")) + else: + print(f"\n note: only {n_text} text positions — the argmax gate is " + f"reported as informational, not binding " + f"({match*100:.2f}% tie-adjusted)") + print("\n" + "=" * 68) + for name, ok, detail in checks: + print(f" [{'PASS' if ok else 'FAIL'}] {name:24s} {detail}") + allok = all(c[1] for c in checks) + print(f"\nGATE 1: {'PASS' if allok else 'FAIL'}") + print("=" * 68) + return 0 if allok else 1 + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/scripts/check_chameleon_thor_precision.py b/scripts/check_chameleon_thor_precision.py new file mode 100644 index 000000000..6840df158 --- /dev/null +++ b/scripts/check_chameleon_thor_precision.py @@ -0,0 +1,204 @@ +#!/usr/bin/env python3 +"""Real-image precision gate for standalone Chameleon-7B on Thor. + +Compares: + * FlashRT FP16 vs HF BF16 last-token logits (cosine, top-k overlap, + greedy next-token equality) + * FlashRT dynamic FP8 vs FlashRT FP16 last-token logits (same metrics) + * optional final-hidden cosine + +Inputs are always real images (from a user-supplied directory), never +synthetic token ids, per the standalone Chameleon-7B optimization plan. + +Usage +----- + PYTHONPATH=. python scripts/check_chameleon_thor_precision.py \\ + --checkpoint /path/to/Chameleon_7B_mGPT \\ + --image-dir /path/to/images \\ + --prompt "Describe the image." \\ + --output /tmp/chameleon_thor_precision.json +""" + +from __future__ import annotations + +import argparse +import json +import pathlib +import sys + +import numpy as np + + +def _load_hf_bf16(checkpoint_dir: pathlib.Path): + """Load the plain HF ChameleonForConditionalGeneration at bf16. + + HF ``from_pretrained`` silently mis-loads q/k norm weights on this + checkpoint's old [1,128] shape; use a naked model + manual + ``load_state_dict`` instead. Requires the model's reference + ``modeling_chameleon`` implementation (optional; ``--skip-hf`` + skips this comparison). + """ + import torch + from transformers import AutoConfig + from transformers import ChameleonForConditionalGeneration as _Cls + from safetensors.torch import load_file + + cfg = AutoConfig.from_pretrained(str(checkpoint_dir)) + cfg.rope_scaling = None + if not hasattr(cfg, "rope_theta") or cfg.rope_theta is None: + cfg.rope_theta = 10000.0 + + model = _Cls(cfg) + sd = {} + for shard in sorted(checkpoint_dir.glob("model-*-of-*.safetensors")): + sd.update(load_file(str(shard))) + missing, unexpected = model.load_state_dict(sd, strict=False, assign=False) + non_vq_missing = [k for k in missing if "vqmodel" not in k] + if non_vq_missing: + print(f"[_load_hf_bf16] WARNING: {len(non_vq_missing)} " + f"non-VQVAE keys missing from ckpt") + if unexpected: + print(f"[_load_hf_bf16] WARNING: {len(unexpected)} unexpected keys") + model = model.to(torch.bfloat16).cuda().eval() + return model + + +def _hf_last_logits(model, input_ids: list[int]) -> np.ndarray: + import torch + ids = torch.tensor([input_ids], dtype=torch.long, device="cuda") + with torch.no_grad(): + out = model(input_ids=ids, use_cache=False) + logits = out.logits[0, -1].float().cpu().numpy() + return logits + + +def _load_real_images(image_dir: pathlib.Path, max_images: int): + from PIL import Image + exts = (".jpg", ".jpeg", ".png", ".bmp") + paths = sorted(p for p in image_dir.iterdir() if p.suffix.lower() in exts) + if not paths: + raise FileNotFoundError(f"No real images found under {image_dir}") + paths = paths[:max_images] + images = [Image.open(p).convert("RGB") for p in paths] + return images, [str(p) for p in paths] + + +def _cosine(a: np.ndarray, b: np.ndarray) -> float: + a = a.astype(np.float64).ravel() + b = b.astype(np.float64).ravel() + denom = (np.linalg.norm(a) * np.linalg.norm(b)) + if denom == 0: + return 0.0 + return float(np.dot(a, b) / denom) + + +def _topk_overlap(a: np.ndarray, b: np.ndarray, k: int = 10) -> float: + top_a = set(np.argsort(-a)[:k].tolist()) + top_b = set(np.argsort(-b)[:k].tolist()) + return len(top_a & top_b) / float(k) + + +def main() -> None: + ap = argparse.ArgumentParser(description=__doc__) + ap.add_argument("--checkpoint", required=True) + ap.add_argument("--image-dir", required=True, + help="Directory of real input images") + ap.add_argument("--prompt", default="Describe the image.") + ap.add_argument("--max-images", type=int, default=1) + ap.add_argument("--target-size", type=int, default=512) + ap.add_argument("--use-trt-vqgan", action="store_true", + help="Use TensorRT VQGAN if compatible engines exist " + "(recommended when available; default is eager VQGAN)") + ap.add_argument("--trt-vqgan-engine-dir", default=None) + ap.add_argument("--topk", type=int, default=10) + ap.add_argument("--output", default="/tmp/chameleon_thor_precision.json") + ap.add_argument("--skip-hf", action="store_true", + help="Skip HF BF16 comparison (FP16 vs FP8 only)") + args = ap.parse_args() + + checkpoint_dir = pathlib.Path(args.checkpoint) + image_dir = pathlib.Path(args.image_dir) + images, image_paths = _load_real_images(image_dir, args.max_images) + print(f"[check] loaded {len(images)} real image(s) from {image_dir}: " + f"{image_paths}") + + from flash_rt.frontends.torch.chameleon_thor import ChameleonTorchFrontendThor + + result: dict = { + "checkpoint": str(checkpoint_dir), + "image_dir": str(image_dir), + "image_paths": image_paths, + "prompt": args.prompt, + "target_size": args.target_size, + "use_trt_vqgan": bool(args.use_trt_vqgan), + "trt_vqgan_engine_dir": args.trt_vqgan_engine_dir, + "vqgan_backend_requested": "trt" if args.use_trt_vqgan else "eager", + } + + print("[check] running FlashRT FP16 reference path...") + fe_fp16 = ChameleonTorchFrontendThor( + str(checkpoint_dir), use_fp8=False, use_cuda_graph=False, + target_size=args.target_size, + use_trt_vqgan=args.use_trt_vqgan, + trt_vqgan_engine_dir=args.trt_vqgan_engine_dir) + out_fp16 = fe_fp16.prefill(args.prompt, images) + logits_fp16 = out_fp16["logits"].numpy().ravel() + hidden_fp16 = out_fp16["hidden"].numpy() + ids_fp16 = out_fp16["input_ids"] + result["vqgan_backend_actual_fp16"] = out_fp16.get("vqgan_backend") + del fe_fp16 + import torch + torch.cuda.empty_cache() + + print("[check] running FlashRT dynamic-FP8 path...") + fe_fp8 = ChameleonTorchFrontendThor( + str(checkpoint_dir), use_fp8=True, use_cuda_graph=False, + target_size=args.target_size, + use_trt_vqgan=args.use_trt_vqgan, + trt_vqgan_engine_dir=args.trt_vqgan_engine_dir) + out_fp8 = fe_fp8.prefill(args.prompt, images) + logits_fp8 = out_fp8["logits"].numpy().ravel() + hidden_fp8 = out_fp8["hidden"].numpy() + result["vqgan_backend_actual_fp8"] = out_fp8.get("vqgan_backend") + del fe_fp8 + torch.cuda.empty_cache() + + fp8_vs_fp16 = { + "logits_cosine": _cosine(logits_fp8, logits_fp16), + "topk_overlap": _topk_overlap(logits_fp8, logits_fp16, args.topk), + "greedy_token_match": bool( + int(np.argmax(logits_fp8)) == int(np.argmax(logits_fp16))), + "hidden_cosine": _cosine(hidden_fp8, hidden_fp16), + } + result["flashrt_fp8_vs_flashrt_fp16"] = fp8_vs_fp16 + print(f"[check] FlashRT FP8 vs FP16: {fp8_vs_fp16}") + + if not args.skip_hf: + print("[check] loading HF BF16 model (this may take a while)...") + try: + hf_model = _load_hf_bf16(checkpoint_dir) + except (ImportError, ModuleNotFoundError) as e: + print(f"[check] HF BF16 reference unavailable ({e}); " + f"rerun with --skip-hf for the FP16-vs-FP8 check only") + hf_model = None + if hf_model is not None: + logits_hf = _hf_last_logits(hf_model, ids_fp16) + del hf_model + torch.cuda.empty_cache() + + fp16_vs_hf = { + "logits_cosine": _cosine(logits_fp16, logits_hf), + "topk_overlap": _topk_overlap(logits_fp16, logits_hf, args.topk), + "greedy_token_match": bool( + int(np.argmax(logits_fp16)) == int(np.argmax(logits_hf))), + } + result["flashrt_fp16_vs_hf_bf16"] = fp16_vs_hf + print(f"[check] FlashRT FP16 vs HF BF16: {fp16_vs_hf}") + + with open(args.output, "w") as f: + json.dump(result, f, indent=2) + print(f"[check] wrote {args.output}") + + +if __name__ == "__main__": + main() diff --git a/scripts/profile_chameleon_thor.py b/scripts/profile_chameleon_thor.py new file mode 100644 index 000000000..cbeaace65 --- /dev/null +++ b/scripts/profile_chameleon_thor.py @@ -0,0 +1,116 @@ +#!/usr/bin/env python3 +"""Nsight Systems profiling helper for standalone Chameleon Thor. + +Use with CUDA profiler capture range, for example: + + nsys profile --force-overwrite=true \ + -o /tmp/chameleon_prefill_only_512 \ + --capture-range=cudaProfilerApi -t cuda,nvtx \ + env PYTHONPATH=. python scripts/profile_chameleon_thor.py \ + --checkpoint /path/to/Chameleon_7B_mGPT \ + --image-dir /path/to/images \ + --target-size 512 --reuse-input-ids --iters 5 +""" + +from __future__ import annotations + +import argparse +import pathlib + + +def _load_real_images(image_dir: pathlib.Path, max_images: int): + from PIL import Image + exts = (".jpg", ".jpeg", ".png", ".bmp") + paths = sorted(p for p in image_dir.iterdir() if p.suffix.lower() in exts) + if not paths: + raise FileNotFoundError(f"No real images found under {image_dir}") + paths = paths[:max_images] + return [Image.open(p).convert("RGB") for p in paths] + + +def _pad_ids(input_ids: list[int], pad_id: int = 1) -> tuple[list[int], int]: + real_len = len(input_ids) + padded = list(input_ids) + rem = len(padded) % 16 + if rem: + padded.extend([pad_id] * (16 - rem)) + return padded, real_len + + +def _run_prefill_body(fe, prompt: str, images, cached_ids, *, use_graph: bool): + import torch + + if cached_ids is None: + ids = fe.encode_prompt(prompt, images) + else: + ids = cached_ids + padded, real_len = _pad_ids(ids) + fe._real_len = real_len + fe.Se = len(padded) + fe._last_input_ids = padded + if fe._use_autotune: + fe._autotune_gemms(fe.Se) + fe._embed_ids(padded) + if use_graph: + fe._capture_graph(fe.Se) + fe._infer_graph.replay() + else: + fe._run_backbone(fe.Se) + fe._project_last() + torch.cuda.synchronize() + + +def main() -> None: + ap = argparse.ArgumentParser(description=__doc__) + ap.add_argument("--checkpoint", required=True) + ap.add_argument("--image-dir", required=True, + help="Directory of real input images") + ap.add_argument("--prompt", default="Describe the image.") + ap.add_argument("--max-images", type=int, default=1) + ap.add_argument("--target-size", type=int, default=512) + ap.add_argument("--use-trt-vqgan", action="store_true") + ap.add_argument("--trt-vqgan-engine-dir", default=None) + ap.add_argument("--use-fp16", action="store_true") + ap.add_argument("--no-graph", action="store_true") + ap.add_argument("--reuse-input-ids", action="store_true") + ap.add_argument("--warmup", type=int, default=3) + ap.add_argument("--iters", type=int, default=5) + args = ap.parse_args() + + import torch + from flash_rt.frontends.torch.chameleon_thor import ChameleonTorchFrontendThor + + images = _load_real_images(pathlib.Path(args.image_dir), args.max_images) + fe = ChameleonTorchFrontendThor( + args.checkpoint, + use_fp8=not args.use_fp16, + use_cuda_graph=not args.no_graph, + target_size=args.target_size, + use_trt_vqgan=args.use_trt_vqgan, + trt_vqgan_engine_dir=args.trt_vqgan_engine_dir, + ) + cached_ids = fe.encode_prompt(args.prompt, images) if args.reuse_input_ids else None + + for _ in range(args.warmup): + _run_prefill_body(fe, args.prompt, images, cached_ids, use_graph=not args.no_graph) + + torch.cuda.synchronize() + torch.cuda.cudart().cudaProfilerStart() + for _ in range(args.iters): + _run_prefill_body(fe, args.prompt, images, cached_ids, use_graph=not args.no_graph) + torch.cuda.synchronize() + torch.cuda.cudart().cudaProfilerStop() + + print({ + "Se": fe.Se, + "real_len": fe._real_len, + "vqgan_backend": fe.vqgan_backend, + "use_fp8": not args.use_fp16, + "graph": not args.no_graph, + "reuse_input_ids": args.reuse_input_ids, + "iters": args.iters, + }) + + +if __name__ == "__main__": + main() diff --git a/tests/test_chameleon_contracts.py b/tests/test_chameleon_contracts.py new file mode 100644 index 000000000..81501366e --- /dev/null +++ b/tests/test_chameleon_contracts.py @@ -0,0 +1,245 @@ +"""Chameleon frontend contract tests — no checkpoint required. + +Covers registry/lazy-import behavior, hardware fail-fast gates, the +Thor prompt pad-to-16 capacity boundary, the Orin generation-parameter +boundary, and the load_model(config="chameleon") redirect. GPU-heavy +eager-vs-graph consistency lives in the precision scripts +(scripts/check_chameleon_thor_precision.py) which need a checkpoint. +""" + +import importlib +import os +from pathlib import Path + +import pytest + +torch = pytest.importorskip("torch") + +from flash_rt.hardware import _PIPELINE_MAP, resolve_pipeline_class + +REPO_ROOT = Path(__file__).resolve().parents[1] +_OPTIONAL_BUILD_MODULES = { + "flash_rt.flash_rt_kernels", + "flash_rt.flash_rt_fa2", +} + + +def _chameleon_build_enabled() -> bool: + build_dir = Path(os.environ.get("FLASHRT_BUILD_DIR", REPO_ROOT / "build")) + cache = build_dir / "CMakeCache.txt" + if not cache.is_file(): + return False + return any( + line.startswith("FLASHRT_ENABLE_CHAMELEON:BOOL=") + and line.rsplit("=", 1)[-1].upper() in {"ON", "TRUE", "1"} + for line in cache.read_text(errors="replace").splitlines() + ) + + +def _load_frontend(module_name: str, class_name: str): + try: + module = importlib.import_module(module_name) + except ModuleNotFoundError as exc: + if _chameleon_build_enabled() or exc.name not in _OPTIONAL_BUILD_MODULES: + raise + return None + return getattr(module, class_name) + + +ChameleonTorchFrontendThor = _load_frontend( + "flash_rt.frontends.torch.chameleon_thor", "ChameleonTorchFrontendThor") +ChameleonTorchFrontendRtxSm87 = _load_frontend( + "flash_rt.frontends.torch.chameleon_rtx_sm87", + "ChameleonTorchFrontendRtxSm87", +) +_THOR_IMPORT = ChameleonTorchFrontendThor is not None +_ORIN_IMPORT = ChameleonTorchFrontendRtxSm87 is not None +PAD_ID = 1 + +needs_thor = pytest.mark.skipif(not _THOR_IMPORT, + reason="chameleon_thor frontend not importable") +needs_orin = pytest.mark.skipif(not _ORIN_IMPORT, + reason="chameleon_rtx_sm87 frontend not importable") + + +# ---------------------------------------------------------------- registry + +def test_registry_maps_to_expected_frontends(): + assert _PIPELINE_MAP[("chameleon", "torch", "thor")] == ( + "flash_rt.frontends.torch.chameleon_thor", + "ChameleonTorchFrontendThor") + assert _PIPELINE_MAP[("chameleon", "torch", "rtx_sm87")] == ( + "flash_rt.frontends.torch.chameleon_rtx_sm87", + "ChameleonTorchFrontendRtxSm87") + + +def test_registry_entries_are_lazy_module_strings(): + for key in (("chameleon", "torch", "thor"), + ("chameleon", "torch", "rtx_sm87")): + mod, cls_name = _PIPELINE_MAP[key] + assert isinstance(mod, str) and isinstance(cls_name, str) + + +def test_frontend_import_link_errors_are_not_swallowed(monkeypatch): + def fail_import(_): + raise ImportError("undefined symbol: chameleon_test_symbol") + + monkeypatch.setattr(importlib, "import_module", fail_import) + with pytest.raises(ImportError, match="undefined symbol"): + _load_frontend("flash_rt.frontends.torch.chameleon_thor", "Frontend") + + +@needs_thor +def test_resolve_thor_pipeline_class(): + cls = resolve_pipeline_class("chameleon", "torch", "thor") + assert cls.__module__ == "flash_rt.frontends.torch.chameleon_thor" + assert cls.__name__ == "ChameleonTorchFrontendThor" + + +@needs_orin +def test_resolve_orin_pipeline_class(): + cls = resolve_pipeline_class("chameleon", "torch", "rtx_sm87") + assert cls.__module__ == "flash_rt.frontends.torch.chameleon_rtx_sm87" + assert cls.__name__ == "ChameleonTorchFrontendRtxSm87" + + +def test_sm87_allowlist_rejects_unsupported_config(): + with pytest.raises(RuntimeError, match="SM87"): + resolve_pipeline_class("groot_n17", "torch", "rtx_sm87") + + +def test_load_model_chameleon_redirects_with_clear_error(): + pytest.importorskip("numpy") + import flash_rt + with pytest.raises(Exception, match="not served through"): + flash_rt.load_model("/nonexistent/fake-ckpt", config="chameleon", + framework="torch", hardware="thor") + + +# ------------------------------------------------------- Thor hardware gate + +def _thor_probe(): + obj = object.__new__(ChameleonTorchFrontendThor) + return ChameleonTorchFrontendThor._require_arch(obj) + + +@needs_thor +def test_thor_rejects_when_cuda_unavailable(monkeypatch): + monkeypatch.delenv("FLASHRT_CHAMELEON_THOR_FORCE", raising=False) + monkeypatch.setattr(torch.cuda, "is_available", lambda: False) + with pytest.raises(RuntimeError, match="CUDA is not available"): + _thor_probe() + + +@needs_thor +def test_thor_rejects_wrong_capability(monkeypatch): + monkeypatch.delenv("FLASHRT_CHAMELEON_THOR_FORCE", raising=False) + monkeypatch.setattr(torch.cuda, "is_available", lambda: True) + monkeypatch.setattr(torch.cuda, "get_device_capability", lambda *_: (8, 7)) + with pytest.raises(RuntimeError, match="targets SM110"): + _thor_probe() + + +@needs_thor +def test_thor_accepts_sm110(monkeypatch): + monkeypatch.delenv("FLASHRT_CHAMELEON_THOR_FORCE", raising=False) + monkeypatch.setattr(torch.cuda, "is_available", lambda: True) + monkeypatch.setattr(torch.cuda, "get_device_capability", lambda *_: (11, 0)) + _thor_probe() # must not raise + + +@needs_thor +def test_thor_documented_env_override_skips_probe(monkeypatch): + monkeypatch.setenv("FLASHRT_CHAMELEON_THOR_FORCE", "1") + # No CUDA mocking: the override must return before touching torch.cuda. + _thor_probe() + + +@needs_thor +@pytest.mark.parametrize("max_seq", [-16, -1, 0, 15]) +def test_thor_rejects_too_small_max_seq(monkeypatch, tmp_path, max_seq): + monkeypatch.setenv("FLASHRT_CHAMELEON_THOR_FORCE", "1") + with pytest.raises(ValueError, match="max_seq must be at least 16"): + ChameleonTorchFrontendThor(tmp_path, max_seq=max_seq) + + +# ----------------------------------------------- Thor prompt padding bounds + +def _bare_thor(se_max): + fe = object.__new__(ChameleonTorchFrontendThor) + fe._Se_max = se_max + fe._use_autotune = False + fe._use_cuda_graph = False + fe._embed_ids = lambda ids: None + return fe + + +@needs_thor +def test_prompt_padding_rejects_overshoot_on_nonaligned_capacity(): + # 30-token prompt pads to 32, which exceeds a 31-token capacity. + fe = _bare_thor(31) + fe.encode_prompt = lambda text, images: list(range(30)) + with pytest.raises(ValueError, match="padded sequence length"): + fe.set_prompt("x") + + +@needs_thor +def test_prompt_padding_accepts_prompt_within_capacity(): + fe = _bare_thor(32) + fe.encode_prompt = lambda text, images: list(range(30)) + ids = fe.set_prompt("x") + assert fe._real_len == 30 + assert fe.Se == 32 and len(ids) == 32 + assert ids[30:] == [PAD_ID, PAD_ID] + + +@needs_thor +def test_prompt_padding_accepts_exact_multiple_of_16(): + fe = _bare_thor(32) + fe.encode_prompt = lambda text, images: list(range(32)) + ids = fe.set_prompt("x") + assert fe.Se == 32 and len(ids) == 32 and fe._real_len == 32 + + +# ------------------------------------------------------ Orin hardware gate + +@needs_orin +def test_orin_rejects_wrong_capability_before_checkpoint_work(monkeypatch): + monkeypatch.delenv("FLASHRT_CHAMELEON_SM87_FORCE", raising=False) + monkeypatch.setattr(torch.cuda, "get_device_capability", lambda *_: (11, 0)) + with pytest.raises(RuntimeError, match="targets SM87"): + ChameleonTorchFrontendRtxSm87("/nonexistent/fake-ckpt") + + +@needs_orin +def test_orin_env_override_bypasses_arch_gate(monkeypatch): + monkeypatch.setenv("FLASHRT_CHAMELEON_SM87_FORCE", "1") + monkeypatch.setattr(torch.cuda, "get_device_capability", lambda *_: (11, 0)) + # The gate must be skipped; construction then fails later for an + # unrelated reason (missing checkpoint), never the SM87 gate error. + with pytest.raises(Exception) as excinfo: + ChameleonTorchFrontendRtxSm87("/nonexistent/fake-ckpt") + assert "targets SM87" not in str(excinfo.value) + + +# --------------------------------------------- Orin generation-param bounds + +@needs_orin +def _bare_orin(): + fe = object.__new__(ChameleonTorchFrontendRtxSm87) + fe._prompt_ready = True + fe.S = 16 + fe.max_seq = 4096 + return fe + + +@needs_orin +def test_generate_negative_max_new_tokens_raises(): + with pytest.raises(ValueError, match="max_new_tokens"): + _bare_orin().generate(max_new_tokens=-1) + + +@needs_orin +def test_generate_zero_max_new_tokens_returns_empty(): + assert _bare_orin().generate(max_new_tokens=0) == "" + assert _bare_orin().generate(max_new_tokens=0, return_ids=True) == [] diff --git a/tests/test_chameleon_thor_vqgan_backend.py b/tests/test_chameleon_thor_vqgan_backend.py new file mode 100644 index 000000000..48711812e --- /dev/null +++ b/tests/test_chameleon_thor_vqgan_backend.py @@ -0,0 +1,32 @@ +from __future__ import annotations + +import inspect +import os +from pathlib import Path + +import pytest + +try: + from flash_rt.frontends.torch.chameleon_thor import ChameleonTorchFrontendThor +except ModuleNotFoundError as exc: + root = Path(__file__).resolve().parents[1] + build_dir = Path(os.environ.get("FLASHRT_BUILD_DIR", root / "build")) + cache = build_dir / "CMakeCache.txt" + chameleon_enabled = cache.is_file() and ( + "FLASHRT_ENABLE_CHAMELEON:BOOL=ON" in cache.read_text(errors="replace") + ) + if chameleon_enabled or exc.name != "flash_rt.flash_rt_kernels": + raise + pytest.skip("flash_rt_kernels is not built", allow_module_level=True) + + +def test_chameleon_trt_vqgan_is_opt_in_by_default(): + sig = inspect.signature(ChameleonTorchFrontendThor.__init__) + assert sig.parameters["use_trt_vqgan"].default is False + + +def test_chameleon_fa4_attn_is_opt_in_by_default(): + sig = inspect.signature(ChameleonTorchFrontendThor.__init__) + assert sig.parameters["use_fa4_attn"].default is None + os.environ.pop("FLASHRT_CHAMELEON_FA4_ATTN", None) + assert bool(os.environ.get("FLASHRT_CHAMELEON_FA4_ATTN", "0") in ("1", "true", "on")) is False diff --git a/tests/test_chameleon_vqvae_hf.py b/tests/test_chameleon_vqvae_hf.py new file mode 100644 index 000000000..cd04f8478 --- /dev/null +++ b/tests/test_chameleon_vqvae_hf.py @@ -0,0 +1,84 @@ +"""Tests for the Apache-2.0 Transformers Chameleon VQ-VAE adapter.""" + +from __future__ import annotations + +from pathlib import Path + +import pytest + +torch = pytest.importorskip("torch") +pytest.importorskip("safetensors") +pytest.importorskip("transformers") + +from safetensors.torch import save_file +from transformers import ChameleonConfig, ChameleonVQVAE +from transformers.models.chameleon.configuration_chameleon import ( + ChameleonVQVAEConfig, +) + +from flash_rt.models.chameleon.vqvae_hf import ( + encode_vqvae_tokens, + load_chameleon_vqvae, + vqvae_checkpoint_digest, +) + + +def test_loads_only_prefixed_vqvae_weights_strictly(tmp_path): + vq_config = ChameleonVQVAEConfig( + embed_dim=32, + num_embeddings=16, + latent_channels=32, + resolution=16, + base_channels=32, + channel_multiplier=[1], + num_res_blocks=1, + attn_resolutions=[], + ) + config = ChameleonConfig( + hidden_size=32, + intermediate_size=64, + num_hidden_layers=1, + num_attention_heads=4, + num_key_value_heads=4, + vq_config=vq_config, + vocabulary_map={}, + ) + config.save_pretrained(tmp_path) + reference = ChameleonVQVAE(vq_config) + state = { + f"model.vqmodel.{key}": value + for key, value in reference.state_dict().items() + } + save_file(state, tmp_path / "model.safetensors") + + loaded, mapping = load_chameleon_vqvae( + tmp_path, device="cpu", dtype=torch.float32) + + assert type(mapping).__name__ == "ChameleonImageVocabularyMapping" + for key, value in reference.state_dict().items(): + assert torch.equal(value, loaded.state_dict()[key]) + assert vqvae_checkpoint_digest(tmp_path) == vqvae_checkpoint_digest(tmp_path) + + +def test_no_research_license_vqgan_is_vendored(): + root = Path(__file__).resolve().parents[1] + vendored = root / "flash_rt/models/chameleon/vqgan" + assert not vendored.exists() or not any(vendored.rglob("*")) + + +def test_encode_tokens_supports_old_and_new_transformers_outputs(): + expected = torch.tensor([1, 2, 3]) + + class OldModel: + def encode(self, _): + return None, None, expected + + class NewOutput: + image_tokens = expected + + class NewModel: + def encode(self, _): + return NewOutput() + + assert encode_vqvae_tokens(OldModel(), None) is expected + assert encode_vqvae_tokens(NewModel(), None) is expected