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)
This commit is contained in:
@@ -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,
|
||||
)
|
||||
|
||||
Reference in New Issue
Block a user