Skip to content

[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
sgl-project:mainfrom
amd-danli103:common-kernel-resources
Open

amd-danli103 wants to merge 3 commits into
sgl-project:mainfrom
amd-danli103:common-kernel-resources

Conversation

@amd-danli103

@amd-danli103 amd-danli103 commented Nov 4, 2026 •

Copy link
Copy Markdown

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, which check pulls out of a FlyDSL launcher, and the occupancy, from max_blocks_per_cu(). So check takes FlyDSL launches; every monokernel in this repository is FlyDSL. Given a launch, check:

  • raises NotResident if 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_RESOURCES says how, set in the server's environment (unset is warn):

    SGLANG_AMD_MONO_RESOURCES=warn    # log each kernel past the budget
    SGLANG_AMD_MONO_RESOURCES=strict  # raise OverBudget instead, and raise when a binary or its occupancy cannot be read
    SGLANG_AMD_MONO_RESOURCES=report  # warn, and also log every kernel within budget

    Any other value raises ValueError. NotResident raises in every mode.

  • with group (the TP CPU group), makes a failure inside check raise 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_cache and the error text of max_blocks_per_cu(). If FlyDSL changes them, two tests in test/unit fail on a machine with FlyDSL and aiter, and at runtime the check warns once and is skipped (strict raises). The gpu.binary decoding is adapted from ROCm/ATOM, which is MIT-licensed; the file header keeps its copyright notice. python -m sgl_amd_mono.common.resources DUMP_DIR reads the kernels of a FLYDSL_DUMP_IR=1 build offline.

DeepSeek-V4.1 (deepseek_v41/gfx950/runner.py)

precompile checks 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

Accuracy Tests

No kernel or forward code changes.
What check reports, on MI355X with lmsysorg/sglang-rocm:v0.5.23-rocm10-mi35x-20261030 (FlyDSL 0.4.0) and main at ca929b7. DeepSeek-V4.1's precompile at TP2 logs, among others:

monokernel spills: dsv41_mono_k1_s18_tp2_kt0_f0_i1_q1_p1: VGPR 256 AGPR 0 SGPR 106, spilled VGPR 34 SGPR 125, scratch 80 B a lane, LDS 149056 B, 1 resident a CU (budget 0 B a lane)

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.

  • DeepSeek-V4.1, every launch precompile builds 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).
  • The other models with monokernels (DeepSeek-V4-Pro, GLM-5.3, GLM-5.3-Flash, MiniMax-M3, MiniMax-M3.1, Muse Spark, Qwen3.8-Flash-Next; 28 builds, compile-only): all read and resident. One GLM-5.3 build (TP4, 8 rows, indexer) spills one VGPR.

Speed Tests and Profiling

Startup time is essentially unchanged: check runs once a launch, after the build, and reads what the build made. DeepSeek-V4.1's precompile of the same 100 launches a TP size (each rank starting at its own share of the widths, as the spec runs it), main and this PR at the same time on one node (separate GPUs and FlyDSL caches), one run each:

precompile wall time TP2 TP4
cold FlyDSL cache: main / this PR 214.6 / 217.5 s 116.2 / 117.5 s
warm FlyDSL cache: main / this PR 1.7 / 3.2 s 1.8 / 3.5 s

check adds 1-3 s, once, at start: about 1% of a cold precompile.

Checklist

  • pre-commit run --all-files passes.
  • Tests under 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 runs check'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:
    pytest test/unit/test_common_resources.py  # 43 passed in about 5 s, with one GPU visible or none (HIP_VISIBLE_DEVICES=-1)
    pytest test/unit                           # 587 passed / 8 failed; main: 544 / 8
    The two FlyDSL tests skip where FlyDSL and aiter are not installed. The 8 failures are the same on main: the MiniMax and Muse Spark contract tests and three glm53_flash spec tests.
  • The SGLang commit or PR this was tested against is named above (the image's SGLang).
  • Accuracy and speed results are attached when kernel or forward code changes. (No kernel or forward code changes.)
  • A new model adds one line to .github/CODEOWNERS that names its owner. (No new model.)

Review and Merge Process

  1. The lint and Check PR title checks must pass.
  2. Ask a code owner from .github/CODEOWNERS to review.
  3. Merge with "Squash and merge": the PR title becomes the commit title on main.

amd-danli103 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.
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