Summary
collector/sglang/collect_dsv4_attn.py::_resolve_local_heads computes the persisted rank-local num_heads with integer division.
Required change
Before calculating native_heads // tp_size, validate that native_heads % tp_size == 0. Raise RuntimeError when the geometry is not exactly divisible.
Rationale
A non-divisible geometry truncates the local head count. Consumers derive the native identity as num_heads * tp_size, so the persisted row can move into an incorrect native-head bucket.
Affected area
collector/sglang/collect_dsv4_attn.py_resolve_local_heads
Acceptance criteria
_resolve_local_heads raises RuntimeError for non-divisible native_heads and tp_size._resolve_local_heads returns the exact quotient for divisible geometry.- The existing
module_heads validation remains in place.
- Add or update unit coverage for non-divisible geometry.
Backlinks
Summary
collector/sglang/collect_dsv4_attn.py::_resolve_local_headscomputes the persisted rank-localnum_headswith integer division.Required change
Before calculating
native_heads // tp_size, validate thatnative_heads % tp_size == 0. RaiseRuntimeErrorwhen the geometry is not exactly divisible.Rationale
A non-divisible geometry truncates the local head count. Consumers derive the native identity as
num_heads * tp_size, so the persisted row can move into an incorrect native-head bucket.Affected area
collector/sglang/collect_dsv4_attn.py_resolve_local_headsAcceptance criteria
_resolve_local_headsraisesRuntimeErrorfor non-divisiblenative_headsandtp_size._resolve_local_headsreturns the exact quotient for divisible geometry.module_headsvalidation remains in place.Backlinks