Workloads and measurement were driven by one of NVIDIA’s benchmarking tools, genai‑perf; the accompanying charts were produced directly from benchmark.py with no manual post‑processing. The GitHub repository for the benchmark code and raw data can be found in this GitHub repo. genai-perf measures several useful metrics, such as Time-To-First-Token (TTFT), Inter-Token Latency, and Tokens/Second - all driven by synthetically generated prompts for various patterned workloads.

Test Methodology

Throughput Results

200 → 200 tokens

• H100 scaled almost linearly up to 500 users, peaking at ≈ 7,000 TPS.
• A100 saturated near ≈ 570 TPS and 50 users; additional users primarily increased queueing delay.

This corresponds to a ≈ 12–14× throughput advantage for the H100 configuration across the sweep.

1,000 → 200 tokens

Longer inputs magnify memory pressure during decoding, yet the relative gap remains wide:

• H100 delivered ≈ 2,600 TPS at 250 concurrent users.
• A100 remained under ≈ 230 TPS at the same load.

Latency Under Load

200 → 200 tokens

• H100 preserved a less than 5s TTFT up to 500 simultaneous sessions.
• A100 quickly surpasses >10s TTFT at 100 users and climbs quickly under high load.

1,000 → 200 tokens

• H100 maintains <5s TTFT latency out to the full user sweep - indiciating, additional capacity beyond 500 users.
• A100 TTFT quickly spikes 10-20 concurrent users above 5s TTFT latency.

Discussion

1. Hopper has numerous enhancements over Ampere aside from just faster clock-speeds, additional SMs, and memory bandwidth with HBM3 vs HBM2e. Those enhancements include Tensor Memory Accelerators (TMAs), FP8, and Transformer Engine. All of these lead to significantly increased throughput and lower latencies at higher user counts.
2. In most cases, H100s will be acheive target latency and concurrent user count requirements at a lower price than A100s. Taking the current cost of an on-demand AWS p4de.24xlarge instance with 8xA100s compared to the p5.48xlarge, the cost increase the two is 2.4x, but the throughput increase can be as high as 14x, which justifies the premium. Likewise, it would take 13 instances p4de.24xlarge instances to match the throughput of the p5.48xlarge instances.
3. Longer input & output sequence lengths will increase latency and decrease concurrency/throughput.
4. Blackwell will offer a similar and more pronounced multiple increase in throughput compared to Hopper. Hopefully, I will be in a position to do Blackwell testing soon as hardware becomes more available.
5. Although the llama-3.3-70b model will fit on 2x80GB A100s or H100s at bfloat16 precision, it leaves very little room for a KV Cache, so the supported minimum per the NVIDIA NIM documentation is 4x80GB A100s or H100s.

Key Takeaways

1. Hopper’s advantage is multiplicative, not incremental. In interactive‑LLM scenarios the H100 system provided an order‑of‑magnitude gain in both throughput and p50 latency.
2. Latency SLOs drive hardware choice. If your service promises sub‑second first‑token latency under real‑world concurrency, A100 will need disproportionate over‑provisioning or sharding; H100 meets the target comfortably.
3. Identify Input/Outut Sequence Length (ISL/OSL) based on workload. Input/output sequence lengths vary widely based on task, such as summarization. Reasoning models also change the dynamic with possibly far more output tokens for what might be a relatively short input sequence.
4. bf16 is the sensible default. It balances memory footprint with performance on both architectures and is the path of least resistance inside NIM. At this time fp8 or int4 profiles are not available from the llama-3.3-70b-instruct model with NVIDIA NIM.
5. Measure under expected concurrency. Single‑request benchmarks tell an incomplete story; backlog effects dominate once users count climbs. 5 seconds is my personal rule of thumb for TTFT.

Reproducing the Experiment

The benchmarks can be reproduced by running using the start_nim.sh & benchmark.sh scripts provided in the above GitHub repository. The “NVIDIA NIM LLMs Benchmarking” documentation on genai-perf is also an excellent resource with some great insights into the various benchmarking metrics, as well as some of NVIDIA’s own benchmarking data with various ISL/OSLs and NIM profiles.

Secondly, there’s the point of problem complexity, which is the subject of this post. I recently fixed a bug in the mlx-lm GPT2 model implementation that I submitted to the MLX project last year. If you’re not familiar with MLX, it is an ML and deep‑learning framework similar to PyTorch, but specific to Apple Silicon. The pull request with information on the bug can be viewed here. In summary, when generating tokens from the GPT2 model using the inference code that includes a KV Cache in mlx-lm, the output was garbled with repetitive text for the smaller GPT2 models like the 128M “small” version. At the time I submitted the original PR, I was mostly testing with the gpt2-xl 1.6B parameter version, which generated output more or less identical to the HuggingFace transformers GPT2 implementation.

Upon analyzing the code, I noticed that positional embeddings were not being added after the initial forward pass of the model, which meant any subsequent tokens effectively had a position of 0 - not ideal! The smaller GPT2 models seem to be more sensitive to the positional embedding layer, while the larger GPT2 models like gpt2-xl still generate acceptable output without it. I would assume that during training the gpt2-xl model with its additional parameters learned positional information in the transformer block layers.

