Implementation:Turboderp org Exllamav2 Ext TP H

Knowledge Sources	Turboderp_org_Exllamav2
Domains	Tensor_Parallelism, C_Extension
Last Updated	2026-02-15 00:00 GMT

Overview

C++ header defining the ExtTPContext class and the tensor parallelism communication API for distributing inference computation across multiple GPUs in ExLlamaV2.

Description

ext_tp.h provides the infrastructure for tensor-parallel (TP) inference across multiple CUDA devices. It defines broadcast type constants, the context class, and communication functions.

Broadcast type constants identify the tensor category being communicated:

• BROADCAST_KV (0) -- Key/Value projections
• BROADCAST_ID (1) -- Intermediate data
• BROADCAST_VC (2) -- Value cache data
• BROADCAST_RS (3) -- Residual stream
• BROADCAST_Q (4) -- Query projections

ExtTPContext is the central context class managing tensor-parallel state:

• kv_split, id_split, vc_split, rs_split, q_split -- Vectors of tuples (device, start, end) defining how each broadcast type is partitioned across devices.
• pinned_temp -- Pinned host memory buffers used as staging areas for cross-device transfers.
• pinned_size -- Size of each pinned buffer.
• streams -- Per-device CUDA streams for asynchronous execution.
• all_devices -- List of all participating device indices.
• thread_pool -- A ThreadPool instance for dispatching parallel device operations.
• tp_data -- Pointer to ExtTPData containing CUDA-level TP state.
• mapped_globals -- Mapped memory for cross-device globals.
• sync_events -- CUDA events for stream synchronization.

make_tp_context is the factory function that constructs an ExtTPContext from split configurations, pinned temp tensors, and stream handles, returning an opaque uintptr_t handle.

Communication functions:

• tp_broadcast(tp_context, buffer, source, broadcast_type, targets, dim, t_device) -- Broadcasts a tensor from a source to multiple target devices, splitting along the specified dimension according to the broadcast type configuration.
• tp_gather(tp_context, buffer, inputs, broadcast_type, targets, broadcast_type_target, dim, t_device) -- Gathers tensor slices from multiple devices into target tensors.
• tp_gather_barrier(...) -- Same as tp_gather but with an additional Barrier parameter for synchronization before the gather completes.
• tp_cross_device_barrier(tp_context, broadcast_type, t_device, stage, next_stage) -- Inserts a synchronization point across devices, with optional stage tracking for pipeline coordination.
• tp_all_reduce(tp_context, buffer, tensors, residuals) -- Performs an all-reduce sum operation across devices, adding results into the residual tensors.

Usage

The TP context is created once during model initialization when tensor parallelism is enabled. The communication functions are called by tp_attn_forward_* and tp_mlp_forward_ to distribute attention and MLP computations across GPUs. The split configurations determine how weight matrices and activations are partitioned.

Code Reference

Source Location

• Repository: Turboderp_org_Exllamav2
• File: exllamav2/exllamav2_ext/ext_tp.h
• Lines: 1-124

Signature

#define BROADCAST_KV 0
#define BROADCAST_ID 1
#define BROADCAST_VC 2
#define BROADCAST_RS 3
#define BROADCAST_Q  4

class ExtTPContext
{
public:
    std::vector<std::tuple<int, int, int>> kv_split;
    std::vector<std::tuple<int, int, int>> id_split;
    std::vector<std::tuple<int, int, int>> vc_split;
    std::vector<std::tuple<int, int, int>> rs_split;
    std::vector<std::tuple<int, int, int>> q_split;
    std::vector<void*> pinned_temp;
    size_t pinned_size;
    std::vector<cudaStream_t> streams;
    std::vector<int> all_devices;
    ThreadPool* thread_pool;
    ExtTPData* tp_data;
    void* mapped_globals;
    std::vector<cudaEvent_t> sync_events;

    ExtTPContext(
        std::vector<std::tuple<int, int, int>> _kv_split,
        std::vector<std::tuple<int, int, int>> _id_split,
        std::vector<std::tuple<int, int, int>> _vc_split,
        std::vector<std::tuple<int, int, int>> _rs_split,
        std::vector<std::tuple<int, int, int>> _q_split,
        std::vector<torch::Tensor> _pinned_temp,
        std::vector<cudaStream_t> _streams
    );
    ~ExtTPContext();
};

