Skip to content

feat(glm): coalesce GB10 continuation prefill with retained checkpoints - #699

Draft
FujitsuPolycom wants to merge 8 commits into
local-inference-lab:dev/jovian-judgementfrom
FujitsuPolycom:feat/gb10-continuation-prefill
Draft

FujitsuPolycom wants to merge 8 commits into
local-inference-lab:dev/jovian-judgementfrom
FujitsuPolycom:feat/gb10-continuation-prefill

Conversation

@FujitsuPolycom

@FujitsuPolycom FujitsuPolycom commented Sep 7, 2026

Copy link
Copy Markdown

Purpose

Extend the maintained single-checkpoint GLM/B12X path to up to four retained
KDA destinations with TP4 and DCP1/2/4, including aligned continuation chunks.
With 512-token physical cache blocks, scheduler alignment is 512/1024/2048
tokens and retention uses two/three/four destinations for DCP1/2/4 respectively.
The scheduler binds an exact query span and checkpoint positions; allocation
preserves physical state ownership; V2/GLM metadata carries the plan in packed
request order; recurrent and convolution exports use the planned destinations.
For an 8K prompt on the physical/lookup-512, scheduler-2048 grid, retained states
are at 4096, 6144, 7168, and 7680 tokens. B12X exports these states within one
bounded prefill invocation, avoiding extra full-model passes for their retention.

The implementation builds on the internal-checkpoint mechanism already present
in the review base: upstream PR #52789
added Kimi-K3/FlashKDA prefix/checkpoint/suffix execution inside one full-model
forward. The maintained GLM/B12X base already reserves and binds one internal
checkpoint. This extension supplies the four-destination B12X export contract and
GLM scheduler-to-allocator-to-metadata ownership needed by DCP4 retention and
MTP3 continuations. Related retention-aware scheduling PR #53479
is open but identifies its branch as a superseded draft; its chunk-stop and
speculative-backoff changes are separate from four-destination B12X export.

VLLM_B12X_KDA_PREFILL_COALESCING defaults off. The opt-in is bounded to the
documented GLM/GB10 BF16 TP4 configuration with DCP1, DCP2 or DCP4 with aligned native prefix
caching and at most 8,192 scheduled tokens. Unsupported, mixed-service,
cache-hit, and preempted cases retain their ordinary fallback rules. The
implementation requires B12X four-checkpoint export, independently of mHC.

Implementation base: 2a979314dc97b03173a0a76fc15664ec924db32b. Standalone implementation:
a6c8407645cf5e751711883f24e8822e73dba9a0; B12X dependency:
70fe41974ef4b18f61caaa2579c81cdc05d1265f. The coalescing branch contains no mHC feature.

Test Plan

Exercise scheduler/allocator ownership, full-prompt admission, retained prefix
boundaries, V2 packed metadata, GLM overrides, B12X binding, and one-slot
compatibility when the feature is disabled. Check actual model-derived cache
geometry through platform alignment, grouping, and KVCacheManager before GPU
state/history checks and serving measurements.

.venv/bin/python -m pytest \
  tests/v1/core/test_recurrent_prefill_checkpoint.py \
  tests/v1/attention/test_gdn_metadata_builder.py \
  tests/v1/worker/test_kda_prefill_checkpoint_binding.py -q
.venv/bin/python -m pytest tests/kernels/mamba/test_causal_conv1d.py \
  -k test_kda_checkpoint_history_excludes_speculative_cells -q

Test Result

62 checkpoint/allocator and worker-binding CPU tests passed at the standalone implementation revision above. Tests cover DCP1/2/4 geometry, retention, failed admission, cache-hit/preemption fallbacks and default-off API compatibility.

Status: implemented for the stated admission scope. Standalone validation uses Python 3.12/Torch 2.13.0+cpu on Windows with uvloop mapped to winloop; GPU operations in the suites use CPU doubles. Executed commands and DCP1/DCP2 evidence identify the complete test conditions. The complete standalone revisions have not run GPU/model validation.

