A Unit Mismatch Bug Discovered Mid-Debugging: Tracing SGLang's Memory Calculation

In the middle of a deep debugging session on an OOM (out-of-memory) failure, the assistant made a critical discovery: a potential unit mismatch in SGLang's memory allocation formula. Message [msg 5151] captures the exact moment of this realization — the assistant tracing through the codebase, connecting two pieces of information that, when placed side by side, reveal an inconsistency that could explain the mysterious OOM.

The Context: Custom Allreduce on PCIe Triggers an OOM

The assistant had been systematically testing allreduce optimization strategies for an 8×RTX PRO 6000 Blackwell system connected via PCIe Gen5. After FlashInfer allreduce fusion failed (SM120 not supported), Torch symmetric memory failed (SM120 not in the architecture lookup table), and Expert Parallelism failed (assertion error and OOM), the assistant turned to the custom allreduce kernel already compiled in sgl-kernel. By setting SGLANG_CUSTOM_ALLREDUCE_ALGO=1stage and SGLANG_FORCE_CUSTOM_AR_PCIE=1, the assistant attempted to force the one-stage allreduce kernel to work on PCIe — bypassing the NVLink check that normally prevents it from running on non-NVLink systems.

The result was an immediate OOM during server startup ([msg 5140]). The error trace showed the scheduler crashing during init_model_worker, which is where memory pools are allocated. This was puzzling because the custom allreduce buffers themselves are small (~24 MB per GPU). The assistant began tracing through the memory allocation code to understand why enabling custom allreduce could cause an OOM when the baseline (without custom allreduce) worked fine.

The Message: Connecting Two Pieces of the Puzzle

Message [msg 5151] is the turning point in this investigation. The assistant writes:

So get_available_gpu_memory returns free_gpu_memory / (1 << 30) — in GiB. And total_gpu_memory in the calling code is in bytes (raw from pytorch). Let me re-check the calling code:

This is followed by a bash command that greps for total_gpu_memory in the relevant source file, revealing:

