perf(qwen): reuse PrefillPagedPlan for uncompiled-GQA decode

docs/index.md

@@ -157,6 +157,7 @@ Organized by domain (model line / subsystem / playbook / lesson) instead of by l
 | `subsystems/kernels/openinfer-kernels-boundary.md` | Architecture decision: reusable frontend/runtime/data-plane layers plus per-model engines; `openinfer-kernels` keeps shared MoE/MLA substrate (`moe`: DeepEP/DeepGEMM/FlashMLA) separate from model-local surfaces such as the narrow GLM5.2 DeepGEMM/FlashMLA wrappers. |
 | `subsystems/kernels/build-rs-submodule-init.md` | `openinfer-kernels/build.rs` initializes missing git submodules automatically for first-time builds before checking vendored third-party kernel headers. |
 | `subsystems/kernels/kernel-op-reports.md` | Qwen3 kernel/report tooling is feature-gated: `qwen3_kernel_report` covers per-op kernel reports, and `qwen3_model_report` emits runtime-traced eager-DAG decode operator rollups with TensorSpec `KernelCall`s, latency stats, tables, and Graphviz DOT; measured FA2 `CTA_TILE_Q=64` prefill default in place. |
+| `subsystems/kernels/reusable-prefill-plan.md` | Issue #711 / PR #714 reuses uncompiled-GQA prefill metadata with a scheduler-derived logical page-table bound and one checked allocation/accounting layout; Qwen3-14B Nsight A/B removes 22 allocation/free calls per step, while pointer stability, shared-prefix mixed load, both production golden suites, and Qwen3.5 scheduler E2E pass. |
 | `subsystems/kernels/typed-forward-pipeline.md` | Reusable typed tensor pipeline macro in `openinfer-kernels` so model crates can express common `typed_ops` chains without model-specific wrapper macros. |
 | `subsystems/kernels/tvm-ffi-mvp.md` | Optional `tvm-ffi-triton-cubin` bridge in `openinfer-kernels` plus a packed TVM wrapper for the Qwen3.5 GDR solve Triton AOT CUBIN launcher. |
 

openinfer-core/src/ops.rs

@@ -81,6 +81,7 @@ pub use openinfer_kernels::ops::single_prefill_nhd_causal_into;
 pub use openinfer_kernels::ops::single_prefill_nhd_noncausal_into;
 pub use openinfer_kernels::ops::write_vec_into;
 pub use paged_plan::PrefillPagedPlan;
+pub use paged_plan::PrefillPagedPlanCapacity;
 #[cfg(feature = "kernel-call-trace")]
 pub use traced::embedding_batch;
 #[cfg(feature = "kernel-call-trace")]

openinfer-core/src/ops/paged_plan.rs

@@ -6,11 +6,108 @@ use cudarc::driver::CudaSlice;
 use crate::kv_pool::KvDesc;
 use crate::tensor::DeviceContext;
 