The TP4/DCP1 and TP4/DCP2 serving comparison toggles coalescing and mHC together. It qualifies bounded combined-runtime checks and cannot isolate either feature's speedup. DCP1 passes the fixed-input on/off and cached/cold tolerances. DCP2 cache-invariant probabilities remain research-only: enabled cached/cold logprob difference 0.127015 exceeds its bound, despite identical answer tokens and passing matched on/off comparisons.

TP4/DCP4 component and serving evidence

The following results belong to the explicitly identified DCP4 source composition; they are separate from the standalone DCP1/2 CPU tests and combined DCP1/2 measurements.

66 tests passed on coalescing revision 1bf41ecb19ff82bd4bca1195cd1ab8f63e7e0717. The combined CPU suite passed
115 tests. Six selected convolution-history GPU tests passed, together with
12 B12X four-checkpoint GPU tests. The model-derived CPU layout check reproduces
the generic 2,560-token block expansion and the native 512-token split layout;
the latter produces one 8K scheduler span with four retained checkpoint targets.

The serving comparison uses four 48-SM NVIDIA GB10 GPUs (compute capability
12.1), GLM-5.3-Flash-NVFP4-Spark weights, BF16 activations, FP8 KV cache,
TP4/DCP4/PP1, MTP3, and an 8,192-token batch ceiling. Native prefix caching and
the fused SparkRing transport configuration are held constant across arms;
the mHC collectives use the existing TP PyNccl communicator. There is no
SparkCache KV connector in this comparison.

Feature-arm values for VLLM_B12X_KDA_PREFILL_COALESCING and
VLLM_GLM53_MHC_PREFILL_SHARD, respectively, are: neither 0/0, coalescing-only
1/0, mHC-only 0/1, and both 1/1.

All arms use these native layout settings:

export VLLM_GLM53_SPLIT_TARGET_BLOCK_SIZE=512
export VLLM_GLM53_SPLIT_MAMBA_BLOCK_SIZE=512

Together with aligned Mamba caching and retention interval zero, they establish
the tested physical/lookup/scheduler grid (512, 512, 2048). CLI block sizes
alone do not preserve that grid: generic hybrid page alignment can raise both
physical block sizes to 2,560, preventing the intended 8K execution path.

The measured image is
sha256:52b207e716a285c16e5e1b14ec2a41f6208b9450c617e7d1cde5a507ea879d7f,
containing vLLM abb715f132bdccb592a34b2596a3d3a8d757ffbc and B12X
70fe41974ef4b18f61caaa2579c81cdc05d1265f. The isolated vLLM PRs compose to this
runtime source; their documentation/evidence commits are not the image revision.
The toolchain includes Python 3.12.3, PyTorch 2.13.0+cu130, Triton 3.7.1,
CUTLASS DSL 4.6.2, and FlashInfer 0.6.17.

Three shape warmups and five exact-answer/cache checks precede the measured
requests. Each prompt size has three unique cold samples with explicit zero
cached tokens. Rates below are prompt tokens divided by median API TTFT, using
the client's QueryPerformanceCounter clock; they include API/network overhead.
The five correctness probes cover cold 8K/16K/32K prompts and repeated/extended
8K prompts with positive cache reuse. All completed arms passed those probes.

VLLM_GLM53_MHC_PREFILL_DIAGNOSTICS=1 is common to all arms. Activation is checked
before timing and again after completion: checkpoint-enabled arms show capacity
four, the required grid, and an 8,192-token dispatch with four checkpoint targets.
The combined and mHC-only arms have all-rank request-time mHC enqueue records
with actual 2,048-row repeated mHC outputs and 90 reduce-scatters/90 gathers per
eligible base forward.
Completed requests corroborate asynchronous completion; logs alone do not prove
GPU state equivalence.

Comparison scope

The “Neither” column is the flag-disabled reference within the specified TP4/DCP4 ring configuration. It retains 512-token target/recurrent cache blocks, aligned checkpoint retention with interval zero, and VLLM_B12X_MLA_CKV_GATHER=1 (VLLM_B12X_MLA_CKV_GATHER_MAX_TOKENS=524288) in every arm. It is not a benchmark of the default dev/jovian-judgement configuration or a comparison against the fastest independently tuned configuration. The results do not establish superiority over a switched TP4/DCP1 setup. Disabling full-CKV gather selects another DCP attention communication algorithm; it is not equivalent to removing communication.

