Repository navigation
spec_decode_verify transforms and qmv_quad gs 128 coverage (gemma4-qat#323, #324) - #13
Merged
Merged
Conversation
… qmv_quad gs 128 The composition gathers the committed token with n_acc as its index and no stop_gradient, so jvp/vjp/grad through spec_decode_verify threw for K > 0. The outputs are integer ids, so the index now carries a zero tangent (numen-tech/gemma4-qat#323). test_qmv_quad_ragged_n adds group size 128 at K = 128, which dispatch_qmv also routes to qmv_quad (numen-tech/gemma4-qat#324): 72 subtests, all pass.
Codex Review SummaryThis comment shows the latest Codex review activity on this pull request.
ℹ️ About Codex in GitHubYour team has set up Codex to review pull requests in this repo. Reviews are triggered when you
Codex reacts with 👀 while any review is running, comments if it has suggestions, and reacts with 👍 once all reviews finish with no findings. |
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.
Two follow-ups to #12 on
prism-v0.32.4-fixes, tracked in gemma4-qat (this repository has issues turned off).spec_decode_verify: transforms throw for K > 0 (numen-tech/gemma4-qat#323)
In the composition, the accept length
n_accis the index of atake_along_axisand carries nostop_gradient. Somx.jvp,mx.vjpandmx.gradthroughmx.fast.spec_decode_verifyall throw[gather_axis] Cannot calculate JVP/VJP with respect to indices. The outputs are integer token ids, so the right tangent is zero. The index now goes throughstop_gradient, which changes no forward value.test_spec_decode_verify_transformschecks that jvp gives zero tangents and that vjp gives a zero cotangent with the logits' shape.qmv_quad ragged-N test: group size 128 (numen-tech/gemma4-qat#324)
dispatch_qmvsends K = 128 toqmv_quadfor any group size, buttest_qmv_quad_ragged_n(patch 23, numen-tech/gemma4-qat#302) covered only group sizes 32 and 64. Group size 128 at K = 128 is now included: 72 subtests, up from 56. They pass on the current kernel, so this adds coverage only.Test
pytest python/tests: 966 passed, 57 skipped (965 before, plus the new test).build-lib/tests/tests(metallib): 311/311 test cases, 6534 assertions.