Kernel shape profiler registry - #1024
Open
mohbasit wants to merge 11 commits into
Open
Conversation
Restores the registry-based kernel_shape_profiler as a standalone, non-invasive tool: wraps kernel *launcher* functions as torch.library custom ops so they appear as cpu_op events with Input Dims / Input type, covering any backend behind the launcher (Triton, ASM, CK, aiter C++ bindings, collectives) rather than only JIT-dispatched kernels. - kernel_shape_profiler.py: explicit registry plus filtered auto-discovery; builds a schema from each launcher's signature, passes non-tensor args and the return value through a thread-local side channel, and rebinds every module-level reference so 'from X import Y' callers are intercepted too. Recovered from the retired SGLang patch and adapted to import standalone. - sitecustomize.py: auto-loaded via PYTHONPATH, drives enable()/disable() from the torch profiler window so nothing is wrapped outside a run. - README.md: approach, safety properties, and a comparison against the JIT-hook tracer on the jit-shape-tracer branch. Two fixes over the recovered original: - Registry refreshed for current layouts. Only 5 of the original 24 entries still resolved against sglang 0.5.18 + aiter (sglang moved Triton kernels to sglang.kernels.ops.*, aiter regrouped its Triton ops). Current paths added alongside the legacy ones, which are kept since unresolvable entries are skipped silently; 21 of 40 now resolve. - Guard against double wrapping. Frameworks keep compat re-exports, so after the first entry is wrapped _patch_all_references rebinds the second path to our own wrapper, which was then wrapped again and produced two nested annotations for one call. Wrappers now carry a _kernel_shape_wrapper marker that enable() skips. Validated on MI300X: synthetic launcher and the real aiter gemm_a8w8_blockscale both annotated with correct fp8/bf16 per-operand dtypes, exactly once per call, with references restored on disable().
SGLang 0.5.18 moved its Triton kernels out of sglang.srt.layers.* into a separate sglang.kernels.ops.* package, which auto-discovery was not scanning. Of the launchers that actually appear in DeepSeek-R1 traces, discovery now finds 6 of 7 instead of 1 of 7: concat_and_cast_mha_k_triton, set_mla_kv_buffer_triton, vocab_parallel_embedding, compute_position_triton and clamp_position_cuda are all gained. The seventh, create_flashinfer_kv_indices_triton, is a @triton.jit kernel rather than a Python launcher, so it is never wrapped by design. Candidate count goes from 1186 to 2012. Requiring a source-level launch pattern in addition to a Tensor annotation would cut that to 413, but it also drops real launchers (compute_position_triton, clamp_position_cuda, write_cache_indices), so the looser filter is kept and the trade-off is documented in the README.
…rvation Port the active-device fix from main's kernel_shape_profiler patch (#980): read torch.utils._device.CURRENT_DEVICE (None when unset) instead of torch.get_default_device() (always concrete). Restoring the raw override — including None — truly clears any default-device mode a module leaks at import time, whereas restoring concrete 'cpu' would install a device mode. Co-authored-by: Cursor <cursoragent@cursor.com>
mohbasit
requested review from
ajassani,
devalshahamd,
gabeweisz and
tsrikris
as code owners
September 11, 2026 03:50
Collaborator
|
Looks good to me but checks fail and this needs tests |
Co-authored-by: Cursor <cursoragent@cursor.com>
Co-authored-by: Cursor <cursoragent@cursor.com>
…com:AMD-AGI/TraceLens into kernel-shape-profiler-registry
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
2 participants
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.
Currently to get shapes in SGLang runs, we have been patching the docker image, this method removes the patching and this tool can be used to get shapes in traces for SGlang framework by just putting the file in the PYTHONPATH and enabling flags.