Repository navigation
Collectives: failure-only pack diagnostics for rejected PCIe all-reduce plans + complete draft-block channel coverage - #21
Open
tunamitom wants to merge 2 commits into
Conversation
…stics When runtime.plan(...) rejects a PCIe all-reduce candidate the raised ValueError names the requirement but not the buffers' actual pack and alignment, which is what every rejection investigation needs. Wrap the call to re-raise ValueError with source/output shape, stride, dtype, byte counts, data_ptr%16, bytes%16 and contiguity beside the runtime's byte capacity, preserving the exception chain (from error). The path is failure-only: zero steady-state cost. This change exists because these diagnostics are how the decode all-reduce channel sizing defect fixed in the following commit was found: the re-raised message showed a 32-row output-projection partial rejected against a 20-row channel. Validation: tests/components/test_speculative_capacity.py on this tree at parity with the release-lane base (CPU only; failure-path code). Not run in this pass: a real runtime rejection against FlashInfer. Material limitation: no test exercises the rejection message end to end. Design-Signed-Off-By: Hermes Agent (hermes-agent); design-review=a38f
Draft graphs reduce batch * (draft_tokens + 1) rows independently of the packed target capacity, so lowering the decode token capacity (32 -> 20) shrank the shared oneshot channel (sized from decode_tokens) below the batch-4 draft block: the 32-row output-projection partial (458,752 B at hidden 7168) exceeded the 20-row channel (286,720 B) and every draft-round plan rejected with 'input does not satisfy PCIe oneshot requirements' (the capacity-32 build passed only because decode_tokens coincided with the 32-row block). Size the channel from all_reduce_rows = max(decode_tokens, draft_reduce_rows), the same draft-block bound the other drafting recipes use for this channel. The cross-recipe check was done: minimax_m3, step5, qwen38, mimo_v2 and deepseek41 already size this channel from draft rows; Kimi was the only outlier. At the default capacity (32 = 4 x 8) the channel size is unchanged, so default packages are identical. Validation: tests/components/test_speculative_capacity.py on this tree (CPU; includes the new all_reduce_rows coverage). Not run in this pass: a CUDA preparation of an undersized-capacity package. Material limitation: none — buffer sizing only. Design-Signed-Off-By: Hermes Agent (hermes-agent); design-review=a38f
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
What / why
Two small collective-planner fixes:
all_reduce_rows = max(decode_tokens, draft_reduce_rows), the same draft-block bound the other drafting recipes use for this channel.Cross-recipe check: the equivalent paths in all drafting recipes were checked — minimax_m3/bootstrap.py:222, step5/bootstrap.py:177, qwen38/bootstrap.py:362, mimo_v2/bootstrap.py:204-206 and deepseek41's
all_reduce_tokensalready size this channel from draft rows; Kimi was the only outlier, so this is not a new pattern. At the default capacity (32 = 4 x 8) the channel size is unchanged and default packages are identical.Commits:
Annotate rejected PCIe all-reduce plans with failure-only pack diagnostics+Cover complete draft proposal blocks in the decode all-reduce channel.Numerics class
NUMERICS-NEUTRAL — channel sizing and failure-path diagnostics; no arithmetic change.
Tests
tests/components/test_speculative_capacity.py— 141 passed (this tree, includes the newall_reduce_rowscoverage).tests/recipes/test_kimi31.py25 andtests/cli/test_models.py40 passed;tests/cli641 passed, same 2 pre-existing failures as the base.Possible follow-up (left out to keep this PR minimal):
getattr(runtime, 'max_size', None)hedges on a known FlashInfer attribute, and no test exercises the rejection message end to end (a CPU test with an undersized fake runtime is cheap).Dependencies: none on the companion FlashInfer PRs.