feat(qwen35): add opt-in SM120 FlashInfer GDN prefill

README.md

@@ -71,7 +71,7 @@ pegainfer --model-path models/Qwen3-4B
 - Rust (2024 edition), CUDA Toolkit (nvcc, cuBLAS), CUDA-capable GPU
 - NVIDIA driver R545 (CUDA 12.3) or newer; `cuFuncGetName` sets this floor, while per-symbol lazy loading keeps the `cuda-12090` cudarc binding from requiring a CUDA 12.9 driver
 - The default build (Qwen3-4B / 8B) is pure Rust + CUDA — no Python at all
-- Python 3 + Triton for `qwen35` feature builds (build-time only — no Python at runtime)
+- Python 3 + Triton for `qwen35` feature builds (build-time only — no Python at runtime). The supported Qwen3.5-4B SM120/Hv32 path can additionally link a pre-generated FlashInfer/CuTe GDN AOT candidate; generation is separate from serving
 - The `kimi-k2` EP path additionally needs NCCL ≥ 2.27 at runtime (`ncclAlltoAll`)
 
 ### Build & Run
@@ -105,7 +105,8 @@ curl -N http://localhost:8000/v1/completions \
 <summary>More options</summary>
 
 ```bash
-# Qwen3.5 requires the feature-gated Triton AOT kernels (Python + Triton at build time)
+# Qwen3.5 uses build-time Triton AOT; a validated SM120/Hv32 FlashInfer GDN
+# candidate can additionally be selected through PEGAINFER_QWEN35_GDN_AOT_BUNDLE
 uv venv && uv pip install triton
 export PEGAINFER_TRITON_PYTHON=.venv/bin/python
 cargo run --release --features qwen35 -- --model-path models/Qwen3.5-4B
@@ -120,6 +121,7 @@ cargo run --release -- --cuda-graph=false
 |----------|-------------|
 | `CUDA_HOME` | CUDA Toolkit path (default: `/usr/local/cuda`) |
 | `PEGAINFER_TRITON_PYTHON` | Python with Triton for `qwen35` build-time AOT compilation |
+| `PEGAINFER_QWEN35_GDN_AOT_BUNDLE` | Validated Qwen3.5-4B SM120/Hv32 FlashInfer GDN candidate directory linked by `pegainfer-kernels/build.rs` |
 | `PEGAINFER_TILELANG_PYTHON` | Python with TileLang for `k3` build-time kernel generation |
 | `PEGAINFER_CUDA_SM` | GPU SM target override when `nvidia-smi` unavailable (e.g. `120`) |
 
@@ -303,14 +305,14 @@ flowchart TB
 **Key design decisions:**
 
 - **GPU-first runtime** — model execution stays in native Rust/CUDA paths
-- **Custom GPU kernels** — CUDA for decode-critical paths, Triton AOT for Qwen3.5 compatibility kernels, FlashInfer for paged attention/sampling, NCCL for multi-GPU reductions, and cuBLAS for matrix multiplication
+- **Custom GPU kernels** — CUDA for decode-critical paths; Triton AOT for the general Qwen3.5 GDN path; a statically linked FlashInfer/CuTe AOT specialization for Qwen3.5-4B SM120/Hv32 when a validated candidate is supplied; FlashInfer for paged attention/sampling; NCCL for multi-GPU reductions; and cuBLAS for matrix multiplication
 - **CUDA Graph** on Qwen decode paths — eliminates kernel launch overhead where enabled
 - **Per-model crate boundary** — Qwen3-4B owns its config, weights, scheduler/executor, tests, benches, and kernel plan in `pegainfer-qwen3`
 
 **Model details:**
 
 - **Qwen3**: 32 Q heads, 8 KV heads (GQA 4:1), head_dim=128
-- **Qwen3.5**: hybrid — 24 linear attention layers (Gated Delta Rule) + 8 full attention layers, head_dim=256
+- **Qwen3.5**: hybrid — 24 linear attention layers (Gated Delta Rule) + 8 full attention layers, head_dim=256. Qwen3.5-4B on single-GPU SM120 uses the validated FlashInfer GDN AOT specialization when linked; unsupported geometry/SM/TP configurations retain the explicit Triton path
 - **DeepSeek V2-Lite**: feature-gated 2-GPU EP2 correctness/attribution path for the HF/host-staged/NCCL narrow greedy gate
 
 ### What's not (yet) implemented
@@ -350,6 +352,8 @@ PEGAINFER_TEST_MODEL_PATH=models/Qwen3.5-4B cargo test --release -p pegainfer-qw
 PEGAINFER_TEST_MODEL_PATH=models/DeepSeek-V2-Lite cargo test --release -p pegainfer-deepseek-v2-lite --features deepseek-v2-lite --test e2e_ep2 -- --nocapture
 ```
 
