Skip to content

[FlashInfer MLA] Use public CUDA graph plan updates - #35555

Draft
saltyminty wants to merge 1 commit into
sgl-project:mainfrom
saltyminty:fix/mingyangw/sglang-mla-api-demonstration
Draft

saltyminty wants to merge 1 commit into
sgl-project:mainfrom
saltyminty:fix/mingyangw/sglang-mla-api-demonstration

Conversation

@saltyminty

@saltyminty saltyminty commented Aug 19, 2026

Copy link
Copy Markdown

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 private
workspace, 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
14096da3f1490d94a640c7d55a07bbe856f20279 directly.

Modifications

  • Opt CUDA-graph MLA wrappers into public graph plan updates.
  • Submit complete MLAPlanMetadata.csr(...) values through
    update_cuda_graph_plan() during decode and target-verify replay.
  • Keep a reusable non-alias page-index source for the public update contract.
  • Remove both private fast_mla_* planner helpers and their plan()
    monkey-patches.
  • Leave existing eager planning and attention runtime calls unchanged to keep
    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.
  • Task-local public-update probe: 20/20 updates passed, capture-stable buffer
    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

  1. Ping Merge Oncalls to start the process. See the PR Merge Process.
  2. Get approvals from CODEOWNERS and other reviewers.
  3. Trigger CI tests with comments or contact authorized users to do so.
    • Common commands include /tag-and-rerun-ci, /tag-run-ci-label, /rerun-failed-ci
  4. After green CI and required approvals, ask Merge Oncalls or people with Write permission to merge the PR.

CI States

Latest PR Test (Base): ❌ Run #34307729743
Latest PR Test (Extra): ❌ Run #34307729362
Latest PR Test (AMD ROCm 7.2): ❌ Run #34307729485

Signed-off-by: mingyangw <mingyangw@nvidia.com>
@saltyminty
saltyminty force-pushed the fix/mingyangw/sglang-mla-api-demonstration branch from 239d77e to b592fcd Compare September 9, 2026 03:35
@saltyminty saltyminty changed the title [FlashInfer MLA] Migrate to unified planning and tensor APIs [FlashInfer MLA] Use public CUDA graph plan updates Sep 9, 2026
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.

1 participant