Observed prefill tokens/s:

Prompt tokens Neither Coalescing only mHC only Both
8,192 1,949.6 2,947.3 1,958.5 3,063.5
16,384 2,358.9 2,928.3 2,390.8 3,060.8
32,768 2,614.5 2,923.5 2,681.8 3,049.6

Coalescing-only rates are 51.18%, 24.14%, and 11.82% higher than neither at
8K, 16K, and 32K, respectively (100 * (coalescing_only / neither - 1)). The combined
row includes the separate mHC transformation and must not be attributed solely
to coalescing.

Observed aggregate decode tokens/s:

Decode cell Neither Coalescing only mHC only Both
C1 / 8K 49.17 47.37 48.46 46.29
C1 / 32K 46.18 45.92 47.31 47.11
C4 / 8K 122.19 120.18 121.17 119.37
C4 / 32K 123.44 125.20 123.90 124.25

Observed MTP-normalized decode steps/s (mean acceptance length in parentheses):

Decode cell Neither Coalescing only mHC only Both
C1 / 8K 17.86 (2.75) 17.91 (2.65) 17.86 (2.71) 17.89 (2.59)
C1 / 32K 17.82 (2.59) 17.73 (2.59) 17.62 (2.68) 17.57 (2.68)
C4 / 8K 43.70 (2.80) 44.54 (2.70) 44.99 (2.69) 42.99 (2.78)
C4 / 32K 44.34 (2.78) 45.64 (2.74) 45.84 (2.70) 45.00 (2.76)

These prefill features do not select their optimized paths during captured decode. Normalized measurements show no consistent decode speedup across cells; the short sequential runs do not establish performance equivalence.

These are sequential arm runs, not interleaved A/B trials. Separate JIT
namespaces and memory-recovery interventions between runs remain recorded
conditions. Three TTFT samples do not establish a confidence interval. Decode
uses one 20-second window per C1/C4 by 8K/32K cell, a five-second warmup, and at
most 1,024 output tokens; acceptance length varies, so the raw bundle also
retains MTP-normalized steps/s. Those short windows do not establish a decode
speedup or absence of a regression.

Status: implemented; research-only serving qualification. All four prefill
arms passed activation and exact-answer/cache checks. All four decode arms
completed their four cells without reported request errors.
The completed 8K-32K smoke coverage does not establish full numerical/model-quality
equivalence, long-duration reliability, or 64K/128K behavior for this source
composition. GSM8K/MRCR evaluation remains unrun.

Evidence: method and conditions,
results report,
and raw samples and activation evidence.
AI assistance was used for implementation, review, and test preparation. Submitted by FujitsuPolycom as a draft for maintainer review.

Reproduction

Model/tokenizer: local-inference-lab/GLM-5.3-Flash-NVFP4-Spark at df116c4fb16b1d37ae43d2cfd624de26ffbc832e. Reproduction instructions include the public source composition, exact decode-client reconstruction, serving settings and standalone prefill client. The prefill client preserves the measured prompt/timing/cache-check algorithm; runtime feature activation still requires separate request-associated worker logs. The reproduction instructions in that directory target the identified TP4/DCP4 measurement composition.

Independent validation requested

Testing on four GB10 Sparks connected through a network switch would help establish portability beyond the ring deployment measured here. Please report exact source/model revisions, TP4 with the selected DCP1/2/4 setting, request-associated feature activation, correctness/cache checks, and cold 8K/16K/32K TTFT with both flags disabled and enabled separately. Include raw and MTP-normalized decode rates. TP2 is not qualified by this submission.

Companion changes

Requires B12X #338 for the four-checkpoint API. vLLM #700 is an independent mHC optimization tested alongside this feature; it is not a dependency.

