Skip to content

Allow platforms to declare Mamba extra-buffer support - #36410

Merged
BBuf merged 2 commits into
sgl-project:mainfrom
xuzijian629:fix/mamba-extra-buffer-platform
Oct 6, 2026
Merged

BBuf merged 2 commits into
sgl-project:mainfrom
xuzijian629:fix/mamba-extra-buffer-platform

Conversation

@xuzijian629

@xuzijian629 xuzijian629 commented Aug 26, 2026 •

Copy link
Copy Markdown
Contributor

Motivation

Mamba extra-buffer validation currently uses a fixed list of device types. This prevents out-of-tree platforms from enabling the feature even when they implement the required cache-state snapshot behavior, and it couples server argument validation to specific hardware names.

Modifications

  • Add SRTPlatform.support_mamba_cache_extra_buffer() as the platform capability used by server argument validation.
  • Preserve the existing behavior for CUDA, ROCm, MUSA, NPU, and XPU platforms.
  • Allow other platforms to opt in by overriding the capability method.
  • Add CPU unit tests covering the default capability behavior, an out-of-tree override, and server argument validation dispatch.

Accuracy Tests

N/A. This changes platform capability validation only and does not modify model computation or numerical behavior.

Speed Tests and Profiling

N/A. This change is outside the inference hot path.

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 #37397056745
Latest PR Test (Extra): ❌ Run #37397056334
Latest PR Test (AMD ROCm 10): ❌ Run #37397056631

@alexnails

Copy link
Copy Markdown
Collaborator

/tag-and-rerun-ci

@github-actions github-actions Bot added the run-ci CI: run the baseline test suite on this PR label Aug 27, 2026
@BBuf

BBuf commented Aug 29, 2026

Copy link
Copy Markdown
Collaborator

Thanks! This branch now needs to be updated onto the latest main. The Mamba extra-buffer validation has moved from ServerArgs._validate_mamba_extra_buffer to python/sglang/srt/arg_groups/mamba_hook.py::validate_mamba_extra_buffer. Please port the current_platform.support_mamba_cache_extra_buffer() check and the corresponding test/patch target to that hook, then re-request review.

@xuzijian629
xuzijian629 force-pushed the fix/mamba-extra-buffer-platform branch from 6ca4a34 to e4237f2 Compare September 3, 2026 21:44
@xuzijian629

Copy link
Copy Markdown
Contributor Author

@BBuf @alexnails thanks for reviewing! Rebased onto master

@xuzijian629
xuzijian629 force-pushed the fix/mamba-extra-buffer-platform branch from ed44692 to d9cc76c Compare September 29, 2026 19:14
@xuzijian629

Copy link
Copy Markdown
Contributor Author

/tag-and-rerun-ci

@xuzijian629

Copy link
Copy Markdown
Contributor Author

@alexnails thanks for reviewing this PR!! I rebased the branch. Are we ready to merge it?

@xuzijian629
xuzijian629 force-pushed the fix/mamba-extra-buffer-platform branch 4 times, most recently from 24c9161 to a79200a Compare October 5, 2026 15:20
# Problem
Mamba extra-buffer validation hardcodes the supported in-tree device types, so an out-of-tree platform cannot opt in even when it implements the required cache-state snapshots.

# Root cause
Built-in hardware facts are read through runtime_context.get_platform(), while current_platform owns out-of-tree plugin behavior. The platform interface had no capability for this Mamba cache behavior.

# Fix
Add a conservative SRTPlatform capability for out-of-tree plugins. Preserve the existing built-in checks through get_platform() and consult the capability only for an out-of-tree platform. Add CPU tests for every supported built-in fact and both out-of-tree capability outcomes.

Test Plan:
To be added
@xuzijian629
xuzijian629 force-pushed the fix/mamba-extra-buffer-platform branch from a79200a to 22b410e Compare October 5, 2026 17:45
@xuzijian629

Copy link
Copy Markdown
Contributor Author

All PR Test Base tests are passing, except for the H200 DeepSeek-v4 FlashInfer/EAGLE ones, which are unrelated to my change.
Our code only runs for Mamba extra-buffer validation, which this DeepSeek test never reaches.

@BBuf
BBuf merged commit 40fb593 into sgl-project:main Oct 6, 2026
144 of 172 checks passed
@xuzijian629
xuzijian629 deleted the fix/mamba-extra-buffer-platform branch October 6, 2026 11:37
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

run-ci CI: run the baseline test suite on this PR

Projects

None yet

Development

Successfully merging this pull request may close these issues.

3 participants