In any case, the fix was relatively simple to implement, but harder to spot initially, so I decided to test various LLMs to see if they would have identified the cause or not. The results were interesting. The following prompt was used for testing:

the token generation w/ KV cache for gpt2 implementation generates garbage output after the first token. W/o KV cache it works fine. Here's the implementation code:

# Copyright © 2023-2024 Apple Inc.

from dataclasses import dataclass
from typing import Any, Dict, Optional, Tuple, Union

import mlx.core as mx
import mlx.nn as nn
import numpy as np

from .base import BaseModelArgs, create_attention_mask, scaled_dot_product_attention


@dataclass
class ModelArgs(BaseModelArgs):
    model_type: str
    n_ctx: int
    n_embd: int
    n_head: int
    n_layer: int
    n_positions: int
    layer_norm_epsilon: float
    vocab_size: int
    num_key_value_heads: int = None

    def __post_init__(self):
        if self.num_key_value_heads is None:
            self.num_key_value_heads = self.n_head


class Attention(nn.Module):
    def __init__(self, args: ModelArgs):
        super().__init__()

        assert args.n_embd % args.n_head == 0, "n_embd must be divisible by n_head"

        self.n_embd = args.n_embd
        self.n_head = args.n_head
        self.head_dim = self.n_embd // self.n_head

        self.scale = self.head_dim**-0.5

        self.c_attn = nn.Linear(self.n_embd, 3 * self.n_embd, bias=True)
        self.c_proj = nn.Linear(self.n_embd, self.n_embd, bias=True)

    def __call__(
        self,
        x: mx.array,
        mask: Optional[mx.array] = None,
        cache: Optional[Any] = None,
    ) -> mx.array:
        B, L, D = x.shape

        qkv = self.c_attn(x)
        queries, keys, values = mx.split(qkv, 3, axis=-1)

        # Prepare the queries, keys and values for the attention computation
        queries = queries.reshape(B, L, self.n_head, -1).transpose(0, 2, 1, 3)
        keys = keys.reshape(B, L, self.n_head, -1).transpose(0, 2, 1, 3)
        values = values.reshape(B, L, self.n_head, -1).transpose(0, 2, 1, 3)

        if cache is not None:
            keys, values = cache.update_and_fetch(keys, values)

        output = scaled_dot_product_attention(
            queries, keys, values, cache=cache, scale=self.scale, mask=mask
        )

        output = output.transpose(0, 2, 1, 3).reshape(B, L, -1)
        return self.c_proj(output)


class MLP(nn.Module):
    def __init__(self, args: ModelArgs):
        super().__init__()

        self.n_embd = args.n_embd
        self.c_fc = nn.Linear(self.n_embd, 4 * self.n_embd)
        self.c_proj = nn.Linear(4 * self.n_embd, self.n_embd)

    def __call__(self, x) -> mx.array:
        return self.c_proj(nn.gelu_approx(self.c_fc(x)))


class TransformerBlock(nn.Module):
    def __init__(self, args: ModelArgs):
        super().__init__()

        self.n_head = args.n_head
        self.n_embd = args.n_embd
        self.layer_norm_epsilon = args.layer_norm_epsilon
        self.attn = Attention(args)
        self.mlp = MLP(args)
        self.ln_1 = nn.LayerNorm(
            self.n_embd,
            eps=self.layer_norm_epsilon,
        )
        self.ln_2 = nn.LayerNorm(self.n_embd, eps=self.layer_norm_epsilon)

    def __call__(
        self,
        x: mx.array,
        mask: Optional[mx.array] = None,
        cache: Optional[Any] = None,
    ) -> mx.array:
        r = self.attn(self.ln_1(x), mask, cache)
        h = x + r
        r = self.mlp(self.ln_2(h))
        out = h + r
        return out


class GPT2Model(nn.Module):
    def __init__(self, args: ModelArgs):
        super().__init__()
        self.n_embd = args.n_embd
        self.n_positions = args.n_positions
        self.vocab_size = args.vocab_size
        self.n_layer = args.n_layer
        self.layer_norm_epsilon = args.layer_norm_epsilon
        assert self.vocab_size > 0
        self.wte = nn.Embedding(self.vocab_size, self.n_embd)
        self.wpe = nn.Embedding(self.n_positions, self.n_embd)
        self.h = [TransformerBlock(args=args) for _ in range(self.n_layer)]
        self.ln_f = nn.LayerNorm(self.n_embd, eps=self.layer_norm_epsilon)

    def __call__(
        self,
        inputs: mx.array,
        mask: mx.array = None,
        cache=None,
    ):
        _, L = inputs.shape

        hidden_states = self.wte(inputs)

        mask = None
        if hidden_states.shape[1] > 1:

            position_ids = mx.array(np.arange(L))
            hidden_states += self.wpe(position_ids)

            if mask is None:
                mask = create_attention_mask(hidden_states, cache)

        if cache is None:
            cache = [None] * len(self.h)

        for layer, c in zip(self.h, cache):
            hidden_states = layer(hidden_states, mask, cache=c)

        return self.ln_f(hidden_states)


