Initial Layerwise Profiler

[LucasWilkinson]

SUMMARY:

Initial layerwise profiler leveraging the kineto base PyTorch profiler.

NOTE: we run in eager mode by default so that the stack-trace/event-tree contains nn.Module, otherwise if cuda-graphs is used all the kernels will be under the CUDAGraphRunner module (this is converted to an nn.Module in this PR so that the _build_module_tree code puts all those kernels under a CUDAGraphRunner)

NOTE: vllm kernels like vllm::reshape_and_cache_kernel have no trace or shape informat because they are not registered as a TorchOp (i.e. using TORCH_LIBRARY, they instead just a raw PYBIND11_MODULE module)

Example on how to use visualization:

pip install -r requirements-dev.txt

python examples/offline_profile.py --model nm-testing/OpenHermes-2.5-Mistral-7B-pruned50 --batch-size 4 --prompt-len 512 --json openhermes7b-dense

### For Breakdown Graphs
# For module level breakdown
python neuralmagic/tools/profiler/visualize_trace.py --json-trace openhermes7b-dense.json --output profile_dense.pdf
# For kernel level breakdown
python neuralmagic/tools/profiler/visualize_trace.py --json-trace openhermes7b-dense.json --output profile_dense_kernels.pdf --level kernel

### For table printing
# Decode Summary Table
python neuralmagic/tools/profiler/print_table.py --json-trace openhermes7b-dense.json --phase decode
# Prefill Summary Table
python neuralmagic/tools/profiler/print_table.py --json-trace openhermes7b-dense.json --phase prefill
# Decode Model Layerwise Table
python neuralmagic/tools/profiler/print_table.py --json-trace openhermes7b-dense.json --phase decode --table model
# Prefill Model Layerwise Table
python neuralmagic/tools/profiler/print_table.py --json-trace openhermes7b-dense.json --phase prefill --table model

Example Output: profile-example-output.txt

TEST PLAN:

• GHA

[robertgshaw2-neuralmagic]

Is this always true?

Does the fact that max_num_batched_tokens < batch_size * prompt_len guarantee this?

[LucasWilkinson]

as of now yes, but its brittle in the sense that logic could change and this would fall out of sync

[robertgshaw2-neuralmagic]

Maybe this should be set to batch_size * prompt_len + eps?

[LucasWilkinson]

personally I like forcing the user to set this (in this case a dev since this tooling is targeting devs) since it means less surprises when the number of GPU blocks decreases and will make the user think twice before setting a very large batch + prompt-len (and/or making it less surprising when we go OOM). But I can be convinced otherwise.