Performance Optimization

This tutorial covers techniques to maximize performance using Tensor Core usage.

Tensor Core Optimization

Understanding Tensor Cores

Modern NVIDIA GPUs (V100, A100, RTX series) include Tensor Cores that can dramatically accelerate matrix operations when certain conditions are met:

• Input dimensions are multiples of 8 (FP16) or 16 (INT8)

• Data types are FP16, BF16, or INT8

• Matrix operations are sufficiently large

def optimize_for_tensor_cores(in_features, out_features, num_terms, low_rank):
    """Optimize dimensions for Tensor Core usage."""

    # Round dimensions to multiples of 8 for FP16 Tensor Cores
    def round_to_multiple(x, multiple=8):
        return ((x + multiple - 1) // multiple) * multiple

    optimized_in = round_to_multiple(in_features)
    optimized_out = round_to_multiple(out_features)
    optimized_rank = round_to_multiple(low_rank)

    # Ensure we don't exceed original dimensions significantly
    if optimized_in > in_features * 1.2:
        optimized_in = in_features
    if optimized_out > out_features * 1.2:
        optimized_out = out_features

    return optimized_in, optimized_out, num_terms, optimized_rank

# Example: Create Tensor Core optimized layers
class TensorCoreOptimizedModel(nn.Module):
    def __init__(self, layer_configs):
        super().__init__()

        layers = []
        for config in layer_configs:
            in_feat, out_feat, terms, rank = optimize_for_tensor_cores(
                config['in_features'], config['out_features'],
                config['num_terms'], config['low_rank']
            )

            layer = pr.nn.SKLinear(in_feat, out_feat, num_terms=terms, low_rank=rank)
            layers.extend([layer, nn.ReLU()])

        self.network = nn.Sequential(*layers)

    def forward(self, x):
        return self.network(x)

# Create optimized model
configs = [
    {'in_features': 1000, 'out_features': 500, 'num_terms': 8, 'low_rank': 64},
    {'in_features': 500, 'out_features': 250, 'num_terms': 6, 'low_rank': 48},
    {'in_features': 250, 'out_features': 10, 'num_terms': 4, 'low_rank': 32}
]

optimized_model = TensorCoreOptimizedModel(configs).to(device)

The next tutorial will cover custom sketching implementations.