class Model(nn.Module):
    def __init__(self, args: ModelArgs):
        super().__init__()
        self.args = args
        self.model_type = args.model_type
        self.model = GPT2Model(args)

    def __call__(
        self,
        inputs: mx.array,
        mask: mx.array = None,
        cache=None,
    ):
        out = self.model(inputs, mask, cache)
        out = self.model.wte.as_linear(out)
        return out

    def sanitize(self, weights):
        new_weights = {}
        for i in range(self.args.n_layer):
            if f"h.{i}.attn.bias" in weights:
                del weights[f"h.{i}.attn.bias"]
            if f"h.{i}.attn.c_attn.weight" in weights:
                weights[f"h.{i}.attn.c_attn.weight"] = weights[
                    f"h.{i}.attn.c_attn.weight"
                ].transpose(1, 0)
            if f"h.{i}.attn.c_proj.weight" in weights:
                weights[f"h.{i}.attn.c_proj.weight"] = weights[
                    f"h.{i}.attn.c_proj.weight"
                ].transpose(1, 0)
            if f"h.{i}.mlp.c_fc.weight" in weights:
                weights[f"h.{i}.mlp.c_fc.weight"] = weights[
                    f"h.{i}.mlp.c_fc.weight"
                ].transpose(1, 0)
            if f"h.{i}.mlp.c_proj.weight" in weights:
                weights[f"h.{i}.mlp.c_proj.weight"] = weights[
                    f"h.{i}.mlp.c_proj.weight"
                ].transpose(1, 0)
        for weight in weights:
            if not weight.startswith("model."):
                new_weights[f"model.{weight}"] = weights[weight]
            else:
                new_weights[weight] = weights[weight]
        return new_weights

    @property
    def layers(self):
        return self.model.h

I wanted to provide as little information as possible in the prompt to not steer the model, but to also relay some of the problematic behaviors. My testing is by no means intended to be comprehensive and the models may fare better or worse for other problems. Two patterns emerge from the testing results:

1. The model correctly identifies the problem as being related to the positional embeddings not being added after the prefill/decode.
2. The model is confused and assumes the KV Cache implementation has a bug or that the KV Cache is being used incorrectly in the attention layers.

Correct solution from o1-2024-12-17:

A key issue is that in your GPT‐2 forward pass, you only add position embeddings and create a causal mask when “hidden_states.shape[1] > 1.” In a typical “one‐token‐at‐a‐time” generation loop (when using a KV cache), each newly generated token has shape [B, 1]. Because of the if hidden_states.shape[1] > 1 check, the code entirely skips adding the new position embedding and never updates the attention mask for that new single token. That mismatch quickly leads to “garbage” outputs after the first step.
In other words, inside GPT2Model.call:

if hidden_states.shape[1] > 1:
	position_ids = mx.array(np.arange(L))
	hidden_states += self.wpe(position_ids)

if mask is None:
    mask = create_attention_mask(hidden_states, cache)

is never invoked on subsequent tokens (where L = 1). To fix this, you must always add the correct position embedding to every token in the sequence—whether you are calling the model on N tokens at once or just on a single newly generated token. You also need to ensure that you create/apply a proper causal mask for incremental decoding.
A standard fix is something along these lines:
• Keep track of how many tokens have been generated so far (e.g., total_seq_len), then compute the new position IDs as range(total_seq_len, total_seq_len + L) rather than just range(L). • Always add position embeddings for the newly generated token (i.e., remove or adjust the if hidden_states.shape[1] > 1 condition). • Always create/apply the causal mask for the new token (so the new position can only attend to previous tokens).
Once that is done, single‐token inference (with KV cache) will produce correct next‐token logits instead of garbled output.

Here is an example of an incorrect solution from llama-3.3-70b-instruct:

# Step-by-step analysis of the problem:
1. **Token Generation Issue**: The issue here seems to be related to the token generation process, specifically when using a KV cache for the GPT2 implementation.
2. **KV Cache Implementation**: The KV cache is used in the attention mechanism to store and update the keys and values for each layer. However, it appears that there might be an issue with how the cache is being updated or accessed.
3. **Garbage Output After First Token**: This suggests that the problem lies in the way the model handles the input sequence beyond the first token.

# Fixed solution:
To fix this issue, we need to review and potentially modify the `cache.update_and_fetch(keys, values)` line in the Attention class's `__call__` method. 

Here's an updated version of the `Attention` class:

