fix: allocate inference workspace buffers outside inference mode

- workspace buffers are transport storage mutated in-place every step by the scheduler loop thread, which holds no ambient inference-mode context because torch.inference_mode is thread-local
- callers building the engine inside torch.inference_mode (scripts/tools/generate.py) produced inference tensors that reject off-thread in-place updates, crashing the first decode fill and aborting tasks after a single token
- force inference mode off around all workspace allocation so every buffer is a plain tensor regardless of caller context

astrai/inference/workspace.py

@@ -52,6 +52,16 @@ class InferenceWorkspace:
         self.device = device
         self.dtype = dtype
 
+        # Invariant: all workspace buffers are plain (non-inference)
+        # tensors.  The scheduler's loop thread mutates them in-place every
+        # step without any ambient inference-mode context — the mode is
+        # thread-local and the loop runs on its own thread.  Inference
+        # identity is fixed at construction and cannot be revoked later,
+        # so allocation must force the mode off: callers that build the
+        # engine inside ``torch.inference_mode()`` (e.g.
+        # ``scripts/tools/generate.py``) would otherwise create inference
+        # tensors that reject off-thread in-place updates.
+        with torch.inference_mode(False):
             # ``position_ids[:, None, None] >= arange`` RHS, reused every step.
             self.arange = torch.arange(max_seq_len, device=device)
             # Decode validity mask: [max_batch, 1, max_seq_len] bool.
@@ -64,46 +74,57 @@ class InferenceWorkspace:
             # CUDA-graph capture) and filled via a host staging buffer.  A
             # double buffer keeps a copy in flight from being overwritten by
             # the next fill.
-        self.input_ids = torch.empty((max_batch_size,), dtype=torch.long, device=device)
+            self.input_ids = torch.empty(
+                (max_batch_size,), dtype=torch.long, device=device
+            )
             self._pin = [
                 torch.empty((max_batch_size,), dtype=torch.long),
                 torch.empty((max_batch_size,), dtype=torch.long),
             ]
             self._pin_idx = 0
 
-        # KV-cache bind metadata (fixed shape, written by ``PagePool.bind_tasks``
+            # KV-cache bind metadata (fixed shape, written by
-        # when the Executor passes this workspace).  Stable addresses make the
+            # ``PagePool.bind_tasks`` when the Executor passes this
-        # decode forward CUDA-graph capturable.
+            # workspace).  Stable addresses make the decode forward
+            # CUDA-graph capturable.
             self.req_pool_indices = torch.empty(
                 (max_batch_size,), dtype=torch.int32, device=device
             )
-        self.seq_lens = torch.empty((max_batch_size,), dtype=torch.long, device=device)
+            self.seq_lens = torch.empty(
+                (max_batch_size,), dtype=torch.long, device=device
+            )
             self.kv_indptr = torch.empty(
                 (max_batch_size + 1,), dtype=torch.int32, device=device
             )
             self.qo_indptr = torch.empty(
                 (max_batch_size + 1,), dtype=torch.int32, device=device
             )
-        max_q_tiles = max_batch_size * ((max_seq_len + Q_TILE_ROWS - 1) // Q_TILE_ROWS)
+            max_q_tiles = max_batch_size * (
+                (max_seq_len + Q_TILE_ROWS - 1) // Q_TILE_ROWS
+            )
             self.q_tile_to_batch = torch.empty(
                 (max_q_tiles,), dtype=torch.int32, device=device
             )
             self.q_tile_to_index = torch.empty(
                 (max_q_tiles,), dtype=torch.int32, device=device
             )
-        self.inc = torch.arange(max_batch_size + 1, dtype=torch.int32, device=device)
+            self.inc = torch.arange(
+                max_batch_size + 1, dtype=torch.int32, device=device
+            )
             self.out_cache_loc = torch.empty(
                 (max_batch_size, 1), dtype=torch.int32, device=device
             )
 
-        # Per-step position IDs (must be at a fixed address for CUDA-graph capture).
+            # Per-step position IDs (must be at a fixed address for
+            # CUDA-graph capture).
             self.position_ids = torch.empty(
                 (max_batch_size,), dtype=torch.long, device=device
             )
 
-        # Split-KV partial-result buffers for decode (persistent, one global
+            # Split-KV partial-result buffers for decode (persistent, one
-        # alloc per process — mirrors FlashInfer's workspace pattern).
+            # global alloc per process — mirrors FlashInfer's workspace
-        # Shape: [max_batch_size, max_q_heads, _MAX_SPLITS, head_dim] (o_part)
+            # 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),