Lower v2 SparseCore kernels through core_map#1
Closed
guowei-dev wants to merge 2 commits into
Closed
Conversation
DescriptionStart with a short description of what the PR does and how this is a change from The rest of the description includes relevant details and context, examples:
If the change fixes a Github issue, please include a link, e.g.,: TestsPlease describe how you tested this change, and include any instructions and/or ChecklistBefore submitting this PR, please make sure:
|
guowei-dev
force-pushed
the
fix/sc-kernel-coremap
branch
from
d3445f1 to
14b7dcd
Compare
guowei-dev
force-pushed
the
fix/v2-coremap
branch
from
02cae09 to
78d3fd9
Compare
New BlockSpec-pipelined gather (ragged_gather_v2) and 2D-grid combine (ragged_gather_reduce_v2) SparseCore kernels for MoE, wired into fused_moe_gmm and used by default. Set RAGGED_GATHER_VERSION=v1 and/or RAGGED_GATHER_REDUCE_VERSION=v1 to fall back to the legacy kernels independently. Includes kernel correctness/perf tests. Signed-off-by: guowei <guoweij@google.com>
Route the v2 ragged_gather and ragged_gather_reduce kernels through the core_map helper added for the v1 kernels, so they keep the fast single-mesh lowering instead of the slower mpmd_map path that pl.kernel uses on jaxlib 0.10.1. Signed-off-by: guowei-dev <rainycat.ava@gmail.com>
guowei-dev
force-pushed
the
fix/v2-coremap
branch
from
78d3fd9 to
327faf3
Compare
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
1 participant
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.
Summary
Adds the v2 SparseCore MoE kernels and routes them through the
core_map_helperintroduced in vllm-project#2887 (the v1 core_map routing fix, now in
main). This keeps the v2kernels on the fast single-mesh
core_maplowering instead of thempmd_mappath thatpl.kerneluses on jaxlib 0.10.1 —mpmd_mapaliases the kernel refs and serializes theSC gather against the gather_reduce preprocess argsort, losing the overlap. The helper is
already in
main, so this PR does not re-introduce it.Two commits:
ragged_gather_v2andragged_gather_reduce_v2,selected at import time via
RAGGED_GATHER_VERSION/RAGGED_GATHER_REDUCE_VERSION(both default to
v2, set either tov1to fall back).core_map_helper.E2E
Qwen3.5-397B-A17B-FP8, tpu7x-8, TP=8, EP + dp_attention, jax 0.10.1 / jaxlib 0.10.1 /
libtpu 0.0.41. Prefill-heavy profile (input 8192 / output 1024, 640 prompts, concurrency
512), 1 warmup + 3 runs averaged, rel stdev < 0.8%.
Same base, same bench, only the lowering differs — the fix moves the v2 kernels the same
way it moved the v1 kernels (+8.25%, vllm-project#2887).