Part of #83 · refines #102
Make TP=1 and TP>1 produce identical reduction order for the WS1 ops (selected-logprob, masked reductions). The reduction contract and the concrete drift example live in #102. Tolerance reuses the #108 per-dtype threshold table — do not define a separate one.
Relationship to #102 (done vs remaining): #102 is the design/contract + reference issue — it defines the TP-invariant reduction semantics and provides a reference helper on simulated TP shards. Its Non-Goals explicitly exclude real vLLM/Megatron/FSDP integration, and it is only forward/contract-focused. #102 is still open (only the reference scaffold PR #103 is in progress), so the contract is not yet fully frozen. This issue (#109)
owns the parts #102 does not cover: the actual tree-reduce kernel implementation, the TP-wrapper including the backward all-reduce path, real multi-GPU consistency tests (beyond simulated shards), and the performance benchmark. In short: #102 defines the correct answer; #109 builds the working multi-GPU implementation and tests against it. Before starting, confirm with @inaniloquentee which parts of #102's contract are frozen.
Planned PRs:
Depends on: WS1 logprob op working draft. The wrapper/test design can start in W1; full work begins once the WS1 op is available.
Part of #83 · refines #102
Make TP=1 and TP>1 produce identical reduction order for the WS1 ops (selected-logprob, masked reductions). The reduction contract and the concrete drift example live in #102. Tolerance reuses the #108 per-dtype threshold table — do not define a separate one.
Relationship to #102 (done vs remaining): #102 is the design/contract + reference issue — it defines the TP-invariant reduction semantics and provides a reference helper on simulated TP shards. Its Non-Goals explicitly exclude real vLLM/Megatron/FSDP integration, and it is only forward/contract-focused. #102 is still open (only the reference scaffold PR #103 is in progress), so the contract is not yet fully frozen. This issue (#109)
owns the parts #102 does not cover: the actual tree-reduce kernel implementation, the TP-wrapper including the backward all-reduce path, real multi-GPU consistency tests (beyond simulated shards), and the performance benchmark. In short: #102 defines the correct answer; #109 builds the working multi-GPU implementation and tests against it. Before starting, confirm with @inaniloquentee which parts of #102's contract are frozen.
Planned PRs:
Depends on: WS1 logprob op working draft. The wrapper/test design can start in W1; full work begins once the WS1 op is available.