perf(comm): route generic MNNVL AR patterns to upstream flashinfer - #1075
Draft
dongjiyingdjy wants to merge 1 commit into
Draft
perf(comm): route generic MNNVL AR patterns to upstream flashinfer#1075dongjiyingdjy wants to merge 1 commit into
dongjiyingdjy wants to merge 1 commit into
Conversation
Replace the private MNNVL kernel with flashinfer's trtllm_mnnvl_ar for the patterns upstream implements (kAllReduce, kARResidualRMSNorm), keeping the private kernel only for the K3-specific epilogues it alone provides (kARResidualAttnResCombine, kAllReduceLatentNorm). A wrapper class adapts flashinfer's MNNVLAllReduceFusionWorkspace to the local supports()/dispatch contract; imports are sentinel-guarded so installs without the module keep the private path (flashinfer 0.6.16, the current pin, has the full API). Workspace creation follows the existing two-phase MIN-vote pattern, and one multicast allocation per process group is shared across every arming site through a group-keyed cache: the McastGPUBuffer granularity (512 MB) dwarfs any requested size, so per-site copies would have cost 2 GB per rank for no benefit. Routing splits by token count (MNNVL_FLASHINFER_MIN_TOKENS, default 32, env-overridable). Isolated graph-replay favors flashinfer at every M, but in-situ (nsys over steady decode, 8x B300 TP8) the two kernels are wait-inclusive equals at decode M while flashinfer's multicast traffic slows overlapped neighbor kernels (~+75 us GPU / +30 us wall per step), so decode-sized calls stay on the private kernel. From M=32 up to the IPC crossover flashinfer wins decisively (1.8-3.2x at M=64-128 for both hidden widths), and cross-node it takes the whole range (1.5-2.1x over the private two-shot at every measured M). Each workspace serves as the other's shape fallback. Measured end to end (K3 TP8 decode, differential TPOT, 9 trials): 10.794 ms vs 10.779 ms same-day baseline (inside the noise band; a full swap without the threshold measured +0.27%). 8-rank numeric checks pass at bf16 rounding level across M=1..2048 for both swapped patterns; the private kernel suite (20 tests x 8 ranks) is unaffected. The private .cuh cannot be removed yet: the two K3 epilogues have no upstream equivalent. Upstreaming them (or a prefix-norm epilogue interface) would allow full retirement. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Signed-off-by: dongjiyingdjy <87510204+dongjiyingdjy@users.noreply.github.com>
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
Replaces the private MNNVL AR kernel with flashinfer upstream
trtllm_mnnvl_arfor the two generic fusion patterns (kAllReduce,kARResidualRMSNorm), keeping the private kernel only for the K3-specific epilogues upstream does not implement (kARResidualAttnResCombine,kAllReduceLatentNorm).Routing (single node, hidden 7168, per-pattern)
Threshold is
MNNVL_FLASHINFER_MIN_TOKENS(default 32, env-overridable). Each mnnvl workspace is the other's shape fallback.Why not a full swap
Isolated graph-replay favors flashinfer at every M, but in-situ (nsys over steady decode, 8x B300 TP8) the two are wait-inclusive equals at decode M, and flashinfer's multicast traffic slows overlapped neighbor kernels (~+75 us GPU / +30 us wall per step → +0.27% TPOT). With the threshold, decode runs the exact same kernel stream as before.
Memory
One multicast allocation per process group, shared across all arming sites via a group-keyed cache (McastGPUBuffer granularity is 512 MB regardless of requested size; per-site copies measured 2 GB/rank).
supports()judges capacity by actual buffer bytes.Validation
Not in this PR
The private
.cuhstays: the two K3 epilogues have no upstream home. Upstreaming them (or a prefix-norm/custom-epilogue interface) would allow full retirement. Cross-node numbers are from single-node standalone benchmarks; true multi-node validation pending.🤖 Generated with Claude Code