FujitsuPolycom and others added 6 commits September 7, 2026 00:43
Export required recurrent states inside aligned prefill chunks up to 8192
tokens. Preserve retained state and append-only worker block tables across
cold chunk continuations, including full-prompt admission accounting.

Carry scheduler checkpoint plans through model runner V2 and the GLM
metadata adapter. Bind B12X two-checkpoint metadata only for the opt-in,
and store causal convolution history independently of speculative capacity.

The feature defaults off, targets GB10 GLM TP4/DCP4, and retains ordinary
cache-hit, mixed-service, preemption and one-checkpoint B12X behavior.
Status: research-only; GPU and full-model evaluation are not executed.

Validation: 107 CPU tests pass; four GPU checkpoint-store tests collect.
Repository checks pass except a mypy error reproduced on the source base;
the Docker graph hook passes through Git Bash on Windows.

Co-authored-by: OpenAI Codex

Signed-off-by: FujitsuPolycom <87842395+FujitsuPolycom@users.noreply.github.com>
Reserve four checkpoint slots so 512-token prefix hits and 2048-token scheduler fallback states fit in one 8192-token forward. Preserve all retention destinations, default-off one-checkpoint bindings, and speculative slot ownership. Log configured grids and scheduled coalescing separately.

DCP4 scheduler, allocator, GDN metadata and warmup CPU contracts pass. Four-slot GPU and model qualification remain pending.

Co-authored-by: Codex <noreply@openai.com>
Signed-off-by: FujitsuPolycom <87842395+FujitsuPolycom@users.noreply.github.com>
Document the maintained GLM split-page settings needed to preserve physical 512-token blocks, 512-token lookup alignment and 2048-token scheduler alignment. Hold this cache layout constant across benchmark arms and require scheduled-path evidence.

Co-authored-by: Codex <noreply@openai.com>
Signed-off-by: FujitsuPolycom <87842395+FujitsuPolycom@users.noreply.github.com>
Signed-off-by: FujitsuPolycom <87842395+FujitsuPolycom@users.noreply.github.com>
Signed-off-by: FujitsuPolycom <87842395+FujitsuPolycom@users.noreply.github.com>
Signed-off-by: FujitsuPolycom <87842395+FujitsuPolycom@users.noreply.github.com>
@coderabbitai

coderabbitai Bot commented Sep 7, 2026

Copy link
Copy Markdown

Important

Draft PR not reviewed

Draft PRs are not automatically reviewed by default.

  • Trigger a manual review

To automatically review draft PRs, update your CodeRabbit configuration:

reviews:
  auto_review:
    drafts: true

Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out.

❤️ Share

Comment @coderabbitai help to get the list of available commands.

@FujitsuPolycom

Copy link
Copy Markdown
Author

Prefill measurements with dual-domain NCCL

Status: research-only performance observations from a composed runtime.

The deployment uses four NVIDIA DGX Sparks serving local-inference-lab/GLM-5.3-Flash-NVFP4-Spark revision df116c4fb16b1d37ae43d2cfd624de26ffbc832e, with TP4/DCP4, native MTP3, BF16 activations, FP8 KV, an 8,192-token scheduler budget, and 512-token target/recurrent cache blocks.

All configurations retain:

  • Full-CKV gather.
  • DCP top-k owner merge with fused mesh endpoints.
  • A hardware-forwarded RoCE ring.
  • Patched NCCL using NIC functions across both host PCIe domains, with Ring and four channels.

SparkCache and compact index-cache gathering are disabled. Only continuation coalescing and mHC token sharding change between columns.

Cold prefill throughput, tokens/s

Prompt tokens Coalescing OFF
mHC sharding OFF
Coalescing ON
mHC sharding OFF
Coalescing ON
mHC sharding ON
8,192 1,893.7 2,938.8 3,406.0
16,384 2,296.3 2,914.1 3,334.7
32,768 2,558.8 2,895.5 3,326.6

Each value is prompt tokens divided by median client-observed TTFT from three cold samples. Startup and request warmups are excluded. Runtime activation checks and bounded exact-answer checks passed.

