ward-clustering

Batched Ward hierarchical clustering on CUDA. One call returns a SciPy-compatible linkage tree, flat cluster labels, and the actual number of clusters per document. Supports variable-length batches and CUDA graph capture.

Usage

Install kernels in an environment with CUDA-enabled PyTorch:

pip install kernels

import torch
from kernels import get_kernel

kernel = get_kernel("sentence-transformers/ward-clustering", version=1, trust_remote_code=True)

x = torch.randn(32, 256, 128, device="cuda")
distances = torch.cdist(x, x)
linkage, labels, counts = kernel.ward(distances, num_clusters=128)

ward accepts square CUDA dissimilarity matrices of shape (N, N) or (B, N, N). Only the strict upper triangle is used. Inputs must be float32 or float64, with finite nonnegative distances whose squares are representable in that dtype.

Single-matrix outputs omit the batch dimension. Outputs remain on CUDA.

For padded batches, provide valid prefix lengths and per-document cluster targets together as CUDA int64 tensors:

lengths = torch.tensor([256, 193, 81, 0], device="cuda")
targets = (lengths // 2).clamp_min(1).minimum(lengths)
linkage, labels, counts = kernel.ward(
    distances[:4], lengths=lengths, cluster_counts=targets
)

Lengths must be between 0 and N. Targets override num_clusters and must be between 1 and the valid length, or 0 for an empty document. Unused output rows are zero-padded. Pass validate=True to check input values, which synchronizes with the host.

Behavior and limits

• Supports 1 to 4096 tokens per padded document and up to 65535 documents per batch, subject to GPU memory. Distance workspace grows quadratically with padded length.
• Tied merges can produce fewer clusters than requested. Requesting one cluster per token returns singletons.
• Float32 rounding and tied distances can produce different trees or partitions from SciPy. Use float64 for tighter reference agreement.
• Euclidean distances give the usual Ward objective. Cosine dissimilarities are also accepted, as used in token pooling, but define a different objective.
• Clustering outputs have no gradient. Distance construction and differentiable cluster means belong to the caller.

Benchmarks

Latency is measured in milliseconds per batch. Lower is better. b128_n512 means 128 documents of 512 tokens, each clustered to N/2 groups. Each CUDA call computes linkage and flat clusters from precomputed float32 cosine distances. Encoding, distance construction, and pooled means are excluded.

Measured on an RTX 3090 with PyTorch 2.10.0+cu128, one CPU thread, TF32 disabled, and asynchronous CUDA launches. CUDA timings average 100 synchronized samples after ten warmups. The SciPy reference uses resident CPU inputs and one timed pass, including the return of merge heights to CUDA for verification.

The standard benchmark entrypoint is benchmarks/benchmark.py, using kernels.benchmark.Benchmark. Direct Hub execution currently needs a CLI update to support native kernel repositories.

Throughput (batch calls per second) and an animated comparison are also available.

Hierarchical Token Pooling in Sentence Transformers

Use the released Sentence Transformers encoding and scoring APIs to pool document embeddings after inference. This example encodes one query and four documents with mLateOn, clusters all four documents in one CUDA call, and averages the vectors in each cluster using PyTorch.

pip install "sentence-transformers>=6.0.1" kernels

import torch
from kernels import get_kernel
from sentence_transformers import MultiVectorEncoder
from torch.nn.utils.rnn import pad_sequence

model = MultiVectorEncoder("lightonai/mLateOn", device="cuda")
kernel = get_kernel("sentence-transformers/ward-clustering", version=1, trust_remote_code=True)

query = "How do heat pumps keep a house warm in winter?"
documents = [
    """A heat pump heats a house by moving thermal energy from the outside air
    into the building. Even on a cold winter day, outdoor air contains energy
    that a refrigerant can absorb. A compressor raises the refrigerant's pressure
    and temperature, and a heat exchanger transfers that warmth to indoor air
    or water circulating through radiators. The refrigerant then expands and
    repeats the cycle. Unlike a gas boiler, the system does not burn fuel to
    create heat. Its efficiency depends on the outdoor temperature, the required
    heating temperature, and the design of the installation. Cold weather can
    require defrost cycles or additional heating capacity.""",
    """Improving insulation helps a house stay comfortable throughout winter.
    Heat can escape through the roof, walls, windows, and gaps around doors.
    Adding loft insulation and sealing draughts reduces these losses, while
    double or triple glazing limits heat transfer through windows. Ventilation
    still matters because an airtight building needs a reliable supply of fresh
    air to manage moisture. Before replacing a heating system, an energy survey
    can identify which improvements will have the largest effect. Better
    insulation usually lowers the heating demand and helps rooms maintain a
    more consistent temperature. It can also make lower-temperature radiators
    and underfloor heating more practical.""",
    """Rooftop solar panels convert sunlight into electricity using photovoltaic
    cells. An inverter converts the panels' direct current into alternating
    current that household appliances can use. Output changes with the time of
    day, cloud cover, shading, and the orientation of the roof. A battery can
    store some daytime generation for use after sunset, although its capacity
    limits how much energy can be shifted. Households may export surplus power
    to the grid under their electricity tariff. Planning an installation means
    checking the roof condition, estimating annual generation, and comparing
    that estimate with the home's electricity consumption. Panels generate
    electricity rather than directly warming the rooms.""",
    """A public library is digitizing its collection of historical newspapers
    so that readers can search local reporting from previous centuries. Staff
    first inspect fragile pages and photograph them using a scanner designed
    to protect the paper. Optical character recognition turns the images into
    searchable text, but volunteers must correct errors caused by faded ink
    and unusual typefaces. Each issue receives metadata describing its date,
    publisher, and place of publication. The library keeps preservation copies
    alongside smaller files for its website. Researchers can then trace changes
    in local businesses, transport, and community life without repeatedly
    handling the original newspapers in the reading room.""",
]

query_embeddings = model.encode_query([query], normalize_embeddings=True)
document_embeddings = model.encode_document(documents, normalize_embeddings=True)


def report(name, embeddings):
    sizes = [len(embedding) for embedding in embeddings]
    scores = model.similarity(query_embeddings, embeddings)[0].tolist()
    print(f"{name}: {sum(sizes)} document token embeddings")
    for i, (size, score) in enumerate(zip(sizes, scores), start=1):
        print(f"  Document {i}: {size:3d} embeddings, similarity {score:.4f}")


report("Before pooling", document_embeddings)

pool_factor = 4
lengths = torch.tensor([len(x) for x in document_embeddings], device="cuda")
padded = pad_sequence(document_embeddings, batch_first=True).float()
distances = (1 - padded @ padded.transpose(1, 2)).clamp(0, 2)
targets = (lengths // pool_factor).clamp_min(1)
_, labels, counts = kernel.ward(distances, lengths=lengths, cluster_counts=targets)

pooled_embeddings = []
for embedding, doc_labels, count in zip(document_embeddings, labels, counts.tolist()):
    doc_labels = doc_labels[: len(embedding)]
    sums = torch.zeros(count, embedding.shape[1], device="cuda", dtype=torch.float32)
    sums.index_add_(0, doc_labels, embedding.float())
    sizes = torch.bincount(doc_labels, minlength=count).unsqueeze(1)
    pooled_embeddings.append((sums / sizes).to(embedding.dtype))

report("After pooling", pooled_embeddings)

Example output:

Before pooling: 618 document token embeddings
  Document 1: 152 embeddings, similarity 12.9408
  Document 2: 154 embeddings, similarity 12.7925
  Document 3: 158 embeddings, similarity 12.6670
  Document 4: 154 embeddings, similarity 12.4875
After pooling: 153 document token embeddings
  Document 1:  38 embeddings, similarity 12.7861
  Document 2:  38 embeddings, similarity 12.6806
  Document 3:  39 embeddings, similarity 12.5529
  Document 4:  38 embeddings, similarity 12.4039

The query embeddings stay unchanged. Cluster means are not renormalized, and pooling can change scores and rankings. Here, pool_factor=4 targets roughly a fourfold reduction in document vectors.