Deduplicate the GLM-5.2 and DSV4 exact LM heads onto shared machinery

[qywu]

Stacked on #74 (base branch issue-71-consolidate-scoring-paths; retarget to main after #74 merges). Second consolidation tranche for #71: the GLM-5.2 (glm5/exact_lm_head_qlora.py) and DSV4 (deepseek_v4/exact_lm_head.py) exact heads ran the same program skeleton in near-duplicate — this PR extracts it once and leaves each family holding only its contract.

Design rule: share the machinery, inject the contract

New xorl/models/transformers/exact_lm_head_shared.py owns the skeleton. The families keep the byte-bearing pieces as closures and constants:

shared (one copy)	family-owned (contract)
ExactLmHeadFunction — the one autograd boundary (gather rows → exact value → local slice; backward: gather grads → surrogate VJP → local slice). Replaces 3 per-family Functions.	base projection kernel (head_v2_full_logits_with_lse FP32 vs F.linear BF16)
ExactHeadRowPlan — row ownership as an injected plan: replicated (GLM local head), equal rank-order blocks (GLM TP16), ragged padded blocks (DSV4 TP8)	temperature dtype boundary (FP32 scale vs BF16 divide/store)
filtered_surrogate_local_grad_logits / surrogate_local_grad_logits — the chunked support+reference-grad loops, parameterized by exact-score/reference-logits closures	native selected-score kernel (head_v2_selected_logprob_from_logits vs BI log_softmax)
exact_lora_local_logits + single_adapter_lora_batch_info — the literal serving A/B SGEMM choreography and its invariants	surrogate VJP formulation (_local_qlora_surrogate_vjp vs _local_surrogate_vjp — different validated dtype boundaries, deliberately NOT merged)
rank-order collectives: vocab gather (dtype/world-parameterized), equal + ragged row gathers, row counts, fp32 sum reduce	geometry constants and operand validation (TP16/9,680-row vs TP8/16,160-row shards)
check_exact_head_tp_group — one TP-group geometry checker with per-family program names	which checks apply (DSV4 adds source_ordinal)

Both filtered forwards now also score through the shared score_with_sampling_transforms instead of hand-rolled support/identity/partitioned sequences, and gather_vocab_shards moved from the loss module to xorl.utils.dist_utils (the placement finding from #74's review).

Byte-safety

Forward value paths are byte-preserved: the same kernels run in the same order with the same dtype boundaries — only plumbing moved. The pieces where the two families genuinely differ in validated numerics (the straight-through surrogate formulations, the base GEMMs, the temperature stores) were deliberately left family-side rather than force-unified. Test-pinned error substrings (requires TP16, gather order, shard/group rank mismatch, must use NCCL, equal source-row counts, collective rank order) survive via the program parameter.

Numbers

Net −575 lines (glm5/exact_lm_head_qlora.py 1298→847, deepseek_v4/exact_lm_head.py 836→612, +369-line shared module). Combined with #74, the exact-head files are down ~40% while gaining the filter capability.

Testing

• Existing coverage carries over: the row plans and shared Function run under real 2-rank gloo collectives (test_glm52_exact_lm_head_qlora_collectives.py), the ragged plan + monkeypatched collectives (test_dsv4_exact_lm_head_temperature.py), the replicated plan + Hopper byte test (test_glm52_exact_lm_head_qlora.py), and the TP-group validation matrix. Test collection counts are identical before/after.
• Full affected sweep green locally: 78 CPU + 15 GPU/distributed, plus test_lm_head_tp_fsdp_e2e.py, test_lm_head_tp_parallel_state.py, test_glm52_exact_qlora_fsdp2.py (17 tests, 2×H100, ~8.5 min).
• The shared module joined the no-dense-support source-inspection test.

[broly-code-security-scanner]

Broly Security Scan

Note

✅ Clean scan
No vulnerabilities detected in this PR.

Note

Re-scan this PR anytime with /broly scan — useful after /broly undismiss, or to refresh findings without a new push.

Broly — SAST (zai-org/GLM-5.2) · Secrets · SCA · IaC · GH Actions · Base Images · Supply Chain Threats · Exploit Chains · Adversarial Verification

We're continuously improving Broly's accuracy and finding quality — your feedback is valuable. False positives, missed findings, bugs, and feature requests all welcome.