class Attention(nn.Module):
    def __init__(self, args: ModelArgs):
        super().__init__()

        assert args.n_embd % args.n_head == 0, "n_embd must be divisible by n_head"

        self.n_embd = args.n_embd
        self.n_head = args.n_head
        self.head_dim = self.n_embd // self.n_head

        self.scale = self.head_dim**-0.5

        self.c_attn = nn.Linear(self.n_embd, 3 * self.n_embd, bias=True)
        self.c_proj = nn.Linear(self.n_embd, self.n_embd, bias=True)

    def __call__(
        self,
        x: mx.array,
        mask: Optional[mx.array] = None,
        cache: Optional[Any] = None,
    ) -> mx.array:
        B, L, D = x.shape

        qkv = self.c_attn(x)
        queries, keys, values = mx.split(qkv, 3, axis=-1)

        # Prepare the queries, keys and values for the attention computation
        queries = queries.reshape(B, L, self.n_head, -1).transpose(0, 2, 1, 3)
        keys = keys.reshape(B, L, self.n_head, -1).transpose(0, 2, 1, 3)
        values = values.reshape(B, L, self.n_head, -1).transpose(0, 2, 1, 3)

        if cache is not None:
            # Check if the cache has been initialized
            if isinstance(cache, dict):
                keys_cache = cache.get('keys', mx.zeros_like(keys))
                values_cache = cache.get('values', mx.zeros_like(values))
                
                # Update the cache with new keys and values
                keys = mx.cat([keys_cache[:, :, -L:, :], keys], dim=-2)
                values = mx.cat([values_cache[:, :, -L:, :], values], dim=-2)
                
                # Update the cache for future use
                cache['keys'] = keys
                cache['values'] = values
                
            else:
                raise ValueError("Cache should be a dictionary")
        else:
            # If no cache is provided, do not update or fetch any cached values
            pass

        output = scaled_dot_product_attention(
            queries, keys, values, scale=self.scale, mask=mask
        )

        output = output.transpose(0, 2, 1, 3).reshape(B, L, -1)
        return self.c_proj(output)

And here is an updated version of the `TransformerBlock` class:

class TransformerBlock(nn.Module):
    def __init__(self, args: ModelArgs):
        super().__init__()

        self.n_head = args.n_head
        self.n_embd = args.n_embd
        self.layer_norm_epsilon = args.layer_norm_epsilon
        self.attn = Attention(args)
        self.mlp = MLP(args)
        self.ln_1 = nn.LayerNorm(
            self.n_embd,
            eps=self.layer_norm_epsilon,
        )
        self.ln_2 = nn.LayerNorm(self.n_embd, eps=self.layer_norm_epsilon)

    def __call__(
        self,
        x: mx.array,
        mask: Optional[mx.array] = None,
        cache: Optional[Any] = None,
    ) -> mx.array:
        r = self.attn(self.ln_1(x), mask, cache)
        h = x + r
        r = self.mlp(self.ln_2(h))
        out = h + r
        
        # Initialize or update the cache for future use
        if cache is None:
            cache = {'keys': mx.zeros((x.shape[0], self.n_head, 0, self.head_dim)), 
                     'values': mx.zeros((x.shape[0], self.n_head, 0, self.head_dim))}
        else:
            pass
            
        return out

# Explanation of changes:
*   **Cache Management**: The primary change made was in how the cache is managed within the `Attention` class. Now, when a cache is provided, it checks if the cache has been initialized and updates it accordingly.
*   **Key and Value Concatenation**: In the updated code, we concatenate the cached keys and values with the new ones before using them for attention computation.

# Tests and example uses:
You can test this implementation by providing a sample input sequence to your `Model` class. Make sure that you have properly initialized all necessary variables and modules.

Here is an example of how you might use it:

args = ModelArgs(model_type='gpt2', n_ctx=1024, n_embd=512, n_head=8, n_layer=12, 
                 n_positions=1024, layer_norm_epsilon=1e-5, vocab_size=50257)

model = Model(args)
input_ids = mx.array(np.random.randint(0, 100, size=(32, 128))) # replace with your own input

output = model(input_ids)
print(output.shape) # Should print (32, 128, 50257)

Here’s a table of my testing results:

In short, reasoning models fared far better than non-reasoning models. Also, I was surprised to see gpt-4.5 fail, while gpt-4.1 succeeded - further proof of why OpenAI may be sunsetting gpt-4.5 so quickly.

I downloaded the Latest release of the library for Mac/Unix, which at the time of this writing is version API 10.34. Once downloaded, I extracted the .zip file:

~/Downloads $ unzip -d twsapi_macunix.1034.02 twsapi_macunix.1034.02.zip

The build documentation for Mac/Unix is available here. The documentation is mostly focused on the building the Python library and the C++ library build documentation is lacking. The API’s Decimal implementation also has an odd dependency on the Intel Decimal Library, which was concerning since I wanted to build this for MacOS on Apple Silicon and Linux/aarch64. I assumed the Intel library is Intel-specific, but much to my pleasant surprise it is not.

After downloading the Intel Decimal Floating-Point Math Library from: https://www.intel.com/content/www/us/en/developer/articles/tool/intel-decimal-floating-point-math-library.html, I extracted that tarball:

Building the Intel Decimal Floating-Point Math Library

~/Downloads $ tar xvzf IntelRDFPMathLib20U2.tar.gz

Interactive Brokers provides steps to build the library contained in IBJts/source/cppclient/Intel_lib_build.txt, which I followed with some necessary modifications below:

