Repository navigation
feat(kda): integrate FlashInfer prefill checkpoints - #41400
Conversation
github-actions
Bot
added
documentation
fe4cfba to
0c7733a
Compare
kpham-sgl
changed the title
ba17536 to
30ec510
Compare
ae9847e to
706048f
Compare
|
Rebased and force-pushed this PR onto main Reran both prefix-cache methods on the exact rebased head, with the full upstream fix and no diagnostic patch: real Kernel/configuration checks: 23 tests passed, including all eight FlashInfer KDA numerical/layout cases, in 229.46 seconds. Pre-commit and focused CPU contract/metadata checks also passed. Cold versus cached continuation: score the same fixed 167-token continuation after 32,769 / 36K / 72K context tokens. All target scores are finite; inputs are identical between backends.
Branching corruption check from #42178: cache an 8,320-token prefix; score copying a 167-token span from branch A before and after an unrelated branch B extends the shared prefix. Three different spans/seeds per backend.
The pre-rebase controls dropped from approximately -0.02 to -1.35 after the competing request. The rebased runs retain high copy scores with both KDA prefill backends and resume at 8,192, keeping the pooled index group private. This reproducer no longer shows the large branching collapse. Warm repeats and the aligned-boundary control are not bit-exact on GLM. The Triton aligned control has a 2.9145 maximum cold/cached difference and a 2.0505 maximum warm/repeat difference despite a small p90; these outliers preclude claiming exact cache parity. These measurements do not isolate the residual drift to a specific operator. HiCache host restore and current-head speed were not tested. |
|
/rerun-test test/registered/unit/layers/attention/test_linear_attn_config.py test/registered/attention/test_kda_prefill_flashinfer.py test/registered/e2e/models/test_glm53_flash_b200.py::TestGLM53FlashB200LowLatency.test_gsm8k |
|
🚀 🚀 |
|
/rerun-test test/registered/unit/layers/attention/test_linear_attn_config.py |
|
🚀 |
|
Retested current opt-in head
Environment: PyTorch 2.14.1+cu130, FlashInfer 0.7.0.post1, CUDA 13.0; FP8 KV and TRTLLM DSA/MoE as in These are local GB300 runs of the B200 recipe, not B200 CI results. The runner explicitly selected Python TreeCore ( Retested Kimi-K3 locally on 8×GB300 across two four-GPU hosts, retaining the registered B300 workloads and thresholds. Added only the explicit FlashInfer prefill selection and two-node launch settings, and selected verified cached snapshots: target
Both hosts use the same package versions and explicit Python TreeCore selection above. These are two-node GB300 runs of the B300 recipes, not B300 CI results. An earlier shared-storage startup failed before inference with CUDA error 719 in an expert-weight CPU-to-GPU copy. The local-NVMe retry passed both accuracy tests above; the slow shard was verified against its checkpoint SHA256 on both hosts. Ran the existing
Both exceed the existing quality threshold 0.75 and acceptance threshold 2.4. FlashInfer's aggregate acceptance is 0.6% higher here; this MMMU-Pro run does not reproduce the single-prompt acceptance drop. The stock preset is stochastic; these are single 200-example evaluations, not an output-parity test. Evaluation latency was 1152.72s / 1112.00s, with 1,086,318 / 1,092,392 generated tokens (FlashInfer / Triton). The separate existing The initial MMMU attempts failed before evaluation because |
kpham-sgl
added
run-ci
|
/tag-and-rerun-ci |
|
cc @leejnau |
Resolve the conflict in kimi_linear.py: keep this PR's logical-mode comment and drop the gate-activation comment that sgl-project#41400 removed on main. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
2 participants
Motivation
FlashInfer #5021 (0.7.0) adds indexed, floor-aligned checkpoints to CuTe DSL KDA prefill, which SGLang needs for radix prefix caching.
Modifications
--linear-attn-prefill-backend flashinferfor BF16, 128-dim safe-gate KDA on SM100/SM103.flashinfer.kda.recurrent_kda(..., backend="cute-dsl")directly. FlashInfer owns scheduling and workspace allocation; SGLang maps radix-cache checkpoint destinations once per tracked batch. Invalid scheduler/checkpoint metadata raises an assertion.This PR leaves dependency pins unchanged; the current tree already pins FlashInfer 0.7.0.post1.
Accuracy Tests
Historical default-selection experiment (reverted by
c0d8467478; FlashInfer remains opt-in): pre-commit passed. The registered CPU rerun passed all 10 tests ona231d574ae, after correcting the policy fixture to publish a resolved CUDA graph config. The B200 rerun passed FlashInfer KDA kernel tests and GLM-5.3-Flash low-latency GSM8K on70b968ef19: score 0.932 over 500 questions. The server log confirmsdecode=triton, prefill=flashinfer, verify=nv_cutedsl. The follow-up commit changes only the unit fixture.Previously reported on 8xGB300, FlashInfer 0.7.0: the original five KDA prefill cases, existing KDA decode/verify (11), and CuTe DSL KDA + GDN checkpoint smoke (2) passed.
The current GPU test file contains six checkpoint/layout cases and two numerical backend cases (tracked prefill and single-token fallback). The direct-call cleanup passes pre-commit and focused CPU checks: eight checked-in test methods, 22 raw-beta contract checks, six checkpoint harness cases exercising the production adapter with a substituted numerical kernel, and six metadata-hook cases. GPU imports and numerical kernels were isolated for these local CPU checks. The eight GPU cases passed on pre-rebase head
ae9847e073on B300 (SM103), FlashInfer 0.7.0.post1 and CUDA 13.0, in 207.87 seconds. The exact rebased-head rerun is reported below.Exact rebased-head GLM-5.3-Flash checks
Reran both prefix-cache methods on the exact rebased head
706048f3fd(before the default-selection change), with the full upstream fix and no diagnostic patch: realzai-org/GLM-5.3-Flashweights (revisioneb9eb208eb0d988989d07a6a12d0fdeb5f52574a), 4xB300 / TP4, BF16 KV and SSM caches, TileLang DSA, Triton decode, 16K prefill chunks, eager execution, and--mamba-track-interval 256. FlashInfer 0.7.0.post1, CUDA 13.0, PyTorch 2.14.1 (Python TreeCore), Transformers 5.17.0. Both arms useSGLANG_JIT_DEEPGEMM_PRECOMPILE=0to compile the exercised GEMM shapes on demand; the initial exhaustive startup compilation was stopped before probe requests.Kernel/configuration checks: 23 tests passed, including all eight FlashInfer KDA numerical/layout cases, in 229.46 seconds. Pre-commit and focused CPU contract/metadata checks also passed.
Cold versus cached continuation: score the same fixed 167-token continuation after 32,769 / 36K / 72K context tokens. All target scores are finite; inputs are identical between backends.
Branching corruption check from #42178: cache an 8,320-token prefix; score copying a 167-token span from branch A before and after an unrelated branch B extends the shared prefix. Three different spans/seeds per backend.
The pre-rebase controls dropped from approximately -0.02 to -1.35 after the competing request. The rebased runs retain high copy scores with both KDA prefill backends and resume at 8,192, keeping the pooled index group private. This reproducer no longer shows the large branching collapse.
Warm repeats and the aligned-boundary control are not bit-exact on GLM. The Triton aligned control has a 2.9145 maximum cold/cached difference and a 2.0505 maximum warm/repeat difference despite a small p90; these outliers preclude claiming exact cache parity. These measurements do not isolate the residual drift to a specific operator. HiCache host restore and current-head speed were not tested.
GLM-5.3-Flash checks before rebase
2026-10-07, head
ae9847e073, real GLM-5.3-Flash weights, 4xB300 TP4, BF16 KV/SSM caches, TileLang DSA, 16K prefill chunks, eager execution, Triton decode. FlashInfer 0.7.0.post1, CUDA 13.0, PyTorch 2.14.1 (Python TreeCore), Transformers 5.17.0. Model revisioneb9eb208eb0d988989d07a6a12d0fdeb5f52574a.Cold versus cached scoring uses identical contexts and a fixed 167-token continuation. Values are absolute logprob difference mean / p90 / max.
The branching check warms a prefix ending inside a 256-token index group, scores copying a span from the first suffix, then rescoring follows a competing request with unrelated text. Values below average three different 167-token copy spans.
Before rebase, head
ae9847e073lacked #42178. Both backends reproduced its pooled-index branching corruption; the device-cache portion of that fix, ported only into the isolated test checkout, removed the collapse. The PR has now been rebased onto main7760d5264a, which contains the full merged #42178 fix. The earlier diagnostic measurements below retain their original source provenance.Warm repeats and the 32,768-token aligned-resume control are not bit-exact on GLM; residual small drift is not isolated to KDA. At 72K on the old layout, warm/repeat cache depths also differ (65,536 / 71,936). Fixed-layout arms use the required 256-token decode tracking interval; old-layout arms use 64. Probe requests generate zero tokens, and prefill tracking follows page/chunk geometry. HiCache host restore and latest-head speed were not tested.
Historical Kimi K3 measurements
The following Kimi K3 results predate the current direct-call cleanup.
Kimi K3 real weights, 8xGB300, BF16 KV, DSpark: cold vs prefix-cache-resume logprobs over a fixed 167-token continuation (|Δlogprob| mean / p90 / max). Warm vs warm is bit-exact in all cases.
FlashInfer matches Triton. The drift is from chunk layout, not state restore: a resume on a cold chunk boundary (32768) is bit-exact.
Speed Tests
These are historical measurements from the wrapper-based integration. The current direct API can read packed sequence offsets on the host to size its workspace and allocates that workspace within each call; performance and D2H behavior need remeasurement on the latest head.
Kimi K3 NVFP4 dummy weights, 8xGB300 TP8/DCP8, bs1, 131k in / 1k out, 95% prefix hit, A/B/A x10:
An intermittent ~6.7 s TTFT stall appears with both backends and is not attributed to KDA.
Kernel-only KDA prefill (checkpoints off), K3 shape on one TP8 rank (12 heads x 128), single GB300, bf16 in / fp32 state, median of 5 (ms). Backends that fall back to Triton at these lengths are omitted:
nvidia_kda(needs >1 sequence),flashkda(caps at 2048 tokens/seq).FlashInfer is 1.3-1.8x Triton on a single sequence and 3.2x on the 8x16k shape chunked prefill actually issues, where it batches the sequences into one kernel.
ptx_kdais ~20% faster on a single sequence but 2.3x slower on that chunked shape; it is also not measured here with checkpoints on, which its own fallback note says routes tracked batches to Triton.cutedslcannot serve K3 at all (no safe gate, no interior checkpoints).Checklist
CI States
Latest PR Test (Base): ⏳ Run #37809707868
Latest PR Test (Extra): ❌ Run #37850689948
Latest PR Test (AMD ROCm 10): ⏳ Run #37809708463