♻️ refactor: remove cross-layer pipeline and rename compute_chunked_prefill
- Remove cross-layer pipeline from OffloadEngine (saves ~1GB GPU memory for long sequences) - Delete layer_k/v_buffer_a/b double buffers - Remove start_decode_pipeline, get_decode_layer_kv, end_decode_pipeline methods - Remove pipeline state tracking variables - Simplify decode to use ring buffer pipeline only (more efficient for long sequences) - Rename compute_chunked_attention → compute_chunked_prefill for clarity - Add mandatory needle test requirements: --enable-offload --input-len 32768 Co-Authored-By: Claude Opus 4.5 <noreply@anthropic.com>
This commit is contained in:
@@ -644,12 +644,6 @@ class ModelRunner:
|
||||
# Get decode start position for accumulated token tracking
|
||||
decode_start_pos = self.kvcache_manager.get_decode_start_pos(seq)
|
||||
|
||||
# Get prefilled CPU blocks for pipeline initialization
|
||||
cpu_block_table = self.kvcache_manager.get_prefilled_cpu_blocks(seq)
|
||||
|
||||
# Start cross-layer pipeline (preloads Layer 0's data)
|
||||
offload_engine.start_decode_pipeline(cpu_block_table)
|
||||
|
||||
# Set up context for chunked decode
|
||||
set_context(
|
||||
is_prefill=False,
|
||||
@@ -666,9 +660,6 @@ class ModelRunner:
|
||||
logits = self.run_model(input_ids, positions, is_prefill=False)
|
||||
reset_context()
|
||||
|
||||
# End cross-layer pipeline
|
||||
offload_engine.end_decode_pipeline()
|
||||
|
||||
# Only offload when block is full (pos_in_block == block_size - 1)
|
||||
# This avoids unnecessary offloading on every decode step
|
||||
if pos_in_block == self.block_size - 1:
|
||||
|
||||
Reference in New Issue
Block a user