feat(csrc): Chameleon-7B kernel layer — fused quant, INT8 GEMMs, FA2 causal attention

CMakeLists.txt

@@ -156,6 +156,13 @@ if((FLASHRT_ENABLE_COSMOS3_EDGE OR FLASHRT_ENABLE_COSMOS3_REASONER) AND
     "for non-Thor builds.")
 endif()
 
+# Chameleon-7B model kernels: QK Norm/RoPE, SM87 INT8/INT4 GEMM + FHT/QuaRot,
+# FA2 FP16 causal instances, and SM100/110 causal FMHA libraries. Off by
+# default so unrelated builds pay no compile/link/symbol cost. Model-neutral
+# fp16 norm/quant/activation helpers stay in the common layer.
+option(FLASHRT_ENABLE_CHAMELEON
+       "Build Chameleon-7B model kernels and bindings" OFF)
+
 # Motus beta integration. Motus-specific kernels are additive and must keep
 # their symbols prefixed with ``motus_``. Keeping a build tag lets the public
 # package compile without Motus kernels when debugging unrelated model paths.
@@ -483,6 +490,71 @@ if(ENABLE_SM100_CUTLASS)
   )
   target_link_libraries(fmha_fp16_strided PRIVATE CUDA::cudart)
   message(STATUS "libfmha_fp16_strided.so: building for sm_${GPU_ARCH} (Thor FMHA for SigLIP)")
+
+  if(FLASHRT_ENABLE_CHAMELEON)
+  # ── libfmha_fp16_causal.so — CUTLASS SM100 FP16 causal FMHA for Chameleon ──
+  # Same as fmha_fp16_strided but with CausalMask<true>/<false> and
+  # CausalIndividualTileScheduler. Used by Chameleon-7B LLM self-attention
+  # (is_causal=True). Loaded at runtime via ctypes in
+  # hardware/thor/attn_backend_chameleon.py.
+  add_library(fmha_fp16_causal SHARED csrc/attention/fmha_fp16_causal.cu)
+  set_target_properties(fmha_fp16_causal PROPERTIES
+    LIBRARY_OUTPUT_DIRECTORY ${CMAKE_CURRENT_SOURCE_DIR}/flash_rt
+    RUNTIME_OUTPUT_DIRECTORY ${CMAKE_CURRENT_SOURCE_DIR}/flash_rt
+    CUDA_STANDARD 17
+    POSITION_INDEPENDENT_CODE ON
+    CUDA_ARCHITECTURES "${GPU_ARCH}a"
+    CUDA_RESOLVE_DEVICE_SYMBOLS ON
+    PREFIX "lib"
+    OUTPUT_NAME "fmha_fp16_causal"
+  )
+  target_include_directories(fmha_fp16_causal PRIVATE
+    ${CMAKE_CURRENT_SOURCE_DIR}/csrc/attention
+    ${CUTLASS_DIR}/examples/77_blackwell_fmha
+    ${CUTLASS_INCLUDE}
+    ${CUTLASS_DIR}/tools/util/include
+  )
+  target_compile_options(fmha_fp16_causal PRIVATE
+    $<$<COMPILE_LANGUAGE:CUDA>:
+      --expt-relaxed-constexpr --expt-extended-lambda -O3
+      ${GPU_GENCODE}
+    >
+  )
+  target_link_libraries(fmha_fp16_causal PRIVATE CUDA::cudart)
+  message(STATUS "libfmha_fp16_causal.so: building for sm_${GPU_ARCH} (Thor causal FMHA for Chameleon)")
+
+  # ── libfmha_fp8_causal.so — CUTLASS SM100 FP8 causal FMHA ──
+  # FP8 (E4M3) input variant of fmha_fp16_causal for Chameleon-7B. Inputs
+  # Q/K/V are FP8, outputs O FP16, accumulators FP32. CUTLASS Sm100 FMHA
+  # mainloop has FP8-aware kPRescale logic that triggers when
+  # ``Element == cutlass::float_e4m3_t``. Loaded via dlopen at runtime
+  # alongside libfmha_fp16_causal.so as a drop-in alternative.
+  add_library(fmha_fp8_causal SHARED csrc/attention/fmha_fp8_causal.cu)
+  set_target_properties(fmha_fp8_causal PROPERTIES
+    LIBRARY_OUTPUT_DIRECTORY ${CMAKE_CURRENT_SOURCE_DIR}/flash_rt
+    RUNTIME_OUTPUT_DIRECTORY ${CMAKE_CURRENT_SOURCE_DIR}/flash_rt
+    CUDA_STANDARD 17
+    POSITION_INDEPENDENT_CODE ON
+    CUDA_ARCHITECTURES "${GPU_ARCH}a"
+    CUDA_RESOLVE_DEVICE_SYMBOLS ON
+    PREFIX "lib"
+    OUTPUT_NAME "fmha_fp8_causal"
+  )
+  target_include_directories(fmha_fp8_causal PRIVATE
+    ${CMAKE_CURRENT_SOURCE_DIR}/csrc/attention
+    ${CUTLASS_DIR}/examples/77_blackwell_fmha
+    ${CUTLASS_INCLUDE}
+    ${CUTLASS_DIR}/tools/util/include
+  )
+  target_compile_options(fmha_fp8_causal PRIVATE
+    $<$<COMPILE_LANGUAGE:CUDA>:
+      --expt-relaxed-constexpr --expt-extended-lambda -O3
+      ${GPU_GENCODE}
+    >
+  )
+  target_link_libraries(fmha_fp8_causal PRIVATE CUDA::cudart)
+  message(STATUS "libfmha_fp8_causal.so: building for sm_${GPU_ARCH} (Thor FP8 causal FMHA for Chameleon)")
+  endif()  # FLASHRT_ENABLE_CHAMELEON
 endif()
 
 # ── CUTLASS SM120a block-128 FP8 GEMM (Path B for Qwen3.6) ──
