Repository navigation
[Feat] common: Check each compiled kernel's residency and scratch where a runner builds it - #101
Open
amd-danli103 wants to merge 3 commits into
Open
amd-danli103 wants to merge 3 commits into
amd-danli103 wants to merge 3 commits into
Conversation
added 3 commits
November 4, 2026 04:24
resources.check reads a FlyDSL launcher's code objects where a runner builds it, raises for a kernel whose grid the CUs cannot hold, and warns (under SGLANG_AMD_MONO_RESOURCES=strict, raises) past the caller's scratch budget; a binary or occupancy it cannot read warns once and is skipped.
Any scratch warns (ROCm/ATOM refuses a width for any); the TP ranks raise together.
A mistyped SGLANG_AMD_MONO_RESOURCES raises ValueError with a group as without one. The gloo test reports exit codes, and the tests read only this module's log records.
amd-danli103
requested review from
HaiShaw,
alexnails,
kevin-mii and
kkHuang-amd
as code owners
November 4, 2026 08:31
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
Nothing checks what a monokernel's binary asks of the GPU once it is built. A kernel whose waves spill VGPRs to scratch is slower, and nothing reports it: from #30 until #95, every DeepSeek-V4.1 K2 build at TP2 spilled to scratch (248-440 B a lane, read at 517040f), which nobody saw until a bisect for #72. A kernel whose grid the CUs cannot hold at once never finishes; SGLang's binding checks the CU count against the grid, not what each kernel's registers and LDS let a CU hold. minimax_m31 checks that for its own launches (
minimax_m31/residency.py); no other model does.This PR reads each binary's registers, spills, scratch and LDS where a runner builds it, in
common/so that any model can call it, and has DeepSeek-V4.1 call it.Modifications
common/resources.py(new)check(compiled, budget=Budget(), log=True, group=None)runs right after a runner builds a launch. The numbers it reads (registers, spills, scratch per lane, LDS) come from the metadata LLVM writes into every AMDGPU code object, whatever the kernel is written in. FlyDSL supplies the rest: the code object, whichcheckpulls out of a FlyDSL launcher, and the occupancy, frommax_blocks_per_cu(). Sochecktakes FlyDSL launches; every monokernel in this repository is FlyDSL. Given a launch,check:raises
NotResidentif the CUs can't hold the whole grid at once. The server then fails to start, on purpose: such a launch would otherwise hang until its deadline.reports scratch above
budget.scratch(default 0, so any scratch, the same line ROCm/ATOM uses).SGLANG_AMD_MONO_RESOURCESsays how, set in the server's environment (unset iswarn):Any other value raises
ValueError.NotResidentraises in every mode.with
group(the TP CPU group), makes a failure insidecheckraise on all ranks together, so the others do not wait at their next collective on the rank that stopped.Two of the FlyDSL pieces it uses are not public: a launcher's
_mem_cacheand the error text ofmax_blocks_per_cu(). If FlyDSL changes them, two tests intest/unitfail on a machine with FlyDSL and aiter, and at runtime the check warns once and is skipped (strictraises). Thegpu.binarydecoding is adapted from ROCm/ATOM, which is MIT-licensed; the file header keeps its copyright notice.python -m sgl_amd_mono.common.resources DUMP_DIRreads the kernels of aFLYDSL_DUMP_IR=1build offline.DeepSeek-V4.1 (
deepseek_v41/gfx950/runner.py)precompilechecks every launch it builds, for the target and the DSpark draft, with the runner's TP group; rank 0 logs.No kernel or forward code changes, and no other model is touched.
Not in this PR
minimax_m31/residency.pyontocommon.resources, its owner's call.strictkernel test for DeepSeek-V4.1, once its builds have no scratch.Accuracy Tests
No kernel or forward code changes.
What
checkreports, on MI355X withlmsysorg/sglang-rocm:v0.5.23-rocm10-mi35x-20261030(FlyDSL 0.4.0) andmainat ca929b7. DeepSeek-V4.1'sprecompileat TP2 logs, among others:The name says which build it is (K1 at 18 rows, TP2, token parallel, with the index source and its index query), and the rest how far over it is: all 256 VGPRs in use and 34 more spilled to 80 B of scratch a lane, with the grid still resident. Until now that took dumping and reading the ISA build by build; with this PR a change that brings a spill in shows it at the next start.
precompilebuilds for the target and the DSpark draft (100 a TP size, TP2 and TP4): all 200 read, each grid resident at one CTA a CU. 17 have scratch and warn: the index-source K1 under token parallel (10), the K2 at 12 rows under token parallel (6) and the draft's K2 at 40 rows (1).Speed Tests and Profiling
Startup time is essentially unchanged:
checkruns once a launch, after the build, and reads what the build made. DeepSeek-V4.1'sprecompileof the same 100 launches a TP size (each rank starting at its own share of the widths, as the spec runs it),mainand this PR at the same time on one node (separate GPUs and FlyDSL caches), one run each:precompilewall timemain/ this PRmain/ this PRcheckadds 1-3 s, once, at start: about 1% of a cold precompile.Checklist
pre-commit run --all-filespasses.test/cover the change:test/unit/test_common_resources.py(43, CPU) reads code objects built in the test (a 64-bit ELF with an AMDGPU metadata note, escaped as MLIR prints it) and one FlyDSL kernel it compiles (compile-only); it runscheck's modes, unreadable code objects, an occupancy FlyDSL cannot give, a grid the CUs cannot hold, a repeated code object, two of one name, a long symbol and ranks that raise together (two gloo processes), and checks that FlyDSL still raises the text read as zero blocks. In the image above, from the repository root:main: the MiniMax and Muse Spark contract tests and threeglm53_flashspec tests..github/CODEOWNERSthat names its owner. (No new model.)Review and Merge Process
lintandCheck PR titlechecks must pass..github/CODEOWNERSto review.main.