Reading TP-sharded fields crashes: reassemblers ignore shard_rank

[zaoxing]

What happens

Reading any tensor-parallel-sharded field from a TP>1 run crashes with a torch
error that gives no hint about the real cause:

RuntimeError: Boolean value of Tensor with more than one value is ambiguous

Why

All three row reassemblers in src/dmi/storage/internals.py group captured
rows by (layer_no, request_id) and sort them by start_token only. They
never look at shard_rank (key[4]).

Under TP each rank writes its own row for the same (layer, request, start_token). So the sort sees two entries whose first element is equal, falls
through to comparing the second element — a torch.Tensor — and torch refuses
to give a truth value for a multi-element tensor.

The pattern appears in three places:

• _reassemble_per_layer()
• _reassemble_attention_per_layer()
• _reassemble_global()

Reproducing it

Against main at 9ba0e8a, no GPU or database needed:

import torch
from dmi.storage.internals import get_internal

class Reader:
    def __init__(self, rows): self._rows = rows
    def prefix_get(self, prefix): return self._rows

for act, field in (("blocks.attn.hook_q", "q"),
                   ("blocks.attn.hook_pattern", "attentions"),
                   ("final_logits", "logits")):
    # two TP shards of the same layer / request / start token
    rows = [(("m", "0:0", act, 0, shard, 0, 2), torch.ones(2, 4))
            for shard in (0, 1)]
    getattr(get_internal("m", Reader(rows)), field)

All three raise.

Which fields are affected

HOOK_DEFS in native/csrc/ring/tensor_meta.h marks seven acts tp_sharded,
and resolve_shard_rank() in p2p_thread.cpp returns a non-zero rank only for
those. So these seven are reachable in a real TP run:

q, k, v, attention_values, mlp_activation, attention_scores,
attentions

Everything else always writes shard_rank = 0, so hidden_states, logits
and token_ids are not affected by sharding. _reassemble_global() shares the
defect but needs duplicate rows from somewhere else to trigger it — worth
hardening anyway, since ClickHouse inserts are not deduplicated and a retried
insert would produce exactly that.

Please don't fix it with just a tiebreaker

The obvious one-line fix — sort on (start_token, shard_rank) — stops the
crash but is worse than the crash, because the per-layer reassembler then
concatenates shards along dim 0, the token axis. That silently produces a
tensor of the wrong shape and meaning rather than an error.

Shards are not slices of the token axis, and not of the trailing axis either.
Derived from compute_hook_shape() (batch dim dropped, since rows store
per-request slices):

field	tp=1	tp=2	sharded axis
q, attention_values	[4, 8, 8]	[4, 4, 8]	1 (heads)
k, v	[4, 2, 8]	[4, 1, 8]	1 (heads)
attentions, attention_scores	[8, 4, 4]	[4, 4, 4]	0 (heads)
mlp_activation	[4, 128]	[4, 64]	1 (features)

So the correct merge axis differs per field, and is never the last one.

There is a further wrinkle for k and v. The shape is
max(1, num_kv_heads // tp), so once tp_size exceeds num_kv_heads the
clamp kicks in and every rank reports one head. With num_kv_heads=2:

tp=2: 1 head/rank x 2 ranks = 2   (partitioned)
tp=4: 1 head/rank x 4 ranks = 4   (replicated -- only 2 real heads)
tp=8: 1 head/rank x 8 ranks = 8   (replicated)

Under GQA with tp > num_kv_heads, shards are copies of each other, so
concatenating them would duplicate data. That case needs dedupe, not concat.

Suggested approach

1. Make the sort total so it can never compare tensors — sort on the key only,
   e.g. key=lambda item: item[0], or (start_token, shard_rank).
2. Raise a clear, named error when more than one shard_rank is present for a
   field and no merge rule is defined for it. That turns an opaque torch error
   into something actionable, and avoids inventing a merge that is wrong for
   most fields.
3. Implement per-field merges deliberately, using the axis table above, and
   handle the GQA replication case for k/v separately.

Step 1 alone is not enough; step 2 is what makes TP failures legible.

Side effect

Reassembly fails before validation runs, so shard-level completeness cannot
currently be checked either — a missing shard is indistinguishable from a
present one once this is fixed, unless the merge rule counts shards.

[zaoxing]

Does this issue by Claude make sense to you @Samfisheryu ?