+/// Checked dimensions for a reusable [`PrefillPagedPlan`].
+///
+/// `max_page_indices` counts logical page-table entries across every request,
+/// not unique physical KV pages. Prefix-cached requests may reference the same
+/// physical page, so the logical bound is `max_batch * pages_per_request`.
+#[derive(Clone, Copy, Debug, Eq, PartialEq)]
+pub struct PrefillPagedPlanCapacity {
+    total_tokens: usize,
+    page_indices: usize,
+    batch: usize,
+    tiles: usize,
+}
+
+impl PrefillPagedPlanCapacity {
+    pub fn for_context(
+        max_total_tokens: usize,
+        max_batch: usize,
+        max_context_tokens: usize,
+        page_size: usize,
+        gqa_group_size: usize,
+    ) -> Result<Self> {
+        anyhow::ensure!(
+            max_total_tokens > 0,
+            "prefill plan token capacity must be positive"
+        );
+        anyhow::ensure!(
+            max_batch > 0,
+            "prefill plan batch capacity must be positive"
+        );
+        anyhow::ensure!(
+            max_context_tokens > 0,
+            "prefill plan context capacity must be positive"
+        );
+        anyhow::ensure!(page_size > 0, "prefill plan page size must be positive");
+        anyhow::ensure!(
+            gqa_group_size > 0,
+            "prefill plan GQA group size must be positive"
+        );
+
+        let max_pages_per_request = max_context_tokens.div_ceil(page_size);
+        let max_page_indices = max_batch
+            .checked_mul(max_pages_per_request)
+            .ok_or_else(|| anyhow::anyhow!("prefill plan logical page capacity overflow"))?;
+        let max_tiles = max_total_tokens
+            .checked_mul(gqa_group_size)
+            .ok_or_else(|| anyhow::anyhow!("prefill plan tile capacity overflow"))?;
+        let capacity = Self {
+            total_tokens: max_total_tokens,
+            page_indices: max_page_indices,
+            batch: max_batch,
+            tiles: max_tiles,
+        };
+        capacity.preallocated_bytes()?;
+        Ok(capacity)
+    }
+
+    pub fn preallocated_bytes(self) -> Result<usize> {
+        PrefillPagedPlan::preallocated_bytes(
+            self.total_tokens,
+            self.page_indices,
+            self.batch,
+            self.tiles,
+        )
+    }
+
+    pub fn max_total_tokens(self) -> usize {
+        self.total_tokens
+    }
+
+    pub fn max_page_indices(self) -> usize {
+        self.page_indices
+    }
+
+    pub fn max_batch(self) -> usize {
+        self.batch
+    }
+
+    pub fn max_tiles(self) -> usize {
+        self.tiles
+    }
+}
+
 pub struct PrefillPagedPlan {
     inner: openinfer_kernels::ops::PrefillPagedPlan,
 }
 
 impl PrefillPagedPlan {
+    /// Exact device bytes reserved by [`Self::new_preallocated`].
+    pub fn preallocated_bytes(
+        max_total_tokens: usize,
+        max_page_indices: usize,
+        max_batch: usize,
+        max_tiles: usize,
+    ) -> Result<usize> {
+        openinfer_kernels::ops::PrefillPagedPlan::preallocated_bytes(
+            max_total_tokens,
+            max_page_indices,
+            max_batch,
+            max_tiles,
+        )
+    }
+
     pub fn new(
         ctx: &DeviceContext,
         desc: &KvDesc<'_>,
@@ -95,21 +192,36 @@ impl PrefillPagedPlan {
     pub fn new_preallocated(
         ctx: &DeviceContext,
         max_total_tokens: usize,
-        max_total_pages: usize,
+        max_page_indices: usize,
         max_batch: usize,
         max_tiles: usize,
     ) -> Result<Self> {
         Ok(Self {
             inner: openinfer_kernels::ops::PrefillPagedPlan::new_preallocated(
                 ctx,
                 max_total_tokens,
-                max_total_pages,
+                max_page_indices,
                 max_batch,
                 max_tiles,
             )?,
         })
     }
 
+    /// Allocate from one checked capacity value so allocation and memory
+    /// accounting cannot silently derive different dimensions.
+    pub fn new_preallocated_for_capacity(
+        ctx: &DeviceContext,
+        capacity: PrefillPagedPlanCapacity,
+    ) -> Result<Self> {
+        Self::new_preallocated(
+            ctx,
+            capacity.total_tokens,
+            capacity.page_indices,
+            capacity.batch,
+            capacity.tiles,
+        )
+    }
+
     /// Refill a pre-allocated plan in place (no allocation, pointers unchanged).
     #[allow(clippy::too_many_arguments)]
     pub fn update_batch_with_cta_tile_q(
@@ -179,3 +291,30 @@ impl Deref for PrefillPagedPlan {
         &self.inner
     }
 }
+
+#[cfg(test)]
+mod tests {
+    use std::collections::HashSet;
+
+    use super::PrefillPagedPlanCapacity;
+
+    #[test]
+    fn logical_page_capacity_counts_shared_prefix_for_each_request() {
+        let capacity = PrefillPagedPlanCapacity::for_context(16, 3, 32, 16, 5)
+            .expect("valid reusable plan capacity");
+        let page_lists = [vec![7, 8], vec![7, 8], vec![7, 8]];
+        let logical_pages: usize = page_lists.iter().map(Vec::len).sum();
+        let physical_pages: HashSet<i32> = page_lists.into_iter().flatten().collect();
+
+        assert_eq!(physical_pages.len(), 2);
+        assert_eq!(logical_pages, 6);
+        assert_eq!(capacity.max_page_indices(), logical_pages);
+    }
+
+    #[test]
+    fn capacity_rejects_zero_page_size() {
+        let error = PrefillPagedPlanCapacity::for_context(1, 1, 1, 0, 1)
+            .expect_err("zero page size must be rejected");
+        assert!(error.to_string().contains("page size"));
+    }
+}