Enable chunking with 8 chunks (the default threshold of 8192 tokens)

---
title: Row-Parallel Chunking
---
[](){ #row_parallel_chunking }

Row-Parallel Chunking is an optimization method for tensor-parallel inference on Intel® Gaudi® that overlaps computation with communication in `RowParallelLinear` layers. In standard tensor-parallel inference, each `RowParallelLinear` layer performs a matrix multiplication followed by a blocking all-reduce across tensor-parallel ranks.

In a standard `RowParallelLinear` forward pass, the execution is sequential:

1. Compute `output = matmul(input, weight)`.
2. Perform a blocking all-reduce of `output` across tensor-parallel ranks.

With chunking enabled, the forward pass becomes:

1. Split the input into *N* chunks along the token dimension.
2. For each chunk *i*:

   - Compute `output_i = matmul(input_i, weight)`.
   - Launch `all_reduce(output_i)` asynchronously.

3. Wait for all async all-reduce operations to complete.
4. Concatenate the chunk outputs.

This pipelining allows the all-reduce of chunk *i* to run concurrently with the matmul of chunk *i+1*, reducing idle time on both the compute engine and the network interface.

### Chunking Conditions

Chunking is only applied when all of the following conditions are met:

- Chunking is configured: `num_chunks > 1`
- All-reduce is required: `reduce_results` is enabled on the layer
- Tensor parallelism is active: `tp_size > 1`
- Input is sufficiently large: `total_tokens >= chunk_threshold`

When any condition is not met, the layer falls back to the standard single-shot computation.

### Input Handling

The implementation handles both 2D `[tokens, hidden]` and 3D `[batch, seq_len, hidden]` inputs:

- 3D with `seq_len > 1` (prefill): Chunks along the sequence dimension
- 3D with `seq_len == 1` (decode): Chunks along the batch dimension
- 2D: Chunks along the first token dimension

### Recommended Usage

Chunking effectiveness depends on balancing communication overlap benefits against computational overhead.

We recommend this feature for:

- Running with tensor parallelism (`TP > 1`)
- Serving workloads with large batch sizes or long prefill sequences
- Workloads where all-reduce communication is a significant fraction of the prefill time

## Configuration

The feature is controlled by two environment variables:

| Environment Variable | Config Name | Default | Description |
|---|---|---|---|
| `VLLM_ROW_PARALLEL_CHUNKS` | `row_parallel_chunks` | `1` (disabled) | The number of chunks to split the input into. Setting the variable to a value greater than 1 enables chunking. |
| `VLLM_ROW_PARALLEL_CHUNK_THRESHOLD` | `row_parallel_chunk_threshold` | `8192` | The minimum number of tokens required to activate chunking. Inputs below this threshold use the standard path. |

The following example shows how to set these variables to enable chunking with different configurations:

```bash
# Enable chunking with 8 chunks (the default threshold of 8192 tokens)
export VLLM_ROW_PARALLEL_CHUNKS=8

# Enable chunking with 16 chunks and a lower threshold
export VLLM_ROW_PARALLEL_CHUNKS=16
export VLLM_ROW_PARALLEL_CHUNK_THRESHOLD=4096
```

## Performance Characteristics

The speedup ratio is the baseline time divided by the chunked time, where values higher than 1.0 indicate speedup. The following tables show the speedup ratio for an isolated `RowParallelLinear` layer measured across different tensor parallelism sizes, chunk counts, and token counts. To interpret the values, consider the following:

- Values greater than 1.0 indicate a speedup over the non-chunked baseline, for example, 1.5 means 50% faster.
- Values below 1.0 indicate a slowdown due to chunking overhead exceeding the overlap benefit.

Tensor parallelism size equal to 2:

| Tokens | 2 chunks | 4 chunks | 8 chunks | 16 chunks | 32 chunks | 64 chunks |
|-------:|:--------:|:--------:|:--------:|:---------:|:---------:|:---------:|
| 1024   | 1.089    | 0.811    | 0.574    | 0.308     | 0.207     | 0.108     |
| 2048   | 1.076    | 1.161    | 0.753    | 0.477     | 0.334     | 0.176     |
| 4096   | 1.322    | 1.393    | 1.431    | 0.810     | 0.656     | 0.351     |
| 8192   | 1.158    | 1.482    | 1.453    | 1.204     | 1.067     | 0.620     |
| 16384  | 1.239    | 1.316    | 1.589    | 1.489     | 1.514     | 1.075     |
| 32768  | 1.246    | 1.460    | 1.434    | 1.649     | 1.563     | 1.569     |
| 65536  | 1.246    | 1.424    | 1.580    | 1.483     | 1.555     | 1.548     |
| 131072 | 1.268    | 1.442    | 1.503    | 1.676     | 1.533     | 1.514     |

Tensor parallelism size equal to 4:

| Tokens | 2 chunks | 4 chunks | 8 chunks | 16 chunks | 32 chunks | 64 chunks |
|-------:|:--------:|:--------:|:--------:|:---------:|:---------:|:---------:|
| 1024   | 0.892    | 0.579    | 0.374    | 0.195     | 0.104     | 0.060     |
| 2048   | 1.035    | 0.888    | 0.509    | 0.307     | 0.142     | 0.088     |
| 4096   | 1.156    | 1.081    | 0.795    | 0.466     | 0.245     | 0.134     |
| 8192   | 1.171    | 1.304    | 1.255    | 0.749     | 0.485     | 0.244     |
| 16384  | 1.162    | 1.309    | 1.416    | 1.216     | 0.780     | 0.496     |
| 32768  | 1.118    | 1.237    | 1.280    | 1.427     | 1.201     | 0.766     |
| 65536  | 1.218    | 1.310    | 1.386    | 1.528     | 1.553     | 1.195     |
| 131072 | 1.193    | 1.387    | 1.332    | 1.353     | 1.560     | 1.545     |

Tensor parallelism size equal to 8:

| Tokens | 2 chunks | 8 chunks | 16 chunks | 32 chunks | 64 chunks |
|-------:|:--------:|:--------:|:---------:|:---------:|:---------:|
| 1024   | 0.656    | 0.253    | 0.132     | 0.075     | 0.045     |
| 2048   | 0.828    | 0.324    | 0.183     | 0.087     | 0.051     |
| 4096   | 0.919    | 0.495    | 0.264     | 0.146     | 0.078     |
| 8192   | 0.993    | 0.705    | 0.470     | 0.240     | 0.157     |
| 16384  | 0.990    | 1.024    | 0.684     | 0.402     | 0.249     |
| 32768  | 0.972    | 1.118    | 0.942     | 0.760     | 0.469     |
| 65536  | 0.989    | 1.164    | 1.129     | 1.179     | 0.758     |
| 131072 | 1.018    | 1.241    | 1.277     | 1.297     | 1.090     |

### Performance Insights

Optimal configuration varies with tensor parallelism size and sequence length. There is no single setting that will benefit all benchmarks, so we recommend experimenting to find which configuration works best for your specific tensor parallelism size and sequence length. A good starting point is 2 chunks with a 4096-token threshold.

Diminishing returns with too many chunks. Excessively fine chunking introduces overhead from graph breaks, kernel launch latency, and reduced per-chunk compute efficiency. The optimal chunk count depends on both tensor parallelism size and typical token count.

## Recommended Settings

The following recommendations are based on isolated layer benchmarks using the Meta Llama 3.3 70B model. In end-to-end inference, the optimal configuration may differ depending on the model architecture, sequence lengths, and workload mix. We recommend benchmarking with your specific setup.

| Tensor parallelism size | Recommended chunks | Notes                                             |
| :---------------------: | :----------------: | ------------------------------------------------- |
|            1            |    1 (disabled)    | No all-reduce needed; chunking adds overhead only |
|            2            |        2-64        | Beneficial for token counts ≥ 4096                |
|            4            |        2-16        | Beneficial for token counts ≥ 8192                |
|            8            |        2-8         | Beneficial for token counts ≥ 16384               |

## Implementation Details

This feature is implemented in `vllm_gaudi/ops/hpu_row_parallel_linear.py` as `HPURowParallelLinear`, which registers as an out-of-tree (OOT) override for vLLM's `RowParallelLinear`. The chunking logic is entirely self-contained in the `forward` method and does not modify any other part of the model or the inference pipeline.

Each chunk boundary introduces a `torch._dynamo.graph_break()` to ensure correct async all-reduce semantics under `torch.compile`. This means the compiled graph will be split at chunk boundaries, which is a necessary trade-off for enabling async communication.