gpt-oss-20b¶

Installation¶

Install vllm-rbln as described in RBLN vLLM. The example scripts on this page live under examples/ in the vllm-rbln repository, so following the source-code install leaves them in place as well.

Activate RBLN Profiler¶

Setting RBLN_PROFILER=1 enables the RBLN Profiler.

1

export RBLN_PROFILER=1

No additional profiler-specific settings are required for Rebel100™.

Extract Profiled Data with Inference¶

See data_parallel_offline.py in vllm-rbln for an example of data-parallel execution. To profile a model, keep its execution settings and pass profiler_config when creating the LLM. After warm-up, use start_profile() and stop_profile() to mark the inference range.

The RBLN Profiler runs on the vllm model implementation, and model_impl="vllm" on LLM selects it.

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17

llm = LLM(
    model=model,
    model_impl="vllm",
    tensor_parallel_size=tp_size,
    enable_expert_parallel=enable_ep,
    block_size=block_size,
    max_num_batched_tokens=128,
    max_num_seqs=1,
    profiler_config={
        "profiler": "torch",
        "torch_profiler_dir": "./profile",
    },
)
llm.generate(prompts, SamplingParams(temperature=0.0, max_tokens=2))
llm.start_profile()
outputs = llm.generate(prompts, sampling_params)
llm.stop_profile()

--dp-size sets the data-parallel size. This page keeps it at 1; profiling a run split over several ranks is covered in gpt-oss-20b (Data Parallelism).

1
2
3
4

$ cd vllm-rbln
$ RBLN_PROFILER=1 python3 examples/vllm/features/data_parallel/data_parallel_offline.py \
    --model openai/gpt-oss-20b \
    --dp-size 1

Which profiler start_profile() and stop_profile() control depends on the profiler_config setting.

• With "profiler": "torch" as above: the PyTorch profiler opens and closes the session, and the RBLN trace is merged into its trace. See PyTorch Profiler Integration.
• Run without profiler_config: the path that drives the RBLN Profiler directly is selected, and the RBLN Profiler session starts and ends at start_profile() and stop_profile().

stop_profile() returns after the trace has been written, so the *.pb file is on disk before the process exits.

When the run finishes, each of the two profilers writes its own output: the RBLN Profiler's traces go under a per-rank directory as .pb files, and the PyTorch profiler's chrometrace goes to the directory named by torch_profiler_dir.

 1
 2
 3
 4
 5
 6
 7
 8
 9
10

# RBLN Profiler
$ ls rank_0/session_0
rbln_{date}_{time}.pb            # the whole trace; this is the file you open in Perfetto
rbln_{date}_{time}_0.pb          # one per execution
rbln_{date}_{time}_1.pb
...

# PyTorch profiler
$ ls profile
dp0_pp0_tp0_dcp0_ep0_rank0.{timestamp}.pt.trace.json.gz

That chrometrace holds not only the framework operations but also the device slices the RBLN Profiler collected. The same run therefore produces both a .pb that holds the device side alone and a chrometrace that holds the framework and the device together. Open the .pb to look at device execution by itself, and the chrometrace to look at it alongside the framework calls — which is what PyTorch Profiler Integration covers.

1

profiler_config = {"profiler": "torch", "torch_profiler_dir": "./profile"}

Asynchronous Dispatch of Decoder and Sampler¶

One decode step splits into executions of distinct compiled models — the decoder and the sampler — each leaving its own slice on the Device Workload track. The slice and command names identify the model associated with each execution; the naming convention is described in How to Analyze.

Commands of one execution carry the same W{workload_index} token in their names. Read the token from the slice's workload_index and search the trace for it to collect the commands that execution ran on each of the four chiplets.

Because the framework dispatches these executions asynchronously, their slices do not necessarily lie strictly one after another. When operations execute asynchronously or concurrently, their slices overlap in time, and Perfetto keeps such overlapping slices on one track and stacks them vertically, so the Device Workload track expands into several rows.

The two stacked slices in the screenshot are an example. The upper one, i2_1_Device Workload_18, is the decoder's execution, and the lower one, i3_2_Device Workload_0, is the sampler's. The sampler's slice starts while the decoder is still running.

The two are dispatched back to back. Early in the sampler's slice, the External HDMA that brings its input in runs first, and the computation on Neural Engine Clusters can be seen starting only after the decoder's commands have all finished.

PyTorch Profiler Integration¶

With the PyTorch profiler enabled as well, the framework operations and the device execution land in one chrometrace. Where the runtime dispatches device work asynchronously, though, the device work does not necessarily fall inside the wall-clock window of the function that launched it — and time alone is then not enough to tell which operation caused which device work.

launch_id is what ties the two levels together in that case. An rbln_run slice marks the point on the host where a runtime call begins, and every device slice that call produced carries the same launch_id. Perfetto draws that correspondence as flow arrows, so selecting a device slice leads back to the host call that launched it even when the two do not overlap in time, and the same arrows can be followed in the other direction.