perf: preallocate inference decode buffers

- add InferenceWorkspace with fixed-shape per-step buffers (input_ids, decode mask, KV bind metadata) for CUDA-graph capture
- bind_tasks derives seq_lens from the pool's own _task_len tracking, dropping the seq_lens parameter
- update decode metadata in-place (position_ids, seq_lens, kv_indptr) instead of re-allocating per step
- task_extend advances _task_len in contiguous mode so the pool tracks current length
- skip log_softmax when logprobs are not requested

astrai/inference/sample.py

@@ -305,12 +305,12 @@ class SamplingPipeline(BaseSamplingStrategy):
             return tokens, chosen
 
         transformed = self.apply(logits, filter_value, input_ids, input_mask)
-        log_probs = torch.log_softmax(transformed.float(), dim=-1)
         tokens = torch.multinomial(
             torch.softmax(transformed, dim=-1), num_samples=1
         ).squeeze(-1)
         if not return_logprobs:
             return tokens
+        log_probs = torch.log_softmax(transformed.float(), dim=-1)
         chosen = torch.gather(log_probs, -1, tokens.unsqueeze(-1)).squeeze(-1)
         return tokens, chosen