perf: Qwen3 offline inference — hybrid attention backend, fused rms_norm_rope, chunked prefill + prompt-lookup spec decoding

.github/ci_config.yaml

@@ -46,7 +46,7 @@ platforms:
           shm_size: 64g
           timeout: 3600
         env:
-          TEST_PARAM: ['default', '--enable-paged-attn', '--enable-paged-attn --enable-graph', '--enable-paged-attn --enable-graph --attn=flash-attn']
+          TEST_PARAM: ['default', '--enable-paged-attn', '--enable-paged-attn --enable-graph', '--enable-paged-attn --enable-graph --attn=flash-attn', '--enable-paged-attn --enable-graph --attn=hybrid']
         stages:
           - name: test
             run: python InfiniLM/examples/bench.py --device nvidia --model=/data-aisoft/mechdancer/models/9g_8b_thinking/ --input-len=256,1024 --output-len=256,1024 --batch-size=8 <TEST_PARAM>

.gitignore

@@ -9,6 +9,9 @@ python/infinilm/lib/*.so
 # Vscode
 .vscode/
 
+# JetBrains
+.idea/
+
 *.sh
 model_weight/
 
@@ -34,3 +37,4 @@ __pycache__/
 *.http
 
 *.nsys-rep
+dev_perf/results/

csrc/backends/attention_backends.hpp

@@ -8,12 +8,21 @@ namespace infinilm::backends {
 
 /**
  * @brief Enumeration of all supported attention backends.
+ *
+ * 各后端说明：
+ * - STATIC_ATTN：静态 attention，prefill/decode 走同一套实现（默认）
+ * - PAGED_ATTN：自研 paged-attention，prefill 为 PagedAttentionPrefill，decode 为 splitkv
+ * - FLASH_ATTN：FlashAttention-2（mha_varlen_fwd / mha_fwd_kvcache）
+ * - FLASHINFER：FlashInfer 后端
+ * - HYBRID：按阶段分离路由——prefill 走 FA2 varlen，decode 走自研 paged kernel
+ *   （见 HybridAttentionImpl）
  */
 enum class AttentionBackend {
     STATIC_ATTN,
     PAGED_ATTN,
     FLASH_ATTN,
     FLASHINFER,
+    HYBRID, // prefill → FlashAttention (FA2 varlen), decode → PagedAttention
     Default = STATIC_ATTN
 };
 
@@ -27,6 +36,8 @@ inline std::ostream &operator<<(std::ostream &os, AttentionBackend backend) {
         return os << "AttentionBackend::FLASH_ATTN";
     case AttentionBackend::FLASHINFER:
         return os << "AttentionBackend::FLASHINFER";
+    case AttentionBackend::HYBRID:
+        return os << "AttentionBackend::HYBRID";
     default:
         throw std::invalid_argument("infinilm::backends: invalid attention backend: " + std::to_string(static_cast<int>(backend)));
         break;
@@ -49,9 +60,13 @@ inline AttentionBackend parse_attention_backend(const std::string &backend) {
     if (backend == "flashinfer") {
         return AttentionBackend::FLASHINFER;
     }
+    if (backend == "hybrid") {
+        // "hybrid"：prefill→FA2 varlen，decode→自研 paged-attention（splitkv）
+        return AttentionBackend::HYBRID;
+    }
 
     throw std::invalid_argument(
-        "Invalid attention_backend: " + backend + ". Valid options are: static-attn, paged-attn, flash-attn, flashinfer");
+        "Invalid attention_backend: " + backend + ". Valid options are: static-attn, paged-attn, flash-attn, flashinfer, hybrid");
 }
 
 } // namespace infinilm::backends

csrc/cache/kv_cache.cpp

@@ -131,8 +131,13 @@ infinicore::Tensor create_layer_kv_cache(
     size_t block_size = config.block_size();
 
     infinicore::Shape kv_shape;
-    if (global_state::get_infinilm_config().attention_backend == backends::AttentionBackend::FLASH_ATTN) {
-        // FLASH_ATTN kernel expects BSHD layout
+    const auto cache_attn_backend = global_state::get_infinilm_config().attention_backend;
+    if (cache_attn_backend == backends::AttentionBackend::FLASH_ATTN || cache_attn_backend == backends::AttentionBackend::HYBRID) {
+        // FLASH_ATTN kernel expects BSHD layout. HYBRID uses the same layout;
+        // its decode paged-attention kernel reads it via strides.
+        // FLASH_ATTN 的 kernel 需要 BSHD 布局（块 → 块内 token → head → dim）；
+        // HYBRID 沿用同一布局——decode 阶段自研 paged kernel 通过 stride
+        // 以 BHSD 逻辑视图零拷贝读取，无需切换布局。
         kv_shape = {2, num_blocks_per_layer, block_size, num_rank_k_heads, k_dim};
     } else {
         kv_shape = {2, num_blocks_per_layer, num_rank_k_heads, block_size, k_dim};

csrc/engine/compiler/general_compiler.cpp

@@ -23,4 +23,13 @@ GeneralCompiler::Compiled GeneralCompiler::get_compiled(const InfinilmModel::Inp
     return result;
 }
 
+std::pair<std::shared_ptr<infinicore::graph::Graph>, infinicore::Tensor>
+GeneralCompiler::get_sampling_compiled(size_t batch_size) {
+    auto result = paged_compiler_->get_sampling_compiled(batch_size);
+    if (result.first != nullptr) {
+        return result;
+    }
+    return static_batching_compiler_->get_sampling_compiled(batch_size);
+}
+
 } // namespace infinilm::engine

csrc/engine/compiler/general_compiler.hpp

@@ -12,6 +12,9 @@ class GeneralCompiler : public GraphCompiler {
 
     Compiled get_compiled(const InfinilmModel::Input &input) override;
 
+    std::pair<std::shared_ptr<infinicore::graph::Graph>, infinicore::Tensor>
+    get_sampling_compiled(size_t batch_size) override;
+
 private:
     std::unique_ptr<StaticBatchingCompiler> static_batching_compiler_;
     std::unique_ptr<PagedCompiler> paged_compiler_;

csrc/engine/compiler/graph_compiler.hpp

@@ -17,6 +17,16 @@ class GraphCompiler {
     virtual void compile() = 0;
     virtual Compiled get_compiled(const InfinilmModel::Input &input) = 0;
 
+    // Optional companion graph that replays captured greedy-sampling kernels
+    // (per-request argmax) for a decode batch size. Only valid right after a
+    // successful get_compiled()+run() for the same batch size, since it reads
+    // that graph's logits blob. Default: no captured sampling.
+    virtual std::pair<std::shared_ptr<infinicore::graph::Graph>, infinicore::Tensor>
+    get_sampling_compiled(size_t batch_size) {
+        (void)batch_size;
+        return {nullptr, {}};
+    }
+
 protected:
     std::shared_ptr<InfinilmModel> model_;
     RankBarrier *barrier_;