116:    def profile_max_num_token(self: ModelRunner, total_gpu_memory: int):
143:        rest_memory = available_gpu_memory - total_gpu_memory * (
339:    def init_memory_pool(self: ModelRunner, total_gpu_memory: int):
342:        self.max_total_num_tokens = self.profile_max_num_token(total_gpu_memory)

The message is deceptively short — just a few lines of reasoning and a single bash command. But it represents a significant conceptual leap. The assistant has been reading through the memory allocation code in previous messages ([msg 5145], [msg 5147], [msg 5148], [msg 5149], [msg 5150]), tracing the flow from init_memory_pool through profile_max_num_token to get_available_gpu_memory. In [msg 5150], the assistant found the critical line in get_available_gpu_memory:

return free_gpu_memory / (1 << 30)

This function returns GiB (gibibytes), not bytes. But the formula in profile_max_num_token at line 143 computes:

rest_memory = available_gpu_memory - total_gpu_memory * (1 - self.mem_fraction_static)

If available_gpu_memory is in GiB and total_gpu_memory is in bytes, this subtraction produces a nonsensical result — a small number (GiB) minus a huge number (bytes) yields a large negative value. The final line of the function then does int(rest_memory * (1 &lt;&lt; 30)) // cell_size, which would amplify the error.

Why This Matters: The Memory Calculation Chain

The memory calculation chain in SGLang works as follows:

  1. init_memory_pool(total_gpu_memory) is called during server initialization, receiving the total GPU memory (presumably in bytes from torch.cuda.get_device_properties().total_memory).
  2. It calls profile_max_num_token(total_gpu_memory), which computes how many tokens the KV cache can hold.
  3. Inside profile_max_num_token, get_available_gpu_memory() is called, which returns free memory in GiB.
  4. The formula rest_memory = available_gpu_memory - total_gpu_memory * (1 - mem_fraction_static) mixes these two units.
  5. The result is converted back with int(rest_memory * (1 &lt;&lt; 30)) to get bytes for the token count calculation. The assistant's discovery is that step 4 mixes GiB and bytes. If total_gpu_memory is indeed in bytes (as the assistant suspects, calling it "raw from pytorch"), then the formula is buggy. The correct formula should either convert total_gpu_memory to GiB first, or convert available_gpu_memory to bytes before the subtraction.

The Reasoning Process: Tracing Through Layers of Code

The assistant's path to this discovery is a textbook example of systematic debugging. Starting from the OOM crash in [msg 5140], the assistant worked backward through the code:

  1. Observing the symptom: The custom allreduce server OOM'd during initialization, while the baseline server (without custom allreduce) worked fine with the same --mem-fraction-static 0.55 setting.
  2. Checking available memory: In [msg 5142], the assistant compared memory values: the custom allreduce server showed 21.70 GB available after weight loading, while the baseline showed 21.71 GB. The 10 MB difference was too small to explain the OOM, suggesting the problem wasn't the custom allreduce buffers themselves.
  3. Examining the KV cache allocation: In [msg 5143], the assistant checked the baseline's KV cache size (10.42 GB) and remaining memory (9.22 GB), confirming the baseline was comfortably within limits.
  4. Reading the memory pool code: In [msg 5145] through [msg 5150], the assistant methodically read through the memory allocation functions, line by line, using sed to extract specific ranges. This is the critical investigative work — reading the source code of the running system to understand what it actually does, rather than guessing.
  5. The moment of insight: In [msg 5150], the assistant found the return statement of get_available_gpu_memory: return free_gpu_memory / (1 &lt;&lt; 30). This immediately triggered the realization that the function returns GiB. The assistant then, in [msg 5151], connected this to the calling code's use of total_gpu_memory, which comes from PyTorch's raw memory functions and is therefore in bytes.

Input Knowledge Required

To understand this message, the reader needs:

  1. Python and PyTorch basics: Understanding that torch.cuda.mem_get_info() returns (free, total) in bytes, and that (1 &lt;&lt; 30) is 1024³ (a gibibyte).
  2. SGLang's architecture: Knowledge that SGLang uses tensor parallelism (TP) across multiple GPUs, with a KV cache memory pool that must be sized before serving begins. The --mem-fraction-static parameter controls what fraction of total GPU memory is reserved for the KV cache and static allocations.
  3. The debugging context: The assistant was investigating why enabling custom allreduce caused an OOM. The memory calculation formula was the hypothesized root cause — a unit mismatch that would produce wildly incorrect token counts.
  4. The code structure: Understanding that profile_max_num_token is called from init_memory_pool, and that total_gpu_memory is passed as a parameter from the caller (which gets it from PyTorch's device properties).

Assumptions Made

The assistant makes a key assumption in this message: that total_gpu_memory in the calling code is "in bytes (raw from pytorch)". This is a reasonable assumption — PyTorch's torch.cuda.get_device_properties(device).total_memory returns bytes, and the function signature def profile_max_num_token(self: ModelRunner, total_gpu_memory: int) doesn't specify units. However, the assistant hasn't yet verified this by reading the caller code. The message ends with "Let me re-check the calling code" — the assistant is about to verify this assumption.

This assumption is well-founded but not yet confirmed. The assistant's reasoning is:

Potential Mistakes or Incorrect Assumptions

There are several ways the assistant's hypothesis could be wrong:

  1. total_gpu_memory might already be in GiB: The caller might convert bytes to GiB before passing the parameter. The function signature total_gpu_memory: int doesn't specify units, and the conversion could happen upstream.
  2. The formula might intentionally mix units: This seems unlikely — it would be a clear bug — but it's possible that available_gpu_memory is later converted back to bytes in a way that compensates. However, the formula as written would produce a negative rest_memory for any reasonable values (21.71 GiB - 96 GiB * 0.45 = 21.71 - 43.2 = -21.49), which would then be multiplied by (1 &lt;&lt; 30) to produce a huge negative number, or clamped to zero.
  3. The OOM might have a different root cause: Even if the unit mismatch exists, it might not be the cause of the OOM. The custom allreduce could be failing for other reasons — IPC buffer allocation failures on PCIe, driver issues, or memory fragmentation.
  4. The mem_fraction_static calculation might work differently: The assistant's interpretation of the formula might be incomplete. There could be additional logic that handles the negative case or adjusts the calculation. The assistant's next steps would be to verify the assumption by reading the caller code that passes total_gpu_memory to init_memory_pool, and then to test whether fixing the unit mismatch resolves the OOM.