[RunningLeon]

Motivation

Reduce redundant recurrent/convolution state storage and checkpoint copies in the PyTorch engine for Qwen3.5 and GLM-5.3 Flash, while preserving speculative decoding, prefix caching, and asynchronous Mooncake Store lifetimes.

The legacy paths maintain runtime state and checkpoint storage separately and may copy state during checkpoint save/restore. They can also split prefill solely to reach a checkpoint boundary. A unified compact state pool lets runtime execution and checkpoints reference the same physical state, with explicit ownership preventing premature reuse.

This implementation is opt-in: set LMDEPLOY_ENABLE_COMPACT_STATE=1 before starting the Python process. Unset or 0 preserves legacy execution.

Modification

• Unify runtime recovery state and checkpoints in one physical pool; keep request-private auxiliary state separate.
• Allocate runtime state by execution phase and permit owner-validated input reuse rather than always reserving disjoint input/output banks.
• Add indexed conv/GDN/KDA kernels, device-side state-ID tracking, and address-stable CUDA Graph metadata.
• Export supported intermediate prefill checkpoints within a forward pass, avoiding checkpoint-only splits.
• Publish local checkpoints without copying state payloads and transfer compact state directly through Mooncake Store.
• Track producer, reader, and I/O ownership through eviction, cancellation, completion, sleep/wakeup, and weight reload.
• Remove redundant KDA decode Q/K/V copies, restrict the offline four-warp prefill optimization to H200 configurations, and avoid live-batch-driven Triton specialization in ID selection/staging.

Legacy execution remains the default. The environment switch derives compact capacity from the existing batch/checkpoint/Store budgets and MTP width; no test hook is required. Unsupported compact execution modes fail explicitly rather than silently falling back.

The current scope remains CUDA Qwen3.5 / GLM-5.3 Flash AR or matching single-layer MTP with UniExecutor/RayExecutor. Existing guards are unchanged: an explicit positive num_gpu_blocks is required; Mooncake Store also requires immutable compact_model_identity and, for MTP, compact_draft_identity in kv_connector_extra_config. PD, MemDecode, deprecated MP, other connectors, and Store with CPU KV paging are not supported.

Implementation references

• Implementation branch: RunningLeon/lmdeploy — feat/compact-unified-state.
• Core implementation: 3802ee441.
• Kernel optimizations and regression tests: 8401699c9.
• Architecture and ownership documentation: lmdeploy/pytorch/engine/cache_engine/README.md.

Performance results

GLM-5.3 Flash, TP4: upstream/main aba987ad2266d2fd2204e9d22d1ce43a5984dc91 versus compact 8401699c92a289441d1ba79aeab1f63a1c5b3a45.

• 1,024 requests per case; client concurrency 1,024.
• MTP5: maximum server batch size 64; target-only: maximum server batch size 256.
• Automatic prefix caching (APC) OFF/ON; Mooncake Store OFF.
• prefill_save_alignment=512; inputs share a 50% prefix.
• KV capacity is fixed across branches, APC settings, and input/output lengths within each MTP setting. AR and MTP use different batch/KV configurations.

All 16 cases completed: 16,384/16,384 measured requests passed exact-length and termination checks, with zero request errors.

MTP / Max batch	Input / Output tokens	APC	Output tok/s: main → compact	Throughput change	TTFT mean reduction	TPOT mean reduction	E2E mean reduction
5 / 64	16384 / 2048	OFF	1719.67 → 1730.63	+0.64%	0.81%	0.70%	0.80%
5 / 64	16384 / 2048	ON	1500.34 → 1830.46	+22.00%	18.73%	17.65%	18.62%
5 / 64	2048 / 16384	OFF	5726.62 → 5893.07	+2.91%	3.62%	2.23%	3.46%
5 / 64	2048 / 16384	ON	5217.31 → 5781.60	+10.82%	9.94%	9.44%	9.88%
0 / 256	16384 / 2048	OFF	1343.17 → 1470.27	+9.46%	7.52%	10.36%	8.31%
0 / 256	16384 / 2048	ON	1972.37 → 2734.51	+38.64%	29.01%	25.01%	27.88%
0 / 256	2048 / 16384	OFF	4048.94 → 5034.62	+24.34%	19.34%	19.10%	19.24%
0 / 256	2048 / 16384	ON	3659.97 → 4977.55	+36.00%	30.10%	23.12%	27.46%

Observed throughput improvements are 0.64%–22.00% with MTP5 and 9.46%–38.64% with target-only decoding in this matrix.

Interpretation and measurement boundaries

• One trial per case; paired TP4 services ran concurrently and shared CPU/I/O resources. All cases produced new compiler artifacts during measurement. These are end-to-end observations, not established steady-state or isolated kernel speedups.
• Both branches had zero prefix-cache hits for 2k inputs, including APC-on cases. Those gains cannot be attributed to prefix reuse.
• For 16k inputs with MTP5/APC-on, cached tokens increased from 917,504 to 1,368,064; the 22.00% gain includes different cache-hit volumes. For target-only/APC-on, both branches cached 8,380,416 tokens, so the 38.64% gain was observed at equal cache-hit volume.
• TTFT includes queueing and ends at the first nonempty text chunk; E2E ends at SSE DONE. TPOT is estimated as (E2E − TTFT) / (output_tokens − 1), not per-token ITL.
• Output lengths were enforced with ignore_eos=True; this performance workload is not an accuracy evaluation. Maximum batch size is a cap, not guaranteed active occupancy.