+The SM120 FlashInfer GDN production boundary has a separate fail-closed five-gate runner because it requires a real generated candidate, the pinned Qwen3.5-4B snapshot, and an SM120 GPU. See [`pegainfer-kernels/tools/flashinfer_gdn/README.md`](pegainfer-kernels/tools/flashinfer_gdn/README.md); the non-default `gdn-validation` feature exists only for that runner and adds no counters or validation API to a default serving build.
+
 The DeepSeek-V2-Lite E2E is a correctness/integration gate. Direct diagnostics and HTTP SLO report commands live in [`benchmarking.md`](docs/models/deepseek-v2-lite/benchmarking.md).
 
 ## License

pegainfer-kernels/Cargo.toml

@@ -15,6 +15,8 @@ tvm-ffi = { version = "0.1.0-alpha.0", optional = true }
 [build-dependencies]
 cc = { workspace = true }
 pegainfer-build = { workspace = true }
+serde_json = { workspace = true, optional = true }
+sha2 = { workspace = true, optional = true }
 
 [features]
 default = []
@@ -24,7 +26,7 @@ tvm-ffi-triton-cubin = ["dep:tvm-ffi", "qwen35"]
 deepseek-v2-lite = []
 # Gemma 4: NVFP4 dequantization, whose conversion intrinsics need CUDA >= 12.8.
 gemma4 = []
-qwen35 = []
+qwen35 = ["dep:serde_json", "dep:sha2"]
 # Shared MoE/MLA third-party substrate: DeepEP, DeepGEMM, and FlashMLA.
 glm52 = ["moe"]
 # Kimi K3: TileLang-generated kernels, FP8xFP4 masked grouped GEMM, and the

pegainfer-kernels/KERNELS.md

@@ -231,7 +231,7 @@ The crate still builds CUDA/Triton symbols needed by the current root binary:
 
 - Qwen3.5 HD256 full-attention kernels: `csrc/qwen35/prefill_attention_hd256.cu`, `csrc/shared/paged_attention.cu`.
 - Qwen3.5 linear-attention decode kernels: `csrc/qwen35/conv1d.cu`, `csrc/qwen35/gated_delta_rule.cu`.
-- Qwen3.5 chunk-wise GDR prefill Triton AOT kernels: `tools/triton/gated_delta_rule_chunkwise_kernels.py`.
+- Qwen3.5 chunk-wise GDR prefill uses the Triton AOT kernels in `tools/triton/gated_delta_rule_chunkwise_kernels.py` generally. A validated build-linked candidate selects the FlashInfer/CuTe AOT specialization in `csrc/qwen35/flashinfer_gdn_aot.c` for single-GPU SM120 Qwen3.5-4B (`Hq/Hk/Hv/D=16/16/32/128`); unsupported SM, geometry, and TP configurations retain Triton.
 
 These are preserved for build compatibility. They are not part of the Qwen3-4B Phase 1 API surface.
 

pegainfer-kernels/build.rs

@@ -9,6 +9,9 @@ use std::sync::Mutex;
 use std::thread;
 use std::time::Instant;
 
