feat(kda): export up to four retained prefill checkpoints on GB10 - #338
FujitsuPolycom wants to merge 6 commits into
Conversation
Save two checkpoint states from the running FP32 recurrence using planned checkpoint capacity and runtime offsets. Preserve the one-checkpoint vector contract by default and validate matrix metadata before cross-index dtype agreement. Use query schema 2 and plan-time GB10 device eligibility. Checkpoint export is research-only: GPU qualification and KDA catalog/provider integration are required. No measured profile or throughput claim is supplied. Validation: 33 focused CPU contract tests pass; Ruff and syntax checks pass. Fifteen GPU tests collect but have not executed, including graph replay and state-pool offsets beyond the signed 32-bit element boundary.
Use the shared GPU runner to forward each input's required gate lower bound. Check transactional poisoning across the full bound output capacity while retaining nonzero error-code and exact recurrent-state preservation assertions. Runtime code is unchanged.
Accept planned capacities one, two and four while preserving vector metadata by default. Validate all active offset pairs and keep multi-checkpoint eligibility on GB10. Query schema 3 separates the four-slot contract from existing profiles. CPU contracts pass; four-slot GPU qualification remains pending. 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>
|
Important Draft PR not reviewedDraft PRs are not automatically reviewed by default.
To automatically review draft PRs, update your CodeRabbit configuration: reviews:
auto_review:
drafts: trueThanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
|
Companion serving measurements Status: research-only. A four-DGX-Spark deployment serving All configurations retain dual-domain NCCL, full-CKV gather, DCP top-k owner merge, and fused mesh endpoints. SparkCache and compact index-cache gathering are disabled. Cold prefill throughput — tokens/s
Each value is prompt tokens divided by median client-observed TTFT from three cold samples, excluding warmup. Runtime activation and bounded exact-answer checks passed. These are complete-serving measurements, not isolated checkpoint-export kernel timings. Configurations ran sequentially rather than as interleaved A/B trials. The results describe a composed runtime, not qualification of the exact PR heads. The runtime identities, measurement conditions, and normalized decode observations are documented in the companion coalescing PR #699 and mHC sharding PR #700. |
|
Checkpoint-export integration at TP4/DCP1 and TP4/DCP2 Status: implemented checkpoint capacity; qualified for the bounded combined-runtime checks documented in the companion vLLM PRs. DCP2 cache-invariant probabilities remain research-only. B12X checkpoint export requires no DCP-specific kernel change. The kernel receives local head geometry, packed token ranges, checkpoint destinations and planned capacity; vLLM determines distributed cache boundaries. TP4 with 512-token physical blocks uses two/three/four active destinations for DCP1/2/4 while retaining capacity four. Three active destinations do not require a capacity-three API. The six executable/metadata files under The serving runtime uses four DGX Sparks, GLM-5.3-Flash-NVFP4-Spark revision Cold prefill throughput, tokens/s
Each value is prompt tokens divided by median TTFT from three cold samples, excluding warmup. Configurations ran sequentially. These are combined serving measurements, not isolated checkpoint-kernel timings or GPU qualification of the complete B12X PR revision. The measured runtime passed 18 GPU recurrent-export/reload and convolution-history component tests with zero skips; this count includes companion vLLM coverage. All four serving configurations passed six exact-answer/cache checks. DCP1 passes fixed-input on/off and cached/cold probability comparisons. DCP2's enabled cached/cold logprob difference of 0.127015 exceeds its bound (disabled: 0.018938), despite matching answer tokens and passing matched on/off comparisons. That numerical limitation is unresolved. Runtime: vLLM Evidence: DCP1/DCP2 conditions, test commands and limitations and numeric evidence and receipt hashes. |
1 participant
Purpose
Export up to four retained interior recurrent states from one planned KDA
prefill on GB10. Callers can preserve multiple prefix states inside a query
without ending the query at each checkpoint.
Caps.max_checkpointsretains itsdefault of one and vector metadata; explicit capacities two/four use contiguous
[sequence_capacity, max_checkpoints]offsets and destination indices withtransactional validation. The plan owns workspace capacity, and live sequence
lengths remain runtime inputs. Invalid checkpoint metadata must preserve the
state pool; pool-scaled offsets use 64-bit addressing.
This implementation is separate from vLLM scheduling and token-sharded mHC.
The extended interface is bounded to SM121, 48-SM GB10 devices. It depends on no
private host inventory or transport plugin.
Implementation base:
06b4de7c723e6f166d65abf5909c5b7d0f8acc68. Runtime implementation:70fe41974ef4b18f61caaa2579c81cdc05d1265f.Test Plan
Check one/two/four-slot planning and metadata contracts on CPU. On GB10, compare
full KDA outputs and each of four exported states with independent FP32 and
prefix-state oracles; cover inactive slots, invalid fourth destinations,
changing live lengths/offsets under frozen graph replay, and live pool indices
past the signed 32-bit offset boundary.
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 defaultdev/jovian-judgementconfiguration 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.Test Result
TP4/DCP1 and TP4/DCP2 integration scope
The checkpoint kernel requires no DCP-specific change: the vLLM adapter selects retained boundaries, while planned B12X capacity remains four for two/three/four active destinations at DCP1/2/4. Combined-runtime tests passed 18 GPU recurrent-export/reload and convolution-history cases, including companion vLLM coverage. The complete B12X PR revision is not separately GPU-qualified by these measurements. DCP2 enabled cached/cold logprob difference 0.127015 exceeds its diagnostic bound; cache-invariant probabilities remain research-only despite matching answer tokens and passing matched on/off checks. The DCP1/DCP2 conditions, test commands and limitations and numeric evidence and receipt hashes identify runtime revisions, conditions and measurement scope.
TP4/DCP4 component and serving evidence
41 CPU tests passed. Twelve selected four-checkpoint GPU tests passed, with
15 other GPU cases deselected. Six companion vLLM convolution-history GPU tests
also passed on the same image. The GPU run checks outputs/state export and
replay contracts; it is not a standalone KDA latency benchmark.
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_COALESCINGandVLLM_GLM53_MHC_PREFILL_SHARD, respectively, are: neither0/0, coalescing-only1/0, mHC-only0/1, and both1/1.All arms use these native layout settings:
Together with aligned Mamba caching and retention interval zero, they establish
the tested physical/lookup/scheduler grid
(512, 512, 2048). CLI block sizesalone 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
abb715f132bdccb592a34b2596a3d3a8d757ffbcand B12X70fe41974ef4b18f61caaa2579c81cdc05d1265f. The isolated vLLM PRs compose to thisruntime 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.
The companion vLLM coalescing implementation uses the four-slot export contract.
Observed serving prefill rates in tokens/s:
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=1is common to all arms. Activation is checkedbefore 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.
A serving rate belongs to the identified B12X/vLLM composition and flag arm;
it does not measure an isolated kernel speedup. Serving and component evidence includes raw measurements, activation records, component checks and source hashes.
Observed MTP-normalized decode steps/s (mean acceptance length in parentheses):
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.
Both base
06b4de7cand feature70fe4197omit KDA registration from thepolicy catalog and offline generator inventory, and contain no embedded KDA
profile. AUTO uses heuristic planning; PREPLANNED_ONLY rejects the KDA query.
This inherited policy-profile limitation is separate from checkpoint-export
correctness and the reported serving measurements. AI assistance was used for implementation, review,
and test preparation. Submitted by FujitsuPolycom as a draft for maintainer review.
Related work
B12X #253 changes direct KDA decode binding; it does not export four interior prefill states. B12X #310 tunes mHC kernel schedules rather than recurrent checkpoint export. The inspected open PRs do not provide this bounded checkpoint interface.
Reproduction
Model/tokenizer:
local-inference-lab/GLM-5.3-Flash-NVFP4-Sparkatdf116c4fb16b1d37ae43d2cfd624de26ffbc832e. 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 original performance measurements are unchanged.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/DCP4 settings, 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
Serving integration: vLLM #699 consumes the four-checkpoint API. vLLM #700 implements independent token-sharded mHC; it is not required for checkpoint export.