uintptr_t make_tp_context(
    const std::vector<std::tuple<int, int, int>> kv_split,
    const std::vector<std::tuple<int, int, int>> id_split,
    const std::vector<std::tuple<int, int, int>> vc_split,
    const std::vector<std::tuple<int, int, int>> rs_split,
    const std::vector<std::tuple<int, int, int>> q_split,
    std::vector<torch::Tensor> pinned_temp,
    std::vector<uintptr_t> streams
);

void tp_broadcast(
    uintptr_t tp_context, int buffer,
    torch::Tensor source, int broadcast_type,
    const std::vector<torch::Tensor>& targets,
    int dim, int t_device = -1
);

void tp_gather(
    uintptr_t tp_context, int buffer,
    const std::vector<torch::Tensor>& inputs, int broadcast_type,
    const std::vector<torch::Tensor>& targets, int broadcast_type_target,
    int dim, int t_device = -1
);

void tp_gather_barrier(
    uintptr_t tp_context, int buffer,
    const std::vector<torch::Tensor>& inputs, int broadcast_type,
    const std::vector<torch::Tensor>& targets, int broadcast_type_target,
    int dim, int t_device = -1, Barrier* barrier = nullptr
);

void tp_cross_device_barrier(
    uintptr_t tp_context, int broadcast_type,
    int t_device = -1, int stage = -1, int next_stage = -1
);

void tp_all_reduce(
    uintptr_t tp_context, int buffer,
    const std::vector<torch::Tensor>& tensors,
    const std::vector<torch::Tensor>& residuals
);

Import

from exllamav2 import exllamav2_ext as ext_c

# Create tensor parallelism context
tp_handle = ext_c.make_tp_context(
    kv_split, id_split, vc_split, rs_split, q_split,
    pinned_temp_tensors, stream_ptrs
)

I/O Contract

Function	Key Parameters	Output	Description
make_tp_context	5 split configs, pinned buffers, streams	uintptr_t handle	Creates TP context with device split configurations
tp_broadcast	source tensor, broadcast_type, target tensors, dim	Targets filled with sliced data	Distributes source tensor slices to devices
tp_gather	input tensors, broadcast_type, target tensors, dim	Targets filled with gathered data	Collects tensor slices from devices into targets
tp_gather_barrier	Same as tp_gather + Barrier*	Targets filled after barrier sync	Gather with synchronization barrier
tp_cross_device_barrier	broadcast_type, stage info	Synchronization point	Blocks until all devices reach the barrier
tp_all_reduce	tensor list, residual list	Residuals updated with summed results	All-reduce sum across devices into residuals

Split Tuple Format	Description
(device_id, start_index, end_index)	Defines which slice of a dimension each device owns

Usage Examples

from exllamav2 import exllamav2_ext as ext_c

# Create split configurations for 2 GPUs
kv_split = [(0, 0, 2048), (1, 2048, 4096)]  # (device, start, end)
id_split = [(0, 0, 2048), (1, 2048, 4096)]
vc_split = [(0, 0, 2048), (1, 2048, 4096)]
rs_split = [(0, 0, 4096), (1, 0, 4096)]     # residual: full on both
q_split  = [(0, 0, 2048), (1, 2048, 4096)]

tp_ctx = ext_c.make_tp_context(
    kv_split, id_split, vc_split, rs_split, q_split,
    pinned_temp, stream_ptrs
)

# Broadcast hidden states to all devices
ext_c.tp_broadcast(tp_ctx, 0, hidden_states, ext_c.BROADCAST_RS, device_tensors, dim=1)

# All-reduce after MLP
ext_c.tp_all_reduce(tp_ctx, 0, output_tensors, residual_tensors)

Related Pages

• Turboderp_org_Exllamav2_ThreadPool -- ThreadPool and Barrier used by ExtTPContext
• Turboderp_org_Exllamav2_Ext_QAttn_H -- Tensor-parallel attention functions (tp_attn_forward_*)
• Turboderp_org_Exllamav2_Ext_QMLP_H -- Tensor-parallel MLP function (tp_mlp_forward_)
• Turboderp_org_Exllamav2_Ext_Norm -- Tensor-parallel RMS norm (rms_norm_tp)