@@ -869,11 +941,18 @@ if(ENABLE_FA2 AND
       csrc/attention/fa2_causal_inst/flash_fwd_split_hdim256_bf16_sm80_causal.cu
     )
   endif()
+  # fp16 hdim=128 serves Chameleon-7B causal prefill/decode on Orin SM87.
+  if(FLASHRT_ENABLE_CHAMELEON AND "128" IN_LIST FA2_HDIMS AND "fp16" IN_LIST FA2_DTYPES)
+    list(APPEND FA2_SRCS
+      csrc/attention/fa2_causal_inst/flash_fwd_hdim128_fp16_sm80_causal.cu
+      csrc/attention/fa2_causal_inst/flash_fwd_split_hdim128_fp16_sm80_causal.cu
+    )
+  endif()
   if(FLASHRT_ENABLE_NATIVE_CPP)
     # The native C boundary has a stable five-symbol surface. Its causal
     # wrapper contains native-only fail-fast dispatch for a slim matrix.
     list(APPEND FA2_SRCS csrc/attention/fa2_wrapper_causal.cu)
-  elseif("bf16" IN_LIST FA2_DTYPES AND
+  elseif(("bf16" IN_LIST FA2_DTYPES OR "fp16" IN_LIST FA2_DTYPES) AND
          ("128" IN_LIST FA2_HDIMS OR "256" IN_LIST FA2_HDIMS))
     # Preserve the existing Python-only source matrix exactly by default.
     list(APPEND FA2_SRCS csrc/attention/fa2_wrapper_causal.cu)
@@ -1582,6 +1661,26 @@ if(ENABLE_SM80_INT8_CUTLASS)
   target_compile_definitions(flash_rt_kernels PRIVATE ENABLE_SM80_INT8_CUTLASS=1)
 endif()
 
+# ── Chameleon-7B model kernels (opt-in, see FLASHRT_ENABLE_CHAMELEON) ──
+# QK Norm/RoPE fused, AWQ FP16 quant, SM80 INT8/INT4 rowwise GEMM fp16-out +
+# FHT/QuaRot rotation. Everything here is Chameleon-specific; model-neutral
+# fp16 norm/quant/activation helpers live in the common layer.
+if(FLASHRT_ENABLE_CHAMELEON)
+  target_sources(flash_rt_kernels PRIVATE
+    csrc/kernels/qk_norm_rope_fused.cu
+    csrc/quantize/awq_quant_fp8_static_fp16.cu)
+  if(ENABLE_SM80_INT8_CUTLASS)
+    target_sources(flash_rt_kernels PRIVATE
+      csrc/gemm/cutlass_sm80_int8_rowwise_fp16out.cu
+      csrc/gemm/cutlass_sm80_int8_rowwise_fp16out_t64x128.cu
+      csrc/gemm/cutlass_sm80_int8_rowwise_fp16out_t256x128.cu
+      csrc/gemm/cutlass_sm80_int4_rowwise.cu
+      csrc/kernels/fht_int4.cu)
+  endif()
+  target_compile_definitions(flash_rt_kernels PRIVATE FLASHRT_ENABLE_CHAMELEON=1)
+  message(STATUS "Chameleon-7B model kernels: ENABLED")
+endif()
+
 # SM120a CUTLASS block-128 FP8 GEMM (Path B for Qwen3.6).
 if(GPU_ARCH STREQUAL "120")
   target_sources(flash_rt_kernels PRIVATE

csrc/attention/fa2_causal_inst/flash_fwd_hdim128_fp16_sm80_causal.cu

@@ -0,0 +1,17 @@
+// FlashRT — FA2 causal instantiation for (fp16, head_dim=128).
+//
+// Sibling of flash_fwd_hdim128_bf16_sm80_causal.cu — adds the fp16
+// specialization needed by the Chameleon-7B (Orin SM87) causal
+// attention path. The vendored launch template already supports
+// Is_causal=true; this file just provides the matching fp16 spec.
+#include "namespace_config.h"
+#include "flash_fwd_launch_template.h"
+
+namespace FLASH_NAMESPACE {
+
+template<>
+void run_mha_fwd_<cutlass::half_t, 128, true>(Flash_fwd_params &params, cudaStream_t stream) {
+    run_mha_fwd_hdim128<cutlass::half_t, true>(params, stream);
+}
+
+}  // namespace FLASH_NAMESPACE

csrc/attention/fa2_causal_inst/flash_fwd_split_hdim128_fp16_sm80_causal.cu

@@ -0,0 +1,13 @@
+// FlashRT — FA2 causal splitkv instantiation for (fp16, head_dim=128).
+//
+// Sibling of flash_fwd_split_hdim128_bf16_sm80_causal.cu — provides
+// the fp16 splitkv dispatch for causal attention. Used by the
+// Chameleon-7B (Orin SM87) path when the splitkv heuristic kicks in.
+#include "namespace_config.h"
+#include "flash_fwd_launch_template.h"
+
+namespace FLASH_NAMESPACE {
+
+template void run_mha_fwd_splitkv_dispatch<cutlass::half_t, 128, true>(Flash_fwd_params &params, cudaStream_t stream);
+
+}  // namespace FLASH_NAMESPACE

csrc/attention/fa2_wrapper.h

@@ -80,6 +80,18 @@ FLASHRT_FA2_NATIVE_API void fvk_attention_fa2_fwd_bf16_causal(
     int o_batch_stride, int o_row_stride, int o_head_stride,
     float softmax_scale, int num_sms, cudaStream_t stream);
 
+FLASHRT_FA2_NATIVE_API void fvk_attention_fa2_fwd_fp16_causal(
+    const void* q_ptr, const void* k_ptr, const void* v_ptr,
+    void* o_ptr, void* softmax_lse_ptr,
+    void* softmax_lse_accum_ptr, void* o_accum_ptr,
+    int batch, int seqlen_q, int seqlen_k,
+    int num_heads_q, int num_heads_kv, int head_dim,
+    int q_batch_stride, int q_row_stride, int q_head_stride,
+    int k_batch_stride, int k_row_stride, int k_head_stride,
+    int v_batch_stride, int v_row_stride, int v_head_stride,
+    int o_batch_stride, int o_row_stride, int o_head_stride,
+    float softmax_scale, int num_sms, cudaStream_t stream);
+
 #ifdef __cplusplus
 }
 #endif

csrc/attention/fa2_wrapper_causal.cu

@@ -7,8 +7,9 @@
 // and is exposed to Python as `flash_rt_fa2.fwd_bf16_causal`
 // (binding added in csrc/fa2_bindings.cpp).
 //
-// Build set is intentionally small: bf16 hdim=128 for Qwen3-8B and
-// bf16 hdim=256 for Qwen3.6 full-attention chunked prefill.
+// Build set: bf16 hdim=128 for Qwen3-8B, bf16 hdim=256 for Qwen3.6
+// full-attention chunked prefill, and fp16 hdim=128 for Chameleon-7B
+// causal attention on Orin SM87.
 //
 // The non-causal wrapper's helpers (fill_params, splitkv heuristic)
 // are duplicated here intentionally to keep this file standalone
@@ -22,6 +23,8 @@
 #include <cstdint>
 #include <cstdio>
 #include <algorithm>
+#include <stdexcept>
+#include <string>
 
 #include "flash_attn_2_src/flash_attn/namespace_config.h"
 #include "flash_attn_2_src/flash_attn/flash.h"
@@ -169,6 +172,7 @@ static int setup_splitkv_causal(FLASH_NAMESPACE::Flash_fwd_params& params,
     return num_splits;
 }
 
+#ifdef FA2_HAS_BF16
 extern "C" void fvk_attention_fa2_fwd_bf16_causal(
     const void* q_ptr, const void* k_ptr, const void* v_ptr,
     void* o_ptr, void* softmax_lse_ptr,
@@ -190,10 +194,9 @@ extern "C" void fvk_attention_fa2_fwd_bf16_causal(
     supported = supported || head_dim == 256;
 #endif
     if (!supported) {
-        fprintf(stderr,
-            "fvk_attention_fa2_fwd_bf16_causal: head_dim=%d not built. "
-            "Enable its FA2_HDIMS entry and rebuild.\n", head_dim);
-        std::abort();
+        throw std::runtime_error(
+            "fvk_attention_fa2_fwd_bf16_causal: head_dim=" + std::to_string(head_dim) +
+            " not built. Enable its FA2_HDIMS entry and rebuild.");
     }
 #else
     if ((head_dim != 128)
@@ -202,15 +205,14 @@ extern "C" void fvk_attention_fa2_fwd_bf16_causal(
 #endif
         ) {
 #ifdef FA2_HAS_HDIM_256
-        fprintf(stderr,
-            "fvk_attention_fa2_fwd_bf16_causal: head_dim=%d not built. "
-            "Only head_dim=128 and 256 are currently instantiated.\n", head_dim);
+        throw std::runtime_error(
+            "fvk_attention_fa2_fwd_bf16_causal: head_dim=" + std::to_string(head_dim) +
+            " not built. Only head_dim=128 and 256 are currently instantiated.");
 #else
-        fprintf(stderr,
-            "fvk_attention_fa2_fwd_bf16_causal: head_dim=%d not built. "
-            "Only head_dim=128 is currently instantiated.\n", head_dim);
+        throw std::runtime_error(
+            "fvk_attention_fa2_fwd_bf16_causal: head_dim=" + std::to_string(head_dim) +
+            " not built. Only head_dim=128 is currently instantiated.");
 #endif
-        std::abort();
     }
 #endif
 
@@ -253,10 +255,9 @@ extern "C" void fvk_attention_fa2_fwd_bf16_causal(
             return;
 #endif
         default:
-            fprintf(stderr,
-                "fvk_attention_fa2_fwd_bf16_causal: head_dim=%d not built "
-                "in this FA2 matrix.\n", head_dim);
-            std::abort();
+            throw std::runtime_error(
+                "fvk_attention_fa2_fwd_bf16_causal: head_dim=" + std::to_string(head_dim) +
+                " not built in this FA2 matrix.");
     }
 #else
     if (head_dim == 128 && num_splits > 1) {
@@ -272,11 +273,87 @@ extern "C" void fvk_attention_fa2_fwd_bf16_causal(
     }
 #else
     else {
-        fprintf(stderr,
-            "fvk_attention_fa2_fwd_bf16_causal: head_dim=%d not built "
-            "(hdim=256 disabled at compile time).\n", head_dim);
-        std::abort();
+        throw std::runtime_error(
+            "fvk_attention_fa2_fwd_bf16_causal: head_dim=" + std::to_string(head_dim) +
+            " not built (hdim=256 disabled at compile time).");
     }
 #endif
 #endif
 }
+#else  // !FA2_HAS_BF16
+extern "C" void fvk_attention_fa2_fwd_bf16_causal(
+    const void*, const void*, const void*, void*, void*,
+    void*, void*,
+    int, int, int, int, int, int,
+    int, int, int, int, int, int,
+    int, int, int, int, int, int,
+    float, int, cudaStream_t)
+{
+    throw std::runtime_error(
+        "fvk_attention_fa2_fwd_bf16_causal: bf16 entry was not compiled. "
+        "Rebuild with -DFA2_DTYPES=\"fp16;bf16\" to enable it.");
+}
+#endif  // FA2_HAS_BF16
+
+// FP16 causal sibling. Only head_dim=128 is instantiated (Chameleon-7B
+// on Orin SM87 is the consumer; bf16 covers the head_dim=256 shapes
+// used by Qwen3.6 chunked prefill).
+#if defined(FA2_HAS_FP16) && defined(FA2_HAS_HDIM_128)
+extern "C" void fvk_attention_fa2_fwd_fp16_causal(
+    const void* q_ptr, const void* k_ptr, const void* v_ptr,
+    void* o_ptr, void* softmax_lse_ptr,
+    void* softmax_lse_accum_ptr, void* o_accum_ptr,
+    int batch, int seqlen_q, int seqlen_k,
+    int num_heads_q, int num_heads_kv, int head_dim,
+    int q_batch_stride, int q_row_stride, int q_head_stride,
+    int k_batch_stride, int k_row_stride, int k_head_stride,
+    int v_batch_stride, int v_row_stride, int v_head_stride,
+    int o_batch_stride, int o_row_stride, int o_head_stride,
+    float softmax_scale, int num_sms, cudaStream_t stream)
+{
+    if (head_dim != 128) {
+        throw std::runtime_error(
+            "fvk_attention_fa2_fwd_fp16_causal: head_dim=" + std::to_string(head_dim) +
+            " not built. Only head_dim=128 is currently instantiated for the "
+            "fp16 causal path. Add a new file under csrc/attention/fa2_causal_inst/ "
+            "and extend the dispatch in fa2_wrapper_causal.cu to support "
+            "additional shapes.");
+    }
+
+    FLASH_NAMESPACE::Flash_fwd_params params;
+    fill_params_causal(params,
+                       q_ptr, k_ptr, v_ptr, o_ptr, softmax_lse_ptr,
+                       batch, seqlen_q, seqlen_k,
+                       num_heads_q, num_heads_kv, head_dim,
+                       q_batch_stride, q_row_stride, q_head_stride,
+                       k_batch_stride, k_row_stride, k_head_stride,
+                       v_batch_stride, v_row_stride, v_head_stride,
+                       o_batch_stride, o_row_stride, o_head_stride,
+                       softmax_scale);
+    // fill_params_causal hardcodes is_bf16=true; flip it for the fp16 path.
+    params.is_bf16 = false;
+
+    int num_splits = setup_splitkv_causal(params, softmax_lse_accum_ptr, o_accum_ptr,
+                                          num_sms, seqlen_q, seqlen_k,
+                                          head_dim, batch, num_heads_q);
+    if (num_splits > 1) {
+        FLASH_NAMESPACE::run_mha_fwd_splitkv_dispatch<cutlass::half_t, 128, true>(params, stream);
+    } else {
+        FLASH_NAMESPACE::run_mha_fwd_<cutlass::half_t, 128, true>(params, stream);
+    }
+}
+#else  // !(FA2_HAS_FP16 && FA2_HAS_HDIM_128)
+extern "C" void fvk_attention_fa2_fwd_fp16_causal(
+    const void*, const void*, const void*, void*, void*,
+    void*, void*,
+    int, int, int, int, int, int,
+    int, int, int, int, int, int,
+    int, int, int, int, int, int,
+    float, int, cudaStream_t)
+{
+    throw std::runtime_error(
+        "fvk_attention_fa2_fwd_fp16_causal: fp16 hdim=128 entry was not "
+        "compiled. Rebuild with -DFA2_DTYPES=\"fp16;bf16\" and "
+        "-DFA2_HDIMS including 128 (and FLASHRT_ENABLE_CHAMELEON=ON) to enable it.");
+}
+#endif  // FA2_HAS_FP16 && FA2_HAS_HDIM_128