+#[cfg(feature = "qwen35")]
+use sha2::Digest as _;
+
 struct TritonKernelSpec {
     artifact_dir: &'static str,
     kernel_path: &'static str,
@@ -39,6 +42,177 @@ struct FlashInferIncludes {
     cccl: Vec<PathBuf>,
 }
 
+#[cfg(feature = "qwen35")]
+const QWEN35_GDN_AOT_ABI_VERSION: u64 = 1;
+#[cfg(feature = "qwen35")]
+const QWEN35_GDN_AOT_ENV: &str = "PEGAINFER_QWEN35_GDN_AOT_BUNDLE";
+
+#[cfg(feature = "qwen35")]
+fn sha256_file(path: &Path) -> String {
+    let bytes =
+        fs::read(path).unwrap_or_else(|error| panic!("failed to read {}: {error}", path.display()));
+    let digest = sha2::Sha256::digest(bytes);
+    let mut hex = String::with_capacity(digest.len() * 2);
+    for byte in digest {
+        write!(&mut hex, "{byte:02x}").expect("write SHA-256 hex to String");
+    }
+    hex
+}
+
+#[cfg(feature = "qwen35")]
+fn json_u64(value: &serde_json::Value, path: &[&str]) -> u64 {
+    let mut cursor = value;
+    for key in path {
+        cursor = &cursor[*key];
+    }
+    cursor.as_u64().unwrap_or_else(|| {
+        panic!(
+            "GDN AOT manifest field {} must be an unsigned integer",
+            path.join(".")
+        )
+    })
+}
+
+#[cfg(feature = "qwen35")]
+fn json_str<'a>(value: &'a serde_json::Value, path: &[&str]) -> &'a str {
+    let mut cursor = value;
+    for key in path {
+        cursor = &cursor[*key];
+    }
+    cursor
+        .as_str()
+        .unwrap_or_else(|| panic!("GDN AOT manifest field {} must be a string", path.join(".")))
+}
+
+/// Validate and attach the release-provided Qwen3.5 GDN object. The generated
+/// object and its native CuTe runtime archive are linked statically; serving
+/// never reads a manifest, loads PTX, or discovers a Python wheel.
+#[cfg(feature = "qwen35")]
+fn build_qwen35_flashinfer_gdn_aot(
+    root: &Path,
+    out_dir: &Path,
+    cuda_include: &Path,
+) -> (Vec<PathBuf>, Option<PathBuf>) {
+    println!("cargo:rerun-if-env-changed={QWEN35_GDN_AOT_ENV}");
+    let shim = root.join("csrc/qwen35/flashinfer_gdn_aot.c");
+    let shim_header = root.join("csrc/qwen35/flashinfer_gdn_aot.h");
+    println!("cargo:rerun-if-changed={}", shim.display());
+    println!("cargo:rerun-if-changed={}", shim_header.display());
+
+    let config_header = out_dir.join("flashinfer_gdn_build_config.h");
+    let mut includes = vec![root.join("csrc/qwen35"), out_dir.to_path_buf()];
+    let mut linked_objects = Vec::new();
+    let mut runtime_dir = None;
+    let mut config = String::from(
+        "#pragma once\n#define PEGAINFER_QWEN35_GDN_ARTIFACT_SHA256 \"unavailable\"\n#define PEGAINFER_QWEN35_GDN_WORKSPACE_BYTES_PER_SM 128u\n",
+    );
+
+    if let Some(bundle) = std::env::var_os(QWEN35_GDN_AOT_ENV) {
+        let bundle = PathBuf::from(bundle);
+        let manifest_path = bundle.join("manifest.json");
+        let manifest_bytes = fs::read(&manifest_path).unwrap_or_else(|error| {
+            panic!("read GDN AOT manifest {}: {error}", manifest_path.display())
+        });
+        let manifest: serde_json::Value = serde_json::from_slice(&manifest_bytes)
+            .unwrap_or_else(|error| panic!("parse GDN AOT manifest: {error}"));
+        assert_eq!(json_u64(&manifest, &["schema_version"]), 3);
+        assert_eq!(json_str(&manifest, &["variant"]), "qwen35_4b_candidate");
+        assert_eq!(json_str(&manifest, &["target", "arch"]), "sm_120a");
+        assert_eq!(
+            json_str(&manifest, &["target", "code_object"]),
+            "embedded_cubin"
+        );
+        assert_eq!(
+            json_u64(&manifest, &["abi", "version"]),
+            QWEN35_GDN_AOT_ABI_VERSION
+        );
+        assert_eq!(json_u64(&manifest, &["geometry", "h_q"]), 16);
+        assert_eq!(json_u64(&manifest, &["geometry", "h_k"]), 16);
+        assert_eq!(json_u64(&manifest, &["geometry", "h_v"]), 32);
+        assert_eq!(json_u64(&manifest, &["geometry", "head_dim"]), 128);
+        assert_eq!(json_str(&manifest, &["tokens", "extent"]), "dynamic");
+        assert_eq!(json_u64(&manifest, &["tokens", "minimum"]), 1);
+        assert_eq!(
+            json_str(&manifest, &["abi", "state_layout"]),
+            "openinfer_hkv_v_contiguous"
+        );
+        assert_eq!(json_str(&manifest, &["workspace", "kind"]), "per_sm");
+        let workspace_bytes_per_sm = json_u64(&manifest, &["workspace", "bytes_per_sm"]);
+        assert_eq!(workspace_bytes_per_sm, 128);
+        assert_eq!(json_u64(&manifest, &["workspace", "alignment_bytes"]), 128);
+
+        let header = bundle.join("kernel.h");
+        let object = bundle.join("kernel.o");
+        let runtime = bundle.join("libcuda_dialect_runtime_static.a");
+        for (label, path, hash_path, size_path) in [
+            (
+                "header",
+                &header,
+                ["artifact", "header", "sha256"],
+                ["artifact", "header", "size_bytes"],
+            ),
+            (
+                "object",
+                &object,
+                ["artifact", "object", "sha256"],
+                ["artifact", "object", "size_bytes"],
+            ),
+            (
+                "native runtime",
+                &runtime,
+                ["artifact", "native_runtime", "sha256"],
+                ["artifact", "native_runtime", "size_bytes"],
+            ),
+        ] {
+            assert!(
+                path.is_file(),
+                "GDN AOT {label} is missing: {}",
+                path.display()
+            );
+            assert_eq!(sha256_file(path), json_str(&manifest, &hash_path));
+            assert_eq!(
+                fs::metadata(path).expect("read GDN AOT metadata").len(),
+                json_u64(&manifest, &size_path)
+            );
+            println!("cargo:rerun-if-changed={}", path.display());
+        }
+        println!("cargo:rerun-if-changed={}", manifest_path.display());
+
+        let object_hash = json_str(&manifest, &["artifact", "object", "sha256"]);
+        config = format!(
+            "#pragma once\n#define PEGAINFER_QWEN35_GDN_ARTIFACT_SHA256 \"{object_hash}\"\n#define PEGAINFER_QWEN35_GDN_WORKSPACE_BYTES_PER_SM {workspace_bytes_per_sm}u\n"
+        );
+        includes.push(bundle);
+        linked_objects.push(object);
+        runtime_dir = runtime.parent().map(Path::to_path_buf);
+    }
+    fs::write(&config_header, config).expect("write GDN AOT build config");
+
+    let shim_obj = out_dir.join("qwen35_flashinfer_gdn_aot.o");
+    let compiler = cc::Build::new().get_compiler();
+    let mut command = compiler.to_command();
+    command
+        .arg("-c")
+        .arg(&shim)
+        .arg("-o")
+        .arg(&shim_obj)
+        .arg("-O3")
+        .arg("-std=c11")
+        .arg("-fPIC")
+        .arg("-isystem")
+        .arg(cuda_include);
+    for include in includes {
+        command.arg("-I").arg(include);
+    }
+    if runtime_dir.is_some() {
+        command.arg("-DPEGAINFER_QWEN35_GDN_AOT");
+    }
+    let status = command.status().expect("compile Qwen3.5 GDN AOT shim");
+    assert!(status.success(), "Qwen3.5 GDN AOT shim compilation failed");
+    linked_objects.push(shim_obj);
+    (linked_objects, runtime_dir)
+}
+
 const GLM52_TRTLLM_FMHA_CUBINS: &[(&str, &str)] = &[
     (
         "kGlm52FmhaSparseSeedQ8",
@@ -1904,6 +2078,14 @@ fn main() {
     // --- k3: DeepGEMM-only, no DeepEP/NCCL dependency ---
     let k3_enabled = cfg!(feature = "k3");
     let qwen35_enabled = cfg!(feature = "qwen35");
+    #[cfg(feature = "qwen35")]
+    let (qwen35_gdn_objects, qwen35_gdn_runtime_dir) = if qwen35_enabled {
+        build_qwen35_flashinfer_gdn_aot(&crate_root(), &out_dir, &cuda_include)
+    } else {
+        (Vec::new(), None)
+    };
+    #[cfg(not(feature = "qwen35"))]
+    let (qwen35_gdn_objects, qwen35_gdn_runtime_dir) = (Vec::new(), None::<PathBuf>);
     if glm52_enabled {
         generate_glm52_trtllm_fmha_cubins(&crate_root(), &out_dir);
         build_glm52_cutedsl_fp8_dsl(&crate_root(), &out_dir, &cuda_include);
@@ -1970,6 +2152,9 @@ fn main() {
             if !kimi_k2_enabled && is_kimi_k2_source(&csrc_dir, path) {
                 return None;
             }
+            if !qwen35_enabled && file_name == "gdn_prepare.cu" {
+                return None;
+            }
             // --- k3 ---
             if !k3_enabled && is_k3_source(&csrc_dir, path) {
                 return None;
@@ -2486,6 +2671,7 @@ fn main() {
     ar_args.extend(
         obj_files
             .into_iter()
+            .chain(qwen35_gdn_objects)
             .map(|path| path.to_string_lossy().to_string()),
     );
 
@@ -2525,6 +2711,10 @@ fn main() {
         toolkit.link_search();
     }
     println!("cargo:rustc-link-lib=static=kernels_cuda");
+    if let Some(runtime_dir) = qwen35_gdn_runtime_dir {
+        println!("cargo:rustc-link-search=native={}", runtime_dir.display());
+        println!("cargo:rustc-link-lib=static=cuda_dialect_runtime_static");
+    }
     println!("cargo:rustc-link-lib=cudart");
     println!("cargo:rustc-link-lib=cublas");
     println!("cargo:rustc-link-lib=cublasLt");