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
This commit is contained in:
@@ -52,77 +52,98 @@ class InferenceWorkspace:
|
|||||||
self.device = device
|
self.device = device
|
||||||
self.dtype = dtype
|
self.dtype = dtype
|
||||||
|
|
||||||
# ``position_ids[:, None, None] >= arange`` RHS, reused every step.
|
# Invariant: all workspace buffers are plain (non-inference)
|
||||||
self.arange = torch.arange(max_seq_len, device=device)
|
# tensors. The scheduler's loop thread mutates them in-place every
|
||||||
# Decode validity mask: [max_batch, 1, max_seq_len] bool.
|
# step without any ambient inference-mode context — the mode is
|
||||||
self.input_mask = torch.empty(
|
# thread-local and the loop runs on its own thread. Inference
|
||||||
(max_batch_size, 1, max_seq_len), dtype=torch.bool, device=device
|
# 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.
|
||||||
|
self.input_mask = torch.empty(
|
||||||
|
(max_batch_size, 1, max_seq_len), dtype=torch.bool, device=device
|
||||||
|
)
|
||||||
|
|
||||||
# Per-step token IDs. Values come from host Python lists every
|
# Per-step token IDs. Values come from host Python lists every
|
||||||
# step, so the device buffer is pre-allocated (stable address for
|
# step, so the device buffer is pre-allocated (stable address for
|
||||||
# CUDA-graph capture) and filled via a host staging buffer. A
|
# CUDA-graph capture) and filled via a host staging buffer. A
|
||||||
# double buffer keeps a copy in flight from being overwritten by
|
# double buffer keeps a copy in flight from being overwritten by
|
||||||
# the next fill.
|
# the next fill.
|
||||||
self.input_ids = torch.empty((max_batch_size,), dtype=torch.long, device=device)
|
self.input_ids = torch.empty(
|
||||||
self._pin = [
|
(max_batch_size,), dtype=torch.long, device=device
|
||||||
torch.empty((max_batch_size,), dtype=torch.long),
|
)
|
||||||
torch.empty((max_batch_size,), dtype=torch.long),
|
self._pin = [
|
||||||
]
|
torch.empty((max_batch_size,), dtype=torch.long),
|
||||||
self._pin_idx = 0
|
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
|
||||||
self.req_pool_indices = torch.empty(
|
# CUDA-graph capturable.
|
||||||
(max_batch_size,), dtype=torch.int32, device=device
|
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.kv_indptr = torch.empty(
|
self.seq_lens = torch.empty(
|
||||||
(max_batch_size + 1,), dtype=torch.int32, device=device
|
(max_batch_size,), dtype=torch.long, device=device
|
||||||
)
|
)
|
||||||
self.qo_indptr = torch.empty(
|
self.kv_indptr = torch.empty(
|
||||||
(max_batch_size + 1,), dtype=torch.int32, device=device
|
(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)
|
self.qo_indptr = torch.empty(
|
||||||
self.q_tile_to_batch = torch.empty(
|
(max_batch_size + 1,), dtype=torch.int32, device=device
|
||||||
(max_q_tiles,), dtype=torch.int32, device=device
|
)
|
||||||
)
|
max_q_tiles = max_batch_size * (
|
||||||
self.q_tile_to_index = torch.empty(
|
(max_seq_len + Q_TILE_ROWS - 1) // Q_TILE_ROWS
|
||||||
(max_q_tiles,), dtype=torch.int32, device=device
|
)
|
||||||
)
|
self.q_tile_to_batch = torch.empty(
|
||||||
self.inc = torch.arange(max_batch_size + 1, dtype=torch.int32, device=device)
|
(max_q_tiles,), dtype=torch.int32, device=device
|
||||||
self.out_cache_loc = torch.empty(
|
)
|
||||||
(max_batch_size, 1), 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.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
|
||||||
self.position_ids = torch.empty(
|
# CUDA-graph capture).
|
||||||
(max_batch_size,), dtype=torch.long, device=device
|
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, 2] (ml_part)
|
# [max_batch_size, max_q_heads, _MAX_SPLITS, head_dim] (o_part)
|
||||||
self.decode_o_part = torch.empty(
|
# [max_batch_size, max_q_heads, _MAX_SPLITS, 2] (ml_part)
|
||||||
(max_batch_size, max_q_heads, _MAX_SPLITS, head_dim),
|
self.decode_o_part = torch.empty(
|
||||||
dtype=torch.float32,
|
(max_batch_size, max_q_heads, _MAX_SPLITS, head_dim),
|
||||||
device=device,
|
dtype=torch.float32,
|
||||||
)
|
device=device,
|
||||||
self.decode_ml_part = torch.empty(
|
)
|
||||||
(max_batch_size, max_q_heads, _MAX_SPLITS, 2),
|
self.decode_ml_part = torch.empty(
|
||||||
dtype=torch.float32,
|
(max_batch_size, max_q_heads, _MAX_SPLITS, 2),
|
||||||
device=device,
|
dtype=torch.float32,
|
||||||
)
|
device=device,
|
||||||
|
)
|
||||||
|
|
||||||
# Decode output buffer (graph-safe pre-alloc). Shape matches the
|
# Decode output buffer (graph-safe pre-alloc). Shape matches the
|
||||||
# decode kernel's output: [batch, q_head, head_dim].
|
# decode kernel's output: [batch, q_head, head_dim].
|
||||||
self.decode_out = torch.empty(
|
self.decode_out = torch.empty(
|
||||||
(max_batch_size, max_q_heads, head_dim),
|
(max_batch_size, max_q_heads, head_dim),
|
||||||
dtype=dtype,
|
dtype=dtype,
|
||||||
device=device,
|
device=device,
|
||||||
)
|
)
|
||||||
|
|
||||||
def fill_input_ids(self, ids: "list[int]") -> Tensor:
|
def fill_input_ids(self, ids: "list[int]") -> Tensor:
|
||||||
"""Write ``ids`` into the device buffer and return ``[B]``.
|
"""Write ``ids`` into the device buffer and return ``[B]``.
|
||||||
|
|||||||
Reference in New Issue
Block a user