Coalescing improves throughput at every measured prompt length. Adding mHC token sharding to coalescing provides an additional 15.9%, 14.4%, and 14.9% at 8K, 16K, and 32K, respectively.

MTP-normalized decode observations, 8,192-token context

Concurrency Coalescing OFF
mHC sharding OFF
Coalescing ON
mHC sharding OFF
Coalescing ON
mHC sharding ON
1 19.40 19.59 19.23
4 46.38 46.07 45.74

Values are derived steps/s using observed speculative counters. Each cell contains one 10-second observation. Decode remains close across configurations; these short measurements do not establish performance equivalence or a repeatable decode gain.

Runtime identity

  • vLLM: 54842c325af3526ad723adf134112b5bb40cedf5
  • B12X: 0b6d61c37c87ae49d2f9d20d38b9da023146e243
  • Docker image: sha256:c389c16e50a5c798e0aa9eec953125f2d69e06f50c90f00a3e518d2db79b42b2

These measurements describe that composed runtime, not qualification of the exact PR heads. Configurations ran sequentially rather than as interleaved A/B trials. The comparison does not isolate CKV gather, top-k owner merge, or dual-domain NCCL, and does not include a coalescing-OFF/mHC-ON configuration.

FujitsuPolycom and others added 2 commits September 8, 2026 11:39
Retain the default-off checkpoint path and its four-slot B12X dependency while accepting DCP1, DCP2 and DCP4. Derive scheduler grids from physical cache geometry and cover checkpoint ownership, failed admission, cache hits and preemption at each supported DCP size.

Validation: 62 CPU checkpoint and binding tests passed on this standalone branch. GPU measurements exist for combined source 8f8ea47; this ported head has not executed GPU or full-model validation.

Co-authored-by: Codex <noreply@openai.com>

Signed-off-by: FujitsuPolycom <87842395+FujitsuPolycom@users.noreply.github.com>
Record source-bound prefill and normalized decode observations, checkpoint and collective invariants, and standalone CPU test commands. Distinguish combined-runtime qualification from standalone GPU validation and retain the unresolved DCP2 cached-versus-cold numerical bound failure.

Validation: 62 checkpoint and binding CPU tests passed; documentation checks passed. Dockerfile graph check passed through Git Bash on Windows.

Co-authored-by: Codex <noreply@openai.com>
Signed-off-by: FujitsuPolycom <87842395+FujitsuPolycom@users.noreply.github.com>
@FujitsuPolycom

Copy link
Copy Markdown
Author

TP4 prefill support with DCP1 and DCP2

Status: implemented; bounded combined-runtime checks are qualified under the conditions below. DCP2 cache-invariant probabilities remain research-only.

Continuation coalescing accepts TP4 with DCP1/2/4. With 512-token physical blocks, scheduler alignment is 512/1024/2048 tokens and retention uses two/three/four destinations respectively. Four-slot B12X capacity, default-off behavior and allocator/fallback invariants are retained. This PR contains no mHC implementation and continues to require B12X #338.

Implementation: a6c8407645cf. 62 checkpoint/allocator and worker-binding tests passed on Python 3.12/Torch 2.13.0+cpu on Windows, using winloop as the uvloop compatibility alias. GPU operations in these suites use CPU doubles.

Serving conditions and evidence

Four NVIDIA DGX Sparks serve local-inference-lab/GLM-5.3-Flash-NVFP4-Spark revision df116c4fb16b1d37ae43d2cfd624de26ffbc832e, using TP4, native MTP3, BF16 activations, FP8 KV, an 8,192-token scheduler budget and 512-token target/recurrent cache blocks. The two flags VLLM_B12X_KDA_PREFILL_COALESCING and VLLM_GLM53_MHC_PREFILL_SHARD are both OFF in controls and both ON in enabled configurations.

The hardware-forwarded RoCE ring and dual-domain NCCL remain enabled. CKV gathering uses the two-rank DCP subgroup at DCP2 and bypasses at DCP1. Cache interleave is 4 at DCP2 and 1 at DCP1. SparkCache, compact index cache and the DCP4-only top-k owner exchange/fused endpoints are disabled throughout.

Cold prefill throughput, tokens/s