IntelRDFPMathLib20U2/LIBRARY modify "makefile":
a)
Line 370:
change
BID_LIB = $(LIB_DIR)/libbid.$A
to
BID_LIB = $(LIB_DIR)/libbid.dylib # or .so when building on Linux

b)
Line 377:
change
	$(AR_CMD) $(AR_OUT)$@ $^
to
	gcc -o $@ $^ -shared

c) 
Line 112:
change
_CFLAGS_OPT := 
to
_CFLAGS_OPT := -fPIC -Wno-implicit-function-declaration # added to avoid issues with newer versions of clang and gcc (>=14), which now treat implicit function declarations as an error

If building on Linux/aarch64, the following error will be encountered about an unknown architecture:

[dlewis@defiant LIBRARY]$ make
makefile.iml_head:356: *** Unknown host architecture aarch64.  Stop.

This can be easily fixed by making some small changes to the makefile.iml_head file around line 347:

# add aarch64 with EFI2 specified for both ARCH_LIST and ARCH_TYPE
ARCH_ALIAS := x86  ia64 EM64T x86_64 i686 amd64 Intel64 sun4u aarch64
ARCH_LIST  := IA32 IA64 EFI2  EFI2   IA32 EFI2  EFI2    EFI2    EFI2
ARCH_TYPE  := IA32 IA64 EFI2  EFI2   IA32 EFI2  EFI2    EFI2    EFI2  
ARCH_TYPES := IA32 IA64 EFI2

Now, the Intel Decimal Library can be built without issues on MacOS or Linux/aarch64:

~/Downloads/IntelRDFPMathLib20U2/LIBRARY $ make CC=gcc CALL_BY_REF=0 GLOBAL_RND=0 GLOBAL_FLAGS=0 UNCHANGED_BINARY_FLAGS=0

Once the build finishes, there will be a .dylib/.so:

~/Downloads/IntelRDFPMathLib20U2/LIBRARY $ ls -l libbid.dylib 
-rwxr-xr-x  1 dlewis  staff  5019408 Feb 11 14:39 libbid.dylib
~/Downloads/IntelRDFPMathLib20U2/LIBRARY $ file libbid.dylib 
libbid.dylib: Mach-O 64-bit dynamically linked shared library arm64

Building the IBKR C++ API Client Library

Back to building the IBKR C++ API client library. The libbid.dylib/.so needs to be copied to IBJts/source/cppclient/client/lib:

~/Downloads/twsapi_macunix.1034.02/IBJts/source/cppclient/client $ mkdir lib
~/Downloads/twsapi_macunix.1034.02/IBJts/source/cppclient/client $ cp ~/Downloads/IntelRDFPMathLib20U2/LIBRARY/libbid.dylib lib/
~/Downloads/twsapi_macunix.1034.02/IBJts/source/cppclient/client $ ls -l lib/libbid.dylib
-rwxr-xr-x  1 dlewis  staff  5019408 Feb 11 14:44 lib/libbid.dylib

I had many issues building with CMake in IBJts/source/cppclient/client, and the sources include a generic makefile, which seemed much easier to modify, so I went that route instead. Here’s my slightly modified makefile:

CXX=g++
CXXFLAGS=-pthread -Wall -Wno-switch -Wno-unused-function -std=c++11 -shared -fPIC
ROOT_DIR=.
BASE_SRC_DIR=${ROOT_DIR}
INCLUDES=-I${ROOT_DIR}
LIB_DIR=lib
LIB_NAME=bid
TARGET=libTwsSocketClient.dylib # or .so if building on Linux

