[Feature Request] Support dynamic token counts in TP communication/GEMM overlap

[ewan0x79]

Problem

Our use case is packed-sequence training with THD-format inputs and dynamically sized microbatches. In TE's THD layout, T is the total number of packed tokens in a microbatch. As the microbatch size and sequence lengths change, the flattened activations used by TP communication/GEMM overlap have shape [T, hidden_size], and T can differ from one microbatch to the next.

The TP-overlap UserBuffer is initialized with a fixed shape. In the v2.11 code on which my prototype is based, overlap paths use the registered buffer size for tensor views and communication counts. As a result, they cannot directly operate on a smaller [T, hidden_size] tensor while retaining an allocation registered for the maximum token count.

I would like to support this workload without allocating and registering new UserBuffers for different microbatch sizes or padding every microbatch to the maximum token count. I have a prototype based on a fixed-capacity allocation and a per-microbatch active shape, described below for design discussion.

This is related to #1303, which asked whether TP overlap supports variable sequence length. This issue adds a concrete workload and an implementation proposal.

Proposed API

# Distributed setup and overlap configuration are omitted here.
# Existing initialization uses the maximum token capacity.
te.initialize_ub(
    [max_tokens, hidden_size],
    tp_size,
    quantization_modes=[te.UserBufferQuantizationMode.NONE],
    dtype=torch.bfloat16,
)

# Before processing a microbatch with fewer tokens:
te.set_ub_active_shape([active_tokens, hidden_size])
# Run the model with tensors sized for active_tokens.

The prototype exposes set_ub_active_shape() on the PyTorch UserBuffer manager. It also exposes set_buffer_active_shape() and active/capacity shape getters on individual CommOverlap and CommOverlapP2P communicators. The manager updates its unquantized bulk, pipeline, and ring-exchange communicators; quantized and external UserBuffers are left at their capacity shape.

Implementation approach in the prototype

Keep allocation capacity separate from the active shape. initialize_ub() registers a buffer sized for the maximum token count. Each communicator stores that capacity shape and a separate active shape. set_buffer_active_shape() updates the logical shape and tensor views; it does not allocate, free, or re-register the underlying UserBuffer. The active data occupies a contiguous prefix of the registered allocation.

Use active sizes throughout the overlap paths. For bulk AG/RS, buffer views, local-chunk offsets, copy-size checks, and communication counts use the active element count rather than the registered capacity. For pipeline and ring exchange, the P2P layout rebuilds its per-rank chunk views and byte offsets using the active first dimension. Reduce-scatter retains its extra P2P chunks within the original allocation.

Constrain the initial behavior. The active shape must be two-dimensional and nonzero, fit within the registered capacity, keep the hidden dimension unchanged, and have a first dimension divisible by TP size. The current manager API applies one active shape to its eligible communicators. Participating TP ranks must therefore use the same active token count before the corresponding overlap operations. The prototype targets unquantized, 2-byte/BF16 UserBuffers for bulk, pipeline, and ring-exchange overlap; dynamic FP8/quantized UserBuffers are outside this initial scope.

The prototype is based on TE v2.11 and has not yet been ported or validated against current main.

Questions for maintainers

1. Would Transformer Engine maintainers be interested in supporting TP communication/GEMM overlap when the token count varies between microbatches, as in the THD packed-sequence workload above?
2. If so, would registering a UserBuffer at maximum capacity and updating its active shape for each microbatch be an acceptable direction for an upstream contribution?

If the overall direction makes sense, I would also appreciate guidance on whether a manager-level update is the right API and whether BF16 bulk, pipeline, and ring-exchange paths are a useful first scope.

I have a prototype implementation and can port it to current main and submit a focused PR if this direction makes sense.