Offline generate_batch lacks token-level chunked prefill → task-ring heap deadlock (507018) on long prompts

[ChaoWao]

Summary

Offline batch generation (LLMEngine.generate_batch → _generate_batch_impl)
does not perform token-level chunked prefill. Each prompt is prefilled in a
single run_prefill call spanning the whole prompt, so for long prompts the
prefill_fwd per-layer staging tensors (which scale with tokens-per-call)
overflow the task-ring heap and trip a Task Allocator Deadlock
(host-side 507018). The async serving path already chunks long prefills
(scheduler.py: enable_chunk_prefill / long_prefill_token_threshold); the
offline path never wired this up.

Repro

Qwen3-14B, a2a3, real batch-16 (16 independent requests), 3338-token prompt
each, --max-num-batched-tokens 1024, PTO2_RING_HEAP=2147483648 (2 GB):

python examples/model/qwen3_14b/npu_generate.py \
  --model-dir /data/models/Qwen3-14B --prompt "<3338-token prompt>" \
  --platform a2a3 --device-id $DEV \
  --max-seq-len 4096 --max-new-tokens 20 --batch-size 16 \
  --max-num-batched-tokens 1024 --profile

Warmup (small tokens) passes; the first real prefill deadlocks:

FATAL: Task Allocator Deadlock  (pto_ring_buffer.h)
  BLOCKED: tasks=3426/131072, heap_used=2143289344/2147483648,
           heap_available=4194304, on=heap        # 99.8% full, stuck
  Heap ring 3: Requested: 11141120 bytes; No reclaim progress for ~500 ms
  Solution: Increase heap (current: 2147483648)

Root cause (from generated prefill_fwd.cpp)

The per-layer staging is sized by toks_pad = ceil(prefill_tokens/128)*128,
i.e. it scales with tokens per run_prefill call:

tensor	shape (3338 tok → toks_pad 3456)	dtype	size/layer
resid1_all	[3456, 5120]	FP32	67.5 MiB
mlp_out_acc	[3456, 5120]	FP32 (manual_dep)	67.5 MiB
post_norm_all	[3456, 5120]	BF16	33.8 MiB
layer_next_hidden	[3338, 5120]	BF16	32.6 MiB

≈ 200 MiB live staging per layer. The 40-layer graph runs ~10–12 layers ahead
of task retirement in the FIFO ring, so the heap fills to ~2.14 GB and
head-of-line reclaim stalls. Since every alloc_* is toks_pad-driven,
tokens-per-call is the only knob that bounds the ring footprint — request-level
grouping (one whole request per call) does not help a single long prompt.

Fix

Wire token-level chunked prefill into _generate_batch_impl: split each
prompt into <= max_num_batched_tokens-token sub-chunks and prefill
incrementally (KV carry-over via the kernel's chunk_lens / chunk_offsets),
reusing the async scheduler's long_prefill_token_threshold mechanism. This
bounds per-call prefill_tokens (e.g. 512 → ~10 MiB/layer staging) and keeps
the ring far below 2 GB.

Environment

• pypto-serving: chunk-offline-batch-prefill (offline request-level grouping;
  the token-level chunking is missing)
• pypto-lib: 4ee8cf4 (#765, Qwen3 CANN FAI static batch-16)
• pypto: ccbf2870
• simpler runtime: worktree feat branch (1e24e1f6)