perf: move decode split partials to InferenceWorkspace

- Replace per-.cu-file static cached tensors with workspace-managed pre-allocated buffers

- InferenceWorkspace now owns decode_o_part / decode_ml_part (mirrors FlashInfer's workspace pattern)

- KVCache carries the buffers through the backend -> C++ kernel chain

- C++ kernels accept optional pre-allocated buffers; fallback to alloc_split_partials for backward compat

- Pre-allocates once at Executor init, zero allocation in the decode hot loop

- Prerequisite for CUDA-graph capture (all kernel addresses are stable)

astrai/inference/core/cache.py

@@ -260,6 +260,9 @@ class KVCache:
         max_len: max(seq_lens) as Python int — avoids GPU sync in decode
         kv_indptr: [batch+1] int32 — prefix sum of seq_lens, precomputed once
             per step so the attention backend avoids rebuilding it per layer.
+        qo_indptr: [batch+1] int32 — prefill qo prefix-sum (None in decode)
+        decode_o_part: split-KV o partial workspace (mirrors FlashInfer)
+        decode_ml_part: split-KV m/l partial workspace (mirrors FlashInfer)
     """
 
     k_buffer: Tensor
@@ -271,6 +274,8 @@ class KVCache:
     max_len: int = 0
     kv_indptr: Optional[Tensor] = None
     qo_indptr: Optional[Tensor] = None
+    decode_o_part: Optional[Tensor] = None
+    decode_ml_part: Optional[Tensor] = None
 
 
 class PagePool:
@@ -565,12 +570,15 @@ class PagePool:
                 torch.arange(b + 1, dtype=torch.int32, device=device) * q_len
             )
             qo_indptr = workspace.qo_indptr[: b + 1]
+            decode_o_part, decode_ml_part = None, None
         else:
             write_pos = seq_lens_t - 1
             loc = self._req_pool.req_to_token[req_pool_indices, write_pos].unsqueeze(-1)
             ocl_buf[:b].copy_(loc)
             out_cache_loc = ocl_buf[:b]
             qo_indptr = None
+            decode_o_part = getattr(workspace, "decode_o_part", None)
+            decode_ml_part = getattr(workspace, "decode_ml_part", None)
 
         return KVCache(
             k_buffer=self._storage.k_buffer,
@@ -582,6 +590,8 @@ class PagePool:
             max_len=max(seq_lens),
             kv_indptr=kv_indptr,
             qo_indptr=qo_indptr,
+            decode_o_part=decode_o_part,
+            decode_ml_part=decode_ml_part,
         )
 
     # ---- internals ----

astrai/inference/core/executor.py

@@ -78,9 +78,14 @@ class Executor:
         # (input_ids, decode mask, KV bind metadata).  Eagerly sized at init
         # so the workspace is CUDA-graph-capture friendly — no allocation
         # during capture.
+        config = model.config
+        max_q_heads = config.num_attention_heads
+        head_dim = config.hidden_size // config.num_attention_heads
         self._workspace = InferenceWorkspace(
             max_batch_size=kv_cache.max_batch_size,
             max_seq_len=kv_cache.max_seq_len,
+            max_q_heads=max_q_heads,
+            head_dim=head_dim,
             device=self.device,
             dtype=self.dtype,
         )

astrai/inference/core/workspace.py

@@ -1,20 +1,16 @@
 """Pre-allocated buffers for the inference decode hot path.
 
-Mirrors SGLang's pre-allocated input buffers (``input_buffers.py``): tensors
-are sized once to the server's maximum dimensions and sliced to the live
-batch each step, so the per-token decode loop never calls
-``torch.empty``/``torch.zeros``/``torch.arange`` for the hot shapes.  Fills
-go through ``out=`` variants (``torch.ge``) which write into the stable
-buffers instead of allocating fresh results.
-
-All buffers are allocated eagerly at init (nothing is lazy), so the
-workspace is CUDA-graph-capture friendly: the decode step reads/writes
-fixed-address tensors with no allocation during capture.
+Mirrors FlashInfer / SGLang's global workspace pattern: all per-step tensors
+are allocated eagerly at init (nothing is lazy), so the decode step
+reads/writes fixed-address tensors with zero ``torch.empty`` calls during
+the hot loop — a prerequisite for CUDA-graph capture.
 """
 
 import torch
 from torch import Tensor
 
+_MAX_SPLITS = 32
+
 
 class InferenceWorkspace:
     """Reusable fixed-shape per-step buffers for decode.
@@ -30,6 +26,11 @@ class InferenceWorkspace:
     - KV-cache bind metadata (``req_pool_indices``, ``seq_lens``,
       ``kv_indptr``, ``inc``, ``out_cache_loc``), written by
       ``PagePool.bind_tasks`` when the Executor passes this workspace.
+    - ``decode_o_part`` / ``decode_ml_part``: split-KV partial result buffers
+      (mirrors FlashInfer's workspace).  One global alloc, reused by every
+      decode step across all layers.  Sliced views are passed to the CUDA
+      attention kernel so its internal ``torch.empty`` hot-path alloc goes
+      through a stable address (CUDA-graph capturable).
 
     No re-allocation while the server's bounds are respected.
     """
@@ -38,11 +39,15 @@ class InferenceWorkspace:
         self,
         max_batch_size: int,
         max_seq_len: int,
+        max_q_heads: int,
+        head_dim: int,
         device: torch.device,
         dtype: torch.dtype,
     ):
         self.max_batch_size = max_batch_size
         self.max_seq_len = max_seq_len
+        self.max_q_heads = max_q_heads
+        self.head_dim = head_dim
         self.device = device
         self.dtype = dtype
 
@@ -83,6 +88,28 @@ class InferenceWorkspace:
             (max_batch_size, 1), dtype=torch.long, device=device
         )
 
+        # Split-KV partial-result buffers for decode (persistent, one global
+        # alloc per process — mirrors FlashInfer's workspace pattern).
+        # Shape: [max_batch_size, max_q_heads, _MAX_SPLITS, head_dim] (o_part)
+        #        [max_batch_size, max_q_heads, _MAX_SPLITS, 2]     (ml_part)
+        self.decode_o_part = torch.empty(
+            (max_batch_size, max_q_heads, _MAX_SPLITS, head_dim),
+            dtype=torch.float32,
+            device=device,
+        )
+        self.decode_ml_part = torch.empty(
+            (max_batch_size, max_q_heads, _MAX_SPLITS, 2),
+            dtype=torch.float32,
+            device=device,
+        )
+
+    def decode_buffers(self, batch: int, q_heads: int):
+        """Return ``(o_part, ml_part)`` view sliced to live dimensions."""
+        return (
+            self.decode_o_part[:batch, :q_heads],
+            self.decode_ml_part[:batch, :q_heads],
+        )
+
     def fill_input_ids(self, ids: "list[int]") -> Tensor:
         """Write ``ids`` into the device buffer and return ``[B]``.