[FlashInfer MLA] Use public CUDA graph plan updates - #35555
Draft
saltyminty wants to merge 1 commit into
Draft
saltyminty wants to merge 1 commit into
saltyminty wants to merge 1 commit into
Conversation
5 tasks
11 tasks
Signed-off-by: mingyangw <mingyangw@nvidia.com>
saltyminty
force-pushed
the
fix/mingyangw/sglang-mla-api-demonstration
branch
from
September 9, 2026 03:35
239d77e to
b592fcd
Compare
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Motivation
Replace SGLang's private FlashInfer MLA CUDA-graph replanning bridge with the
public plan-update API introduced by
FlashInfer #5041.
The current implementation monkey-patches
plan()and reaches into privateworkspace, metadata-buffer, and generated-module attributes. This draft shows
the public replacement while preserving the existing eager and full-plan paths.
This PR remains a draft until SGLang can consume a FlashInfer release containing
the new API. Validation below used FlashInfer commit
14096da3f1490d94a640c7d55a07bbe856f20279directly.Modifications
MLAPlanMetadata.csr(...)values throughupdate_cuda_graph_plan()during decode and target-verify replay.fast_mla_*planner helpers and theirplan()monkey-patches.
this downstream demonstration narrowly scoped.
Accuracy Tests
Targeted B200/SM100 validation against the FlashInfer commit above:
test/registered/attention/unittests/mla/test_flashinfer.py: 6 passed,1 registered skip, and 15 subtests passed. The skip is the existing SM100
EAGLE draft CUDA-graph known failure.
pointers remained unchanged, the page-index source was non-aliasing, and no
private bridge symbols remained.
Speed Tests and Profiling
No standalone SGLang speed benchmark was run for this downstream demonstration.
The sync-free implementation and performance validation live in FlashInfer
#5041; this PR only replaces SGLang's private bridge with that public API.
Checklist
Review and Merge Process
/tag-and-rerun-ci,/tag-run-ci-label,/rerun-failed-ciCI States
Latest PR Test (Base): ❌ Run #34307729743
Latest PR Test (Extra): ❌ Run #34307729362
Latest PR Test (AMD ROCm 7.2): ❌ Run #34307729485