feat(csrc): Chameleon-7B kernel layer — fused quant, INT8 GEMMs, FA2 causal attention - #165
feat(csrc): Chameleon-7B kernel layer — fused quant, INT8 GEMMs, FA2 causal attention#165DXICM wants to merge 7 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 Chameleon kernel layer. The new kernels are generally well organized and named, the implementation comments are clear, and the intended Thor and Orin paths are easy to understand. Codex performed a strict review based on FlashRT's long-term maintenance and incremental-build principles. Before merging, we recommend addressing the following items. 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 feature direction is sound; the main request is to close the model build boundary and make the public kernel contracts safe before these interfaces become part of the long-lived common surface. |
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).
|
Thanks for the detailed review. All four required changes are addressed, and re-validated on a production Jetson Orin (sm_87, CUDA 12.2, torch 2.3.0). 1. Model-level build isolation —
2. 3. No 4. Expanded validation — committed test files, production Orin results:
Combined production run with the HyVLA suites: 43 passed / 0 failed / 5 skipped (all skips are intended hardware gates). Checklist items requiring SM110 hardware (Thor FP16/FP8 causal FMHA coverage, Chameleon ON + SM110 build/import/symbol checks) are queued for the Thor device run; they will be reported here once available. |
2 participants
Summary
Adds the csrc kernel layer required by the Chameleon-7B model integration (model frontends land in a follow-up PR based on this branch).
New kernels:
quantize_fp8_device_fp16,awq_quant_fp8_static_fp16) and INT8 rowwise norms (rms_norm_quantize_dynamic_fp8_fp16,residual_add_rms_norm_quantize_dynamic_fp8_fp16,gate_geglu_quantize_dynamic_fp8_fp16) incsrc/kernels/cutlass_sm80_int8_rowwise_fp16out*.cu,cutlass_sm80_int4_rowwise.cu) withGemmRunnerFP8 FP16-out dispatchfht_int4.cu) — residual+rms_norm+rotate+quantize in one passqk_norm_rope_fused.cu) for Chameleon's per-head Q/K LayerNorm layoutflash_fwd_hdim128_fp16_sm80_causal.cu) and Thor CUTLASS causal FMHA shared libs (fmha_fp16_causal.cu,fmha_fp8_causal.cu)CMake/binding compliance:
#ifdef ENABLE_SM80_INT8_CUTLASS … #else throw #endif; sources compile under the same gateflash_rt_kernelstargetENABLE_SM100_CUTLASS, ctypes/dlopen-loaded (no pybind)Validation:
tests/test_chameleon_thor_fused_kernels.py: fused-vs-unfused bitwise equality regression (no checkpoint required)flash_rt_kernelsimports verified unaffectedWhat is not changed
flash_rt_kernelssymbol surface is additive onlyTest plan
cmake -B build -S . -DGPU_ARCH=87 && cmake --build build -j$(nproc) --target flash_rt_kernels(SM87)cmake -B build -S . -DGPU_ARCH=110 && cmake --build build -j$(nproc) --target flash_rt_kernels(SM110)python -m pytest tests/test_chameleon_thor_fused_kernels.py -qpython -c "from flash_rt import flash_rt_kernels; print(flash_rt_kernels.__file__)"on both archs