Skip to content

[Feature] GLM-5.3 internal prefill checkpoints - #5048

Closed
yidingcheng0206 wants to merge 4 commits into
InternLM:mainfrom
yidingcheng0206:feat/glm53-internal-checkpoint
Closed

yidingcheng0206 wants to merge 4 commits into
InternLM:mainfrom
yidingcheng0206:feat/glm53-internal-checkpoint

Conversation

@yidingcheng0206

Copy link
Copy Markdown
Collaborator

[Feature] GLM-5.3 internal prefill checkpoints

What changes

  • Enforce prefill_save_alignment for both internal and prompt-end local
    prefill checkpoints. Default 8192 is unchanged; smaller endpoint saves
    previously admitted are now rejected, including with IC OFF.
  • Add an optional native FlashKDA adapter using its existing
    checkpoint_state / checkpoint_offsets outputs.
  • Add opt-in internal checkpoint capture for GLM-5.3 AR and deepseek_mtp,
    preserving reservation, restore/pinning, publication and cleanup.
  • Retain MTP lookahead checks and absolute-position conv/recurrent ring writes.
    No FlashKDA, KPool, convolution, or TileLang decode kernel sources change.
IC OFF: forward [start,H) -> save S_H -> forward [H,E)
IC ON:  forward [start,E), exporting S_H; runtime advances to S_E

The same checkpoint H is preserved, rather than discarded for a faster run.
PR #5037 is already in main and is not reapplied.

Enablement and scope

export LMDEPLOY_KDA_PREFILL_BACKEND=flash_kda
export LMDEPLOY_INTERNAL_PREFILL_CHECKPOINT=1

Requires local prefix caching, BF16 inputs, FP32 recurrent state and
FlashKDA with the native checkpoint ABI. Validated FlashKDA:
ee0be888cd0e972f9409bf53756f8c38c6652173.

IC defaults to OFF; FLA remains the default prefill entry. Current scope is
exclusive single-sequence text prefill, one internal checkpoint, within the
existing prefill budget. Required endpoint saves and ineligible requests keep
the original path. IC does not support connectors/Mooncake, multimodal,
scoring/routed-expert output, or mixed local decode/prefill.
Intervals within the PCG cap are not merged; IC currently executes outside
prefill PCG while preserving decode graphs. Extra state export and larger
forwards have memory/time costs.

Validation

Submission head: a447f3768c5a7818563f85a40b9e4d5e55a0ec78.
Tested implementation: 170bc8b33b38412d0cbdb56b0e08ed041d804640 on
main 52e57d8f0d3a00fc53010f34bb5fe9109e73c560. The only change since
validation is documentation cleanup; implementation and tests are identical.

Check Result
Targeted CPU regression 366 passed, 9 skipped, 1 CUDA-dependent test deselected
CUDA restoration/ring/routing regression 13 passed, none skipped; non-MTP and K=1/3/5
TP1/DP8/EP8 model check, K=3 IC OFF/ON: 11 requests each; 22/22 final answers correct
Checkpoint evidence Expected split/merged shapes, save/restore and lookahead rejection passed
Decode graph Eight workers per service logged capture and use_graph=True

Model checks used identical data, native FlashKDA, BF16 KV, prefix ON, prefill16384,
alignment8192, PCG2048, C1, greedy/seed42 and natural EOS with a 512-token cap.
Full responses matched verbatim in only 4/11 pairs; this is not bitwise
equivalence or full IC accuracy acceptance. No separate replay trace was added.
Existing Transformers-version and MP-shutdown warnings were retained.
Both services stopped and GPUs were released.

No new full GSM8K, performance matrix or production-stress claim is made.
Historical performance results are not substituted for this submission-base check.

Review order

Four commits: alignment fix → native backend → non-MTP IC → MTP IC.
User-facing details: docs/en/advance/internal_checkpoint.md.

# named views for operator-owned block caches and configured state caches
block_caches: Mapping[str, torch.Tensor] | None = None
named_state_caches: Mapping[str, torch.Tensor] | None = None
state_capture: 'PrefillStateCapture | None' = None

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

for this type hint, should it be 'PrefillStateCapture' ?

@@ -0,0 +1,58 @@
# GLM-5.3 internal prefill checkpoints

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

please add zh_cn docs

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

3 participants