feat(chameleon): Chameleon-7B Thor FP8 and Orin SM87 INT8/QuaRot frontends - #166
feat(chameleon): Chameleon-7B Thor FP8 and Orin SM87 INT8/QuaRot frontends#166DXICM wants to merge 13 commits into
Conversation
…T4 GEMMs Generic FP16-backbone kernels for the upcoming Chameleon-7B paths: - Fused dynamic-FP8 quantization (graph-replay safe, host-scale free): rms_norm_quantize, gate_geglu_quantize and residual_add_rms_norm_quantize_dynamic_fp8_fp16, plus quantize_int8_rowwise_fp16. - INT8 rowwise norms: residual_add_rms_norm_fp16, rms_norm_int8_rowwise_fp16, residual_add_rms_norm_int8_rowwise_fp16. - clamp_inplace_fp16 (FP16 overflow guard for late FFN down-projections), qk_norm_rope_fused_fp16 (per-head QK-LayerNorm + RoPE in one pass) and awq_quant_fp8_static_fp16 (AWQ per-channel activation quantization). - SM80 CUTLASS GEMMs under ENABLE_SM80_INT8_CUTLASS: INT8 rowwise FP16-out (base + T64x128/T256x128 tile variants) and INT4 rowwise, plus the radix-16 FHT kernels (fht_int4.cu) used by the QuaRot-Hadamard tier. Output row stride == N is a contract: Orin KV-cache writes rely on it. - GemmRunner: FP8_NN_DEV_FP16 (=6, FP8_NT_DEV=5 already taken), fp8_nn_dev_fp16 and the autotune_fp8_nn_dev_fp16 / autotune_fp8_nn_bias entry points. All kernels are bound unconditionally (SM80 GEMMs behind the existing ENABLE_SM80_INT8_CUTLASS inline-ifdef guards) and are exercised by checkpoint-free bit-exact tests in a follow-up commit.
Attention backends for the Chameleon-7B frontends:
- FA2: fp16 hdim128 sm80 causal forward + split-KV instantiations wired
into the FA2_HDIMS/FA2_DTYPES matrix, a new fvk_attention_fa2_fwd_fp16_causal
entry point in fa2_wrapper_causal.cu (FA2_HAS_FP16 && FA2_HAS_HDIM_128
guarded, stubbed otherwise) and the fwd_fp16_causal pybind binding. The
existing bf16 causal wrapper gains the same FA2_HAS_BF16 guard so an
fp16-only slim matrix still links.
- Thor SM110: libfmha_fp16_causal.so and libfmha_fp8_causal.so shared
targets inside ENABLE_SM100_CUTLASS, mirroring the fmha_fp16_strided
target (same output dir, "${GPU_ARCH}a" archs, install rules). The FP16
library exports fmha_fp16_causal and the bottom-right-aligned
fmha_fp16_causal_br used by incremental KV-cache decode, where a
top-left causal mask would be silently wrong.
Checkpoint-free correctness tests for the new fused kernels: - test_chameleon_thor_fused_kernels.py: bit-exact comparison of the three fused dynamic-FP8 quantize kernels against their unfused kernel sequences (runs on any built flash_rt_kernels). - test_fp4_chameleon_layer16.py: NVFP4 FFN tier microbenchmarks for the L31 overflow layer; skips cleanly when flash_rt_fp4 (sm_120+ gated build) is not importable.
- Move test_fp4_chameleon_layer16.py out of tests/ (it is a benchmark script with no test_ functions; will land in benchmarks/ with the model PR) - Drop dangling docs/chameleon7b_rtx_sm87.md reference from fht_int4.cu (that doc ships with the model PR, not this kernels PR) - Replace internal roadmap language in awq_quant_fp8_static_fp16.cu header with neutral technical description
|
Thank you for contributing the complete Chameleon-7B Thor FP8 and Orin SM87 frontends. The responsibilities of the model layer, pipelines, and hardware backends are generally separated clearly, and we did not find an existing model being redirected to the Chameleon runtime. Codex reviewed this PR against FlashRT's long-term maintenance standards. We recommend addressing the following items after #165 has been corrected and merged. Required changes
Pre-merge checklist
Please also refer to the repository's PR Review Checklist and Adding a New Model guides. This is a Codex-assisted maintainability review. The architecture is moving in a reasonable direction; the dependency order, model build isolation, input boundaries, and third-party distribution requirements need to be completed before these hardware frontends are suitable for long-term support on |
Address review feedback on build boundary and public-kernel safety: Build isolation (FLASHRT_ENABLE_CHAMELEON, OFF by default): - New CMake option gates all Chameleon-specific TUs, libraries, and symbols together: QK Norm/RoPE, AWQ FP16 quant, SM80 INT8/INT4 rowwise GEMM fp16-out + FHT/QuaRot, FA2 FP16 causal instances, and the SM100/110 causal FMHA shared libraries. - Chameleon-specific kernel definitions inside common norm.cu / quantize.cu are wrapped in #ifdef FLASHRT_ENABLE_CHAMELEON; the corresponding extern declarations and m.def blocks in bindings.cpp use the same guard (combined with ENABLE_SM80_INT8_CUTLASS / FLASHRT_HAVE_MOTUS_VAE_FP8 where those gates already applied). - Model-neutral fp16 norm/quant/activation helpers (residual_add_rms_ norm_fp16, *_quantize_dynamic_fp8_fp16, clamp_inplace_fp16) remain in the common layer with model-neutral docstrings. - Preprocessor simulation with the option OFF confirms zero gated symbols survive; with ON, all 18 gated bindings are active. Kernel contracts: - qk_norm_rope_fused_fp16 now enforces dim==128 (the only shape the RoPE writeback fully covers) and validates seq_len>0, num_heads>0, eps>0, raising py::value_error instead of risking partial output. - fa2_wrapper_causal.cu no longer calls std::abort() from any Python-reachable path: all 7 unsupported-shape / not-compiled sites now throw std::runtime_error (surfaced as Python RuntimeError).
Adds the reference and contract tests requested in review: - test_qk_norm_rope_fused.py: torch reference for per-head LayerNorm + rotate-half RoPE at dim=128, plus the dim/seq_len/eps contract (invalid inputs raise ValueError before launch). - test_sm80_int8_int4_gemm_fht.py: INT8 rowwise fp16-out GEMM vs dequant reference at production shapes, INT4 variant, and FHT/QuaRot norm-preservation checks. - test_fa2_fp16_causal.py: prefill vs torch SDPA causal reference, q_len=1 decode vs full-row softmax reference, and a causality leak check (perturbing the last key must not change earlier rows). - test_thor_causal_fmha.py: ctypes-loaded libfmha_fp16_causal.so vs torch SDPA reference (MHA and GQA shapes). - test_gemm_runner_dispatch.py: fp8_nn_dev_fp16 device-descale path vs dequant reference and autotune-cached re-run equivalence. All tests skip cleanly without CUDA, the built module, or the FLASHRT_ENABLE_CHAMELEON-gated symbols. Also updates the stale "abort" wording in the fa2_bindings.cpp causal docstring.
Production validation on Jetson Orin (SM87, CUDA 12.2) found 3 failing tests; all were test bugs, not kernel regressions (build matrix 10/10, INT8 rowwise GEMM and FHT numerics passed): - test_int4_rowwise_fp16out: the kernel is W4A4 — A must also be packed s4 with per-row /7.0 scales, not int8. The test now quantizes and packs both operands with the production layout (even index in the low nibble, cutlass::int4b_t order) and asserts err==0 instead of skip-on-error. - test_q_len_1_decode: the reference used a [B,S,NH,HD] batched matmul that broadcast incorrectly (512 vs 131072 elements); replaced with explicit einsum score/reference computation. - test_gemm_runner_dispatch: cuBLASLt FP8 matmul requires sm_89+ tensor cores (CUBLAS_STATUS_NOT_SUPPORTED on sm_87); the module now skips cleanly below capability (8, 9).
Standalone Chameleon-7B (image+text) prefill/decode frontend for Jetson AGX Thor: - flash_rt/models/chameleon/pipeline_thor.py: 32-layer Chameleon forward with runtime dynamic per-tensor FP8 (fused quantize kernels), cuBLASLt per-shape autotune, selective L31 ffn_down clamp, optional AWQ V-proj and NVFP4 FFN tiers, CUDA-graph capture with re-embed before replay, and incremental KV-cache decode over fmha_fp16_causal_br. - Vendored Meta Chameleon VQ-GAN tokenizer (flash_rt/models/chameleon/vqgan, Meta Chameleon License headers retained; see the package docstring) with an eager default path and an opt-in TensorRT engine backend (hardware/thor/vqgan_trt_backend.py). - ChameleonTorchFrontendThor (frontends/torch/chameleon_thor.py): checkpoint_dir is a required argument with a clear error when missing; declarative weight spec in _chameleon_thor_spec.py. - hardware/thor/attn_backend_chameleon.py: CUTLASS causal FMHA backend with optional FA4 fast path, loading libfmha_fp16_causal.so from the package directory.
Jetson AGX Orin (SM87) Chameleon-7B path aligned with the upstream rtx_sm87 naming: - flash_rt/models/chameleon/pipeline_rtx.py: one chameleon_forward serving prefill and decode on the SM80 CUTLASS INT8/INT4 rowwise GEMMs with QuaRot-Hadamard rotations (correctness requirement, not an optimization) and the ffn_down clamp on the last 4 layers (FP16 65504 overflow guard). - _chameleon_quant.py: INT8/INT4 weight quantization + Hadamard packing from the BF16 checkpoint. - _chameleon_spec.py: declarative weight spec with an inlined, Chameleon-specific _llm_block (no bias terms, no FP8 scales). - ChameleonTorchFrontendRtxSm87 (chameleon_rtx_sm87.py): set_prompt / prefill / decode_step / generate, FLASHRT_CHAMELEON_SM87_FORCE escape hatch. - hardware/rtx/attn_backend_chameleon.py: FA2 fwd_fp16_causal is mandatory for decode (bottom-right causal semantics); the backend raises rather than falling back to a top-left cuBLAS mask, which would be silently wrong. Runtime numbers (21.07 tok/s, 16/16 bit-identical greedy vs HF BF16) were measured on Orin hardware in the derivative repo and still need SM87 validation here.
- Register ("chameleon", "torch", "thor") and ("chameleon", "torch",
"rtx_sm87") in _PIPELINE_MAP and allow the SM87 key in _SM87_ALLOWED.
- api.load_model redirect for config="chameleon" (chat-style VLM, same
pattern as qwen3_vl): raises NotImplementedError pointing at the two
direct-instantiation frontends.
- tests/test_chameleon_thor_vqgan_backend.py: eager-vs-TRT VQGAN backend
contract test.
- scripts/: bench_chameleon_thor.py, check_chameleon_thor_precision.py,
profile_chameleon_thor.py, chameleon_orin_check.py (Gate-1 harness) and
build_vqgan_trt.py (now driven by the vendored
flash_rt.models.chameleon.vqgan package); HF BF16 reference rows use
transformers' ChameleonForConditionalGeneration directly.
- examples/thor/chameleon_quickstart.py + README entry,
benchmarks/chameleon_thor_latency.py.
- Docs: chameleon_usage.md, chameleon_thor_sm110.md and
chameleon7b_rtx_sm87.md; Chameleon rows in USAGE.md, README.md and
docs/benchmark_comparison.md.
Thor numbers were measured on Jetson AGX Thor (sm_110). All SM87 runtime
numbers in the Orin doc (21.07 tok/s, 16/16 bit-identical greedy vs HF
BF16) come from Orin hardware in the derivative repo and still need SM87
validation in this tree.
- FLASHRT_RYNNVLA2_FP4_LAYERS -> FLASHRT_CHAMELEON_FP4_LAYERS (the env var was inherited from the RynnVLA port with its old name) - replace "001"/"002"/"vendor bf16" comments with plain Chameleon / HF-reference wording in pipeline_thor.py and chameleon_thor.py
- Add #!/usr/bin/env python3 shebangs to 4 scripts and 1 benchmark - Rename _chameleon_spec.py to _chameleon_rtx_sm87_spec.py (hardware suffix per adding_new_model.md convention) and update the import - Add config="chameleon" to docs/stable_api.md (config enum, redirect bullet, resolve_pipeline_class registration) - Translate docs/chameleon_thor_sm110.md from Chinese to English - Remove all internal "derivative repo" / "RynnVLA" provenance references from both engineering docs (42 occurrences) - Move fp4_chameleon_layer16 benchmark from tests/ to benchmarks/ (it has no test_ functions; misfiled in the kernels branch)
…SM110 fail-fast, max_new_tokens contract, VQGAN license, contract tests Address the flashrt-project#166 maintainability review: - Thor prompt-pad boundary: allocation floors capacity to a multiple of 16 and set_prompt validates the PADDED length, so a non-aligned max_seq can never let pad-to-16 overshoot the buffers/KV cache. - Thor hardware gate: ChameleonTorchFrontendThor checks device capability before checkpoint load / CUDA allocation; documented dev override FLASHRT_CHAMELEON_THOR_FORCE=1. - Generation boundary: Thor generate_greedy and Orin generate both reject negative max_new_tokens (ValueError); Orin returns an empty result for zero instead of running prefill and emitting one token. - VQGAN licensing: vendored Meta Chameleon files now carry the full Chameleon Research License (LICENSE) plus a NOTICE recording provenance (incl. the upstream CompVis MIT attribution), the inference-only modifications, and a compatibility notice. Documented in chameleon_usage.md. - Contract tests (tests/test_chameleon_contracts.py): registry + lazy-import, load_model chameleon redirect, Thor/Orin hardware fail-fast, prompt padding bounds, and generation-parameter bounds.
a57db4c to
193797e
Compare
|
Thanks for the review. The branch has been rebased onto the corrected #165 head and all six items are addressed in 193797e; once #165 merges, this PR will be rebased onto 1. Dependency/rebase — rebased onto the corrected #165 head; will rebase onto 2. Thor prompt-padding capacity boundary — capacity is now floored to a multiple of 16 at allocation ( 3. SM110 fail-fast — 4. Orin generation-parameter boundary — 5. VQGAN licensing — the vendored files are from Meta Chameleon ( 6. Contract tests — Production Orin re-test of this branch (build matrix + contract tests) is queued; results will be reported here. |
2 participants
Summary
Adds standalone Chameleon-7B (text + image VLM) to FlashRT on both Jetson targets. Depends on #165 (csrc kernel layer).
Jetson Thor (SM110) — new.
ChameleonTorchFrontendThorwith all-32-layer runtime dynamic per-tensor FP8, cuBLASLt per-shape autotune, L31 selective clamp, fused SwiGLU/quantize/norm kernels, FA4 attention opt-in, and KV-cache incremental decode (30.4 tok/s). VQGAN image tokenizer defaults to eager; TensorRT VQGAN is explicit opt-in (use_trt_vqgan=True). Seedocs/chameleon_thor_sm110.md.Jetson Orin (SM87) — new.
ChameleonTorchFrontendRtxSm87with INT8/QuaRot weight-only quantization and INT8 KV cache. QuaRot Hadamard rotation conditions Chameleon's massive-activation channels for clean INT8 GEMMs. Seedocs/chameleon7b_rtx_sm87.md.Architecture (per
docs/adding_new_model.mdrules 1-4):models/chameleon/pipeline_thor.py+pipeline_rtx.py— one compute path per (model, hardware)frontends/torch/chameleon_thor.py+chameleon_rtx_sm87.py— one frontend per (model, framework, hardware)_PIPELINE_MAP: two entries, strictly one-to-one, no shared classesif arch == ...)load_model(config="chameleon")raisesNotImplementedErrorwith direct-construction instructions (chat-style VLM, not VLApredict()surface; same pattern as Qwen3-VL)Performance (Thor, real image, target_size=512):
Precision: FlashRT FP8 vs FP16 logits cosine 0.99999999, greedy next-token exact match. FP16 vs HF BF16 cosine 0.9999997.
What is not changed
shared_primitives.pyuntouched (model-specific backends in dedicatedattn_backend_chameleon.pyfiles)FLASHRT_CHAMELEON_*prefix, off by defaultdocs/stable_api.mdupdated withconfig="chameleon"entryTest plan
python -m pytest tests/test_chameleon_thor_vqgan_backend.py -qpython -m pytest tests/test_chameleon_thor_fused_kernels.py -q(from feat(csrc): Chameleon-7B kernel layer — fused quant, INT8 GEMMs, FA2 causal attention #165)python examples/thor/chameleon_quickstart.py --checkpoint /path/to/Chameleon_7B_mGPT(Thor)python scripts/chameleon_orin_check.py --checkpoint /path/to/Chameleon_7B_mGPT(Orin SM87)resolve_pipeline_class("chameleon", "torch", "thor")and("chameleon", "torch", "rtx_sm87")resolve correctlyload_model(config="chameleon")raisesNotImplementedErrorwith construction instructions