Validation and remaining scope

• Environment-switch follow-up (local changes after 8401699c9): 224 tests passed, including boolean parsing/worker-environment forwarding, AR/MTP capacity selection with APC/Store roles, fail-closed guards, and real CUDA allocation. This does not rerun full-model quality or serving performance.
• Environment-switch E2E follow-up: GLM-5.3 Flash TP4, 4/4 cases passed — pipeline AR with compact ON/OFF, and API server AR/MTP3 with compact ON; 12 requests completed without request errors. Every case reused 1,024 prefix tokens on the repeated long prompt. Actual pipeline state-manager/capacity checks and server startup logs confirm selection without configuration-injection hooks. This is a functional smoke test, not a new accuracy or performance study.
• Post-rebase component/kernel/configuration validation: 1,229 tests passed, plus GLM TP4 serving stability, directed Store GET, and sleep/wakeup checks.
• Final scoped regression: 353 tests passed, with no failures or skips, including real cross-process CUDA IPC, B64/H64 KDA decode, prefill boundary lengths, and dynamic-batch CUDA Graph tests. These counts describe separate validation rounds and are not summed as unique coverage.
• The earlier full GSM8K study passed the observed 0.5-percentage-point loss limit for Qwen3.5 and GLM-5.3 Flash, both AR and MTP3. It predates the final R2/R3 changes and is not relabeled as final-commit accuracy evidence. Store PUT was exercised, but not GET; Qwen had no actual prefix hits in that workload.
• Historical Qwen AR cold-tail regression and a DP8 non-graceful shutdown remain separate open items. This TP4 performance matrix does not close them or establish a complete RL-step speedup.
• Pre-commit checks passed. Broader hardware/topology validation remains necessary before considering default enablement.

Use cases (Optional)

The following examples use GLM-5.3 Flash, TP4, compact state ON, prefix caching ON, and Mooncake Store OFF. Run them separately on four available GPUs with sufficient memory. Use the branch containing the environment-switch follow-up and replace the model path with your local checkpoint directory.

export CUDA_VISIBLE_DEVICES=0,1,2,3
export LMDEPLOY_ENABLE_COMPACT_STATE=1
export MODEL_PATH=/path/to/GLM-5.3-Flash

Python pipeline

cat > glm53_compact_pipeline.py <<'PYTHON'
import os

from lmdeploy import GenerationConfig, PytorchEngineConfig, pipeline

backend_config = PytorchEngineConfig(
    tp=4,
    distributed_executor_backend='ray',
    dtype='bfloat16',
    language_model_only=True,
    max_batch_size=8,
    session_len=8192,
    num_gpu_blocks=2048,
    max_prefill_token_num=2048,
    enable_prefix_caching=True,
    prefix_cache_state_budget=8,
    prefill_save_alignment=512,
    block_size=64,
)
if __name__ == '__main__':
    with pipeline(os.environ['MODEL_PATH'], backend_config=backend_config,
                  log_level='INFO') as pipe:
        response = pipe('What is 2 + 2? Answer briefly.',
                        gen_config=GenerationConfig(max_new_tokens=512, do_sample=False))
        print(response.text)
PYTHON
python glm53_compact_pipeline.py

OpenAI-compatible API server

Using the same environment variables above:

lmdeploy serve api_server "$MODEL_PATH" \
  --backend pytorch \
  --tp 4 \
  --distributed-executor-backend ray \
  --dtype bfloat16 \
  --language-model-only \
  --max-batch-size 8 \
  --session-len 8192 \
  --num-gpu-blocks 2048 \
  --max-prefill-token-num 2048 \
  --enable-prefix-caching \
  --prefix-cache-state-budget 8 \
  --prefill-save-alignment 512 \
  --cache-block-seq-len 64 \
  --model-name glm53-flash \
  --server-name 127.0.0.1 \
  --server-port 23333 \
  --log-level INFO

Send a request from another terminal:

curl http://127.0.0.1:23333/v1/chat/completions \
  -H 'Content-Type: application/json' \
  -d '{
    "model": "glm53-flash",
    "messages": [{"role": "user", "content": "What is 2 + 2? Answer briefly."}],
    "max_tokens": 512,
    "temperature": 0
  }'

Optional MTP3: for the pipeline, import SpeculativeConfig from lmdeploy.messages and pass speculative_config=SpeculativeConfig(method='deepseek_mtp', num_speculative_tokens=3) to pipeline(...). For the server, append --speculative-algorithm deepseek_mtp --speculative-num-draft-tokens 3 to the launch command. Leave these options out for target-only decoding.

num_gpu_blocks=2048 is an explicit example capacity, not a universal hardware recommendation or the performance-table configuration. Compact currently requires a positive explicit KV block count; adjust it to the available memory. To use legacy execution, set LMDEPLOY_ENABLE_COMPACT_STATE=0 and restart the process. The startup log reports Enable unified compact state when the compact path is selected.

Checklist

• Pre-commit and linting checks passed.
• Unit and GPU regression tests cover the modified paths; validation scope and remaining gaps are documented above.
• Compatibility across all supported downstream dependency versions has not been exhaustively validated; current results are limited to the tested environment.
• Internal architecture documentation and relevant docstrings have been updated.