$(TARGET):
	$(CXX) $(CXXFLAGS) $(INCLUDES) $(BASE_SRC_DIR)/*.cpp -L$(LIB_DIR) -l$(LIB_NAME) -o$(TARGET)

clean:
	rm -f $(TARGET) *.o

Once that is modified, make can simply be run without issue:

~/Downloads/twsapi_macunix.1034.02/IBJts/source/cppclient/client $ make
g++ -pthread -Wall -Wno-switch -Wno-unused-function -std=c++11 -shared -fPIC -I. ./*.cpp -Llib -lbid -olibTwsSocketClient.dylib

In preparation for building the sample client IBKR provides, I copied libTwsSocketClient.dylib to lib, where previously libbid.dylib was copied to:

~/Downloads/twsapi_macunix.1034.02/IBJts/source/cppclient/client $ cp libTwsSocketClient.dylib lib/
~/Downloads/twsapi_macunix.1034.02/IBJts/source/cppclient/client $ ls -l lib/
total 12200
-rwxr-xr-x  1 dlewis  staff  1221496 Feb 11 14:57 libTwsSocketClient.dylib
-rwxr-xr-x  1 dlewis  staff  5019408 Feb 11 14:44 libbid.dylib

Building the C++ Sample Application

Now, the client can be built in IBJts/samples/Cpp/TestCppClient. First, some similar makefile modifications need to be made, and again I avoided using CMake because of problems encountered with the build process. It can be made to work, but it is much simpler to just use the provided makefile. Here is my diff of the makefile after modifications:

~/Downloads/twsapi_macunix.1034.02/IBJts/samples/Cpp/TestCppClient $ diff -u makefile makefile.orig
--- makefile	2025-02-11 15:03:37
+++ makefile.orig	2025-02-11 15:02:31
@@ -4,10 +4,10 @@
 BASE_SRC_DIR=${ROOT_DIR}/client
 INCLUDES=-I${BASE_SRC_DIR} -I${ROOT_DIR}
 SOURCE_DIR=${BASE_SRC_DIR}
-SOURCE_LIB=TwsSocketClient
+SOURCE_LIB=libTwsSocketClient.so
 LIB_DIR=$(SOURCE_DIR)/lib
 LIB_NAME_A=libbid.a
-LIB_NAME=bid
+LIB_NAME_SO=libbid.so
 TARGET=TestCppClient
 
 $(TARGET)Static:
@@ -17,7 +17,7 @@
 	$(CXX) $(CXXFLAGS) $(INCLUDES) $(BASE_SRC_DIR)/*.cpp ./*.cpp -L$(LIB_DIR) -l:$(LIB_NAME_SO) -o$(TARGET)
 
 $(TARGET)Dynamic:
-	$(CXX) $(CXXFLAGS) $(INCLUDES)  ./*.cpp -L$(LIB_DIR) -l$(LIB_NAME) -L$(SOURCE_DIR) -l$(SOURCE_LIB) -o$(TARGET)Dynamic
+	$(CXX) $(CXXFLAGS) $(INCLUDES)  ./*.cpp -L$(LIB_DIR) -l:$(LIB_NAME_SO) -L$(SOURCE_DIR) -l:$(SOURCE_LIB) -o$(TARGET)Dynamic
 
 run$(TARGET)Static:
 	./$(TARGET)Static

The makefile can be easily modified or used to generate static application binaries, but in this case I opted for dynamic linking. If you go the statically linked route, then you can also easily make the necessary modifications to the C++ API library makefile in

Finally, we can build and run the TestCppClientDynamic:

~/Downloads/twsapi_macunix.1034.02/IBJts/samples/Cpp/TestCppClient $ make TestCppClientDynamic
g++ -pthread -Wall -Wno-switch -Wpedantic -Wno-unused-function -std=c++11 -I../../../source/cppclient/client -I../../../source/cppclient  ./*.cpp -L../../../source/cppclient/client/lib -lbid -L../../../source/cppclient/client -lTwsSocketClient -oTestCppClientDynamic
~/Downloads/twsapi_macunix.1034.02/IBJts/samples/Cpp/TestCppClient $ DYLD_LIBRARY_PATH=../../../source/cppclient/client/lib ./TestCppClientDynamic 
Start of C++ Socket Client Test 0
Attempt 1 of 50
Connecting to 127.0.0.1:7496 clientId:0
Error. Id: -1, Time: Tue Feb 11 15:08:17 2025, Code: 502, Msg: Couldn't connect to TWS. Confirm that "Enable ActiveX and Socket Clients" is enabled and connection port is the same as "Socket Port" on the TWS "Edit->Global Configuration...->API->Settings" menu. Live Trading ports: TWS: 7496; IB Gateway: 4001. Simulated Trading ports for new installations of version 954.1 or newer:  TWS: 7497; IB Gateway: 4002
Cannot connect to 127.0.0.1:7496 clientId:0
Sleeping 10 seconds before next attempt

Success! Hopefully, this will be helpful to someone looking to build the IBKR C++ API library on MacOS or Linux/aarch64.

Understanding how CUDA and efficient kernels work under the hood enables us to make informed decisions about model architecture, optimize critical operations, and squeeze maximum performance from our GPU hardware. One powerful technique is using PyTorch CUDA extensions to leverage custom CUDA kernels. In this series of posts, we’ll explore how to create, profile, and optimize these extensions, using a simple, straightforward matrix multiplication example as a guide.

The Power of Custom CUDA Kernels

PyTorch allows us to write custom CUDA kernels and integrate them seamlessly into our Python code. This capability is particularly useful when we need to optimize specific operations that are critical to an application’s performance.

Our Example: Matrix Multiplication

We’ll use matrix multiplication as our example. While PyTorch already has highly optimized matrix multiplication routines, implementing our own allows us to understand the process and potentially optimize for specific use cases. For the first post in this series, we’ll start with 3 basic matrix multiplication kernels that can be optimized and vary based on quality of implementation: the Element-wise Kernel, the Row-wise Kernel, and the Column-wise Kernel. Each of these approaches distributes the computation across GPU threads in a different manner, offering varying levels of parallelism and memory access patterns.

Creating the CUDA Extension

Let’s break down the key components of our CUDA extension:

The CUDA Kernels (matrixMultiply.cu)

// element-wise matrix multiplication
__global__ void matrixMulKernel(float *m, float *n, float *p, int size)
{
    int i = blockDim.y * blockIdx.y + threadIdx.y;
    int j = blockDim.x * blockIdx.x + threadIdx.x;
    if (i < size && j < size)
    {
        float pValue = 0;
        for (int k = 0; k < size; ++k)
        {
            pValue += m[i * size + k] * n[k * size + j];
        }
        p[i * size + j] = pValue;
    }
}

// row-wise matrix multiplication
// thread 0:
// p_0,0 p[0]
// p_0,1 p[1]
// ..
// thread 1:
// p_1,0 p[4]
// p_1,1 p[5]
// ..
__global__ void matrixMulKernelRow(float *m, float *n, float *p, int size)
{
    int row = blockDim.x * blockIdx.x + threadIdx.x;

    if (row < size)
    {
        for (int col = 0; col < size; ++col)
        {
            float pValue = 0;
            for (int i = 0; i < size; ++i)
            {
                pValue += m[row * size + i] * n[i * size + col];
            }
            p[row * size + col] = pValue;
        }
    }
}

// column-wise matrix multiplication
// thread 0:
// p_0,0 p[0]
// p_1,0 p[4]
// ..
// thread 1:
// p_0,1 p[1]
// p_1,1 p[5]
// ..
__global__ void matrixMulKernelCol(float *m, float *n, float *p, int size)
{
    int col = blockDim.x * blockIdx.x + threadIdx.x;

    if (col < size)
    {
        for (int row = 0; row < size; ++row)
        {
            float pValue = 0;
            for (int i = 0; i < size; ++i)
            {
                pValue += m[row * size + i] * n[i * size + col];
            }
            p[row * size + col] = pValue;
        }
    }
}

The Element-wise kernel, while simple, offers true parallelism by utilizing the GPU’s massive thread capacity. However, the Row-wise and Column-wise kernels, despite appearing to offer a different parallelization strategy, actually introduce serialization and will perform poorly. They fail to fully utilize the GPU’s parallel architecture, with each thread sequentially computing an entire row or column. This approach negates much of the performance benefit of using a GPU. The Element-wise kernel, while parallel, is far from optimal. It lacks shared memory usage, efficient memory access patterns, and coalescing, all critical for peak GPU performance. In a future post we’ll profile these kernels to quantify their performance differences.

The C++ Wrapper (torchMatrixMultiply.cu)

__global__ void matrixMulKernel(float *m, float *n, float *p, int size);
__global__ void matrixMulKernelRow(float *m, float *n, float *p, int size);
__global__ void matrixMulKernelCol(float *m, float *n, float *p, int size);

using KernelFunc = void (*)(float *, float *, float *, int);

torch::Tensor cuda_matrixMultiply(const torch::Tensor &a, const torch::Tensor &b, KernelFunc kernel)
{
    TORCH_CHECK(a.sizes() == b.sizes());
    TORCH_CHECK(a.dtype() == torch::kFloat);
    TORCH_CHECK(b.dtype() == torch::kFloat);
    TORCH_INTERNAL_ASSERT(a.device().type() == torch::DeviceType::CUDA);
    TORCH_INTERNAL_ASSERT(b.device().type() == torch::DeviceType::CUDA);

    torch::Tensor a_contiguous{a.contiguous()};
    torch::Tensor b_contiguous{b.contiguous()};
    torch::Tensor result{torch::empty(a_contiguous.sizes(), a_contiguous.options())};

    float *a_ptr = a_contiguous.data_ptr<float>();
    float *b_ptr = b_contiguous.data_ptr<float>();
    float *result_ptr = result.data_ptr<float>();

    // Assumes square matrices and we cast to int for simplicity
    // and compatibility with our existing kernel code. In practice,
    // we would need to handle non-square matrices and use an unsigned long
    // to match PyTorch's tensor sizes.
    int dim{static_cast<int>(a.sizes()[0])};

    dim3 blockSize(16, 16);
    dim3 gridSize((dim + blockSize.x - 1) / blockSize.x, (dim + blockSize.y - 1) / blockSize.y);

    kernel<<<gridSize, blockSize>>>(a_ptr, b_ptr, result_ptr, dim);
    checkCudaError(cudaGetLastError(), "Kernel launch failed");

    return result;
}

torch::Tensor matrixMultiply(const torch::Tensor &a, const torch::Tensor &b, const std::optional<std::string> &kernel_type)
{
    if (kernel_type.has_value())
    {
        if (kernel_type == "row")
        {
            return cuda_matrixMultiply(a, b, matrixMulKernelRow);
        }
        else if (kernel_type == "col")
        {
            return cuda_matrixMultiply(a, b, matrixMulKernelCol);
        }
        else
        {
            throw std::invalid_argument("Invalid kernel type");
        }
    }
    else
    {
        return cuda_matrixMultiply(a, b, matrixMulKernel);
    }
}

TORCH_LIBRARY(myextension, m)
{
    m.def("mymatrixmultiply(Tensor a, Tensor b, str? kernel_type = None) -> Tensor");
}

TORCH_LIBRARY_IMPL(myextension, CUDA, m)
{
    m.impl("mymatrixmultiply", TORCH_FN(matrixMultiply));
}

This wrapper handles the conversion between PyTorch tensors and CUDA array pointers, launches each kernel, and registers the function with PyTorch’s dispatcher. We use an optional parameter and function pointer to specify the kernel at runtime, allowing us to switch between the Element-wise, Row-wise, and Column-wise kernels in our Python code.

Building the Extension

Building a CUDA extension requires careful attention to linking and include paths. Our Makefile handles this by specifying the compiler, compiler flags, and paths to the necessary libraries and include directories. It also ensures the build directory exists, compiles the CUDA source files into object files, and links them into a shared library that will be loaded by PyTorch as our extension:

NVCC := nvcc
NVCC_FLAGS := -g -G
BUILD_DIR := build
PYTHON_VER := 3.12
MLENV_DIR := ~/mlenv/lib/python$(PYTHON_VER)/site-packages/torch
TORCH_INCLUDE := $(MLENV_DIR)/include
TORCH_LIB := $(MLENV_DIR)/lib

# Find all .cu files in the current directory
SRCS := $(filter-out torchMatrixMultiply.cu, $(wildcard *.cu))
# Generate corresponding object file names in the build directory
OBJS := $(patsubst %.cu,$(BUILD_DIR)/%.o,$(SRCS))
# Generate executable names from the object files
EXECS := $(patsubst $(BUILD_DIR)/%.o,$(BUILD_DIR)/%,$(OBJS))

# Ensure the build directory exists
$(shell mkdir -p $(BUILD_DIR))

.PHONY: all clean

all: $(EXECS) $(BUILD_DIR)/torchMatrixMultiply.so

# Rule to build executables from object files
$(BUILD_DIR)/%: $(BUILD_DIR)/%.o
	$(NVCC) $(NVCC_FLAGS) $< -o $@

# Rule to compile .cu files into object files
$(BUILD_DIR)/%.o: %.cu
	$(NVCC) $(NVCC_FLAGS) --compiler-options '-fPIC' -c $< -o $@

.SECONDARY: $(OBJS)

$(BUILD_DIR)/torchMatrixMultiply.so: torchMatrixMultiply.cu
	$(NVCC) $(NVCC_FLAGS) -shared --compiler-options '-fPIC' -L ~/mlenv/lib/python3.12/site-packages/torch/lib \
	-lc10 -ltorch_cpu -ltorch -ltorch_python -lc10_cuda -ltorch_cuda \
	-isystem $(TORCH_INCLUDE)/torch/csrc/api/include \
	-isystem $(TORCH_INCLUDE) \
	torchMatrixMultiply.cu $(BUILD_DIR)/matrixMultiply.o -o $(BUILD_DIR)/torchMatrixMultiply.so \
	-DTORCH_API_INCLUDE_EXTENSION_H -DPYBIND11_COMPILER_TYPE=\"_gcc\" -DPYBIND11_STDLIB=\"_libstdcpp\" \
	-DPYBIND11_BUILD_ABI=\"_cxxabi1011\" -isystem /usr/include/python$(PYTHON_VER) -D_GLIBCXX_USE_CXX11_ABI=0 -std=c++17

clean:
	rm -rf $(BUILD_DIR)

Key points:

• We link against PyTorch libraries
• We include PyTorch and Python headers
• We set necessary compilation flags for PyTorch and pybind11

Using the Extension in PyTorch

Once built, we can use our extension in PyTorch:

import torch

torch.ops.load_library("build/torchMatrixMultiply.so")

torch.manual_seed(42)
a = torch.rand(1000, 1000, device="cuda")
print(a)
b = torch.rand(1000, 1000, device="cuda")
print(b)

print("row kernel:")
print(torch.ops.myextension.mymatrixmultiply(a, b, "row"))
print("col kernel:")
print(torch.ops.myextension.mymatrixmultiply(a, b, "col"))
print("default kernel:")
print(torch.ops.myextension.mymatrixmultiply(a, b))

Profiling with NVIDIA Compute Command Line Profiler (ncu)

To optimize our kernel, we need to understand its performance characteristics. NVIDIA’s ncu tool is useful for this:

ncu --kernel-name regex:"matrixMulKernel.*" python torchMatrixMultiply.py

This command profiles our kernel, providing insights into metrics like:

• SM occupancy
• Memory throughput
• Instruction throughput

In the next post, we’ll review the profiling results and optimization strategies, such as those below to improve our kernel’s performance.

Optimization Strategies

Based on the profiling results, we can apply various optimization strategies:

1. Shared Memory: If memory bandwidth is a bottleneck, we can use shared memory to reduce global memory accesses.

2. Loop Unrolling: This can increase instruction-level parallelism.

3. Tiling: Dividing the matrices into smaller tiles can improve cache utilization.

4. Vectorization: Using vector loads and stores can increase memory throughput.

5. Warp-level Primitives: For certain operations, warp-level primitives can be faster than block-level synchronization.

Example GitHub Repository

Complete example code is available on GitHub.

Conclusion

Custom CUDA extensions offer a powerful way to optimize applications utilizing PyTorch. By understanding the nuances of CUDA programming, profiling, and optimization techniques, we can significantly improve the performance of our machine learning workloads.

In future posts, we’ll dive deeper into the profiling process for each of our kernels and explore optimization strategies to enhance their performance.