Prompt tokens DCP1, both OFF DCP1, both ON DCP2, both OFF DCP2, both ON
8,192 2,263.0 3,616.3 1,938.7 3,493.9
16,384 2,615.7 3,615.1 2,345.5 3,461.8
32,768 2,828.4 3,586.1 2,619.0 3,445.9

Each value is prompt tokens divided by median client-observed TTFT from three cold samples after excluded warmups. Six exact-answer/cache checks and request-associated feature activation checks passed in each configuration. Same-DCP comparisons measure the combined feature effect; they do not isolate coalescing or mHC, and cross-DCP comparisons also change the attention/cache algorithm.

MTP-normalized decode observations, steps/s

Context / concurrency DCP1, both OFF DCP1, both ON DCP2, both OFF DCP2, both ON
8,192 / C1 21.17 21.40 19.73 20.05
32,768 / C1 21.23 21.17 19.86 19.80
8,192 / C4 47.93 48.00 45.41 46.62
32,768 / C4 46.21 49.93 45.01 44.47

Values divide aggregate emitted tokens/s by observed speculative acceptance length. C1/C4 denotes one/four concurrent requests. Each cell is one approximately 10-second observation at temperature 1 with a 2,048-token output cap. Configurations ran sequentially; these observations establish neither repeatable decode gains nor performance equivalence.

Numerical scope

The combined runtime passed 18 GB10 recurrent-export/reload and convolution-history component tests with zero skips. Fixed-input fixtures produced identical answer tokens and passed all six matched on/off logprob comparisons using abs(on-off) <= 0.1 + 0.02*abs(off).

DCP1 also passed cached-versus-cold comparisons. DCP2 cache-invariant probabilities remain research-only: the enabled cached/cold maximum logprob difference is 0.127015, outside its diagnostic bound; the disabled difference is 0.018938. The DCP2 extended-cold on/off difference of 0.100147 narrowly passes its approximately 0.100867 bound. Matching answers and on/off passes do not resolve the cached/cold failure. No broad accuracy or isolated PR-head GPU qualification is claimed.

Runtime identity

  • vLLM: 8f8ea47be212bbdd91b2172d5958ea2aae2b0e50
  • B12X: 0b6d61c37c87ae49d2f9d20d38b9da023146e243
  • Docker image ID: sha256:e585e8b3ebeeea2852011d319f4a4d781bf2ae33388210b7cd16ff16d3765d9b

These identities describe the measured combined runtime. Each standalone code revision has CPU coverage; neither complete standalone revision has GPU/model validation. AI assistance was used for implementation, testing and evidence preparation.

Evidence: DCP1/DCP2 conditions, test commands and limitations and numeric evidence and receipt hashes.

@original-el8

Copy link
Copy Markdown

Two follow-ups to the continuation-coalescing plan, measured on four DGX Sparks (GLM-5.3-Flash-NVFP4-Spark, TP4/DCP4, MTP3, 8,192 budget, 512-token blocks), are up as a draft against this PR's branch: FujitsuPolycom#14

  1. Let prefill_checkpoint_plan accept a final span that ends at the prompt end off the 512 grid. A 16,228-token prompt then schedules as 8192 + 8036 instead of 8192, 4096, 2048, 1024, 512, 356. Cold TTFT at 16,228 tokens: 5.67 s -> 4.83 s (-14.8%) in an A/B/A/B; 16,383 -5.6%, 32,319 -3.2%, aligned controls flat.
  2. Admit coalescing when at most one request still has prompt tokens to compute rather than when the engine holds exactly one request, so decoding requests do not push every new prompt down the ladder. Cold prefill beside three decoding 8K requests: 8,192 tokens 3.41 s -> 2.71 s (-20.6%), 16,228 tokens 6.02 s -> 5.31 s (-11.8%).

Answer and prefix-cache checks passed in every arm; the checkpoint test file gains cases for the unaligned final plan, its DCP4 allocator columns, the single-token tail and prompt-only exclusivity (60 passed on this head). Happy to move this wherever you prefer it to land.

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.

2 participants