big M opt: splitT1 + L2 scatter + 2 stream - #5
Merged
Merged
Conversation
…L2 scatter + 2 stream) Kernel-only. Base: megamoe_nvfp4_dev @ c277f44. One weight format, one entry point. Every forward picks the kernel family from the routed load actually seen by this rank: rho = M * topk / local_experts rho <= 192 -> BN256 family, small-M fused path rho > 192 -> BN128 family, split L1 + split L2 Selection happens in the C++ entry from this call's y.size(0), so decode, small prefill and large prefill can take different families inside one serving process; model name and hidden size no longer participate. Both families read the same packed-B bytes — block_n selects only the retained scale-metadata view. kernel_family="auto"|"fused"|"split" is exposed for cross-validation only; production must use "auto". Large-M split always enables all three optimizations together, emitted as compile-time defines into the JIT key so no environment variable can silently select an unvalidated variant: 1. splitT1 — compile specialization of the NVLink barrier timeout path: DG_NVLINK_BARRIER_TRAP_ONLY_TIMEOUT is now emitted for the split JIT source, replacing the 7-argument device printf in the barrier spin-wait timeout branch with a bare trap. dev already did this for fused; split did not; 2. 2 stream — L1 dispatch-dequant half-stream mode4 + physical-warp remap (DG_NVFP4_L1_DEQUANT_HALF_STREAM=4, DG_NVFP4_L1_DEQUANT_WARP_REMAP=1, with a static assert binding remap to mode4); 3. L2 scatter — split L2 (down) warp-private contiguous 8-row shared-memory staging, so each destination row leaves as one 256-byte burst instead of eight scattered 16-byte requests billed as 32-byte sectors. Large-M result versus dev (Flash and Pro, M=2048/4096/8192): faster at all 6 points, +7.54% to +13.61%, equal-weight geometric mean -10.35% latency. Still slower than the W8A8 calibration reference at all 6 points, geometric mean +14.19%. The measurement protocol behind those numbers was not supplied and is not reproducible from this repository; README_BIGMOPT.md records that gap explicitly. No fused-path measurements are included. Small M keeps dev's fused device body byte-identical (sm90_nvfp4_mega_moe_small_m_fused_body.inl is unchanged), but the fused path is not bit-for-bit dev behaviour — three things differ: - heuristics/sm90_nvfp4_mega_moe_small_m.hpp: the rho<=6 bucket flips use_mode2_lop3_decoder false -> true, because one common Mode2 braided packed-B copy now serves both families; - impls/sm90_nvfp4_mega_moe_small_m.hpp: the fused launch passes launch_pdl=false; - impls/sm90_nvfp4_mega_moe_mode2_dequant.cuh is included by both the split and small-M .cuh, so its changes are on the small-M compile path too. The other five small-M tuning buckets are unchanged. Reverting the first and third would restore stock dev small-M behaviour but would break the single-weight-copy premise. On the weight side, choose_nvfp4_block_n_for_mega_moe_sm90 drops the intermediate_hidden>=3072 special case (cutoff 190) and always uses family_threshold=192. The SGLang integration overlay, built wheels and raw gate evidence are not part of this branch; README_BIGMOPT.md records where the full sealed delivery is retained.
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
2 participants
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
Base
megamoe_nvfp4_dev@c277f44. Three large-M split optimizations, plus a unified entry that keeps one weight copy and picks the kernel family per forward. Kernel only, no SGLang integration. 16 files, +1628 / -183.Dynamic family selection
Decided in the C++ entry from this call's
y.size(0), so decode, small prefill and large prefill can take different families in one process. Model name and hidden size no longer participate. Both families read the same packed-B bytes;block_nonly selects the retained scale-metadata view.kernel_family="auto"|"fused"|"split"— production uses"auto".What changed — large M (split)
Always enabled together, emitted as compile-time defines into the JIT key.
DG_NVLINK_BARRIER_TRAP_ONLY_TIMEOUTfor the split JIT source, replacing the 7-arg deviceprintfin the NVLink barrier timeout branch with a bare trap. dev already did this for fused; split did not.barrier.cuhunchangedsplit_l1_body.inl983 -> 1679 linessplit_l2_body.inl659 -> 751 linesWhat changed — small M (fused)
Fused device body is byte-identical to dev. The dequant functions it calls (
dequant_braided_quad*,dequant_braided_selector_word,dequant_smem_b_from_packed_braided_lut_window) are unchanged — the +295 lines inmode2_dequant.cuhare split-only functions fused never calls.Two host lines differ:
heuristics/..._small_m.hpp:130use_mode2_lop3_decoderfalse -> true, only in therho <= 6bucketimpls/..._small_m.hpp:287launch_pdl=false(dev defaults totrue)choose_nvfp4_block_n_for_mega_moe_sm90drops theintermediate_hidden >= 3072special case (cutoff 190) and always uses 192. Pro's crossover moves M=1520 -> 1536; Flash and MiMo unaffected.Performance
Microseconds, 8-GPU rank-MAX median.
dev=c277f44,W8A8=deepseek-ai/DeepGEMMPR deepseek-ai#383 @bc4f33a.Large M — split (the bigMopt work)
bigMopt vs devpositive = bigMopt faster. Geometric mean vs dev: Flash -12.56%, Pro -8.09%, all 6 points -10.35% (1.1155x). Faster than dev at 6/6, still slower than W8A8 at 6/6 (geometric mean +14.19%).Small M — fused
bigMopt runs dev's fused body byte-for-byte, so this is a precision comparison, not a bigMopt speedup claim. Negative = NVFP4 faster.
Geometric mean: M<=256 (12 points) NVFP4 faster by 11.55%; M>=512 (4 points) W8A8 faster by 20.84%. Crossover between M=256 and M=512. PR deepseek-ai#383's FP8 path is always split L1+L2 while NVFP4 here is fused, so these rows compare two kernel structures.
Protocol
Small-M table:
bench_kineto, 8 GB L2 flush,num_tests=20internal mean is one observation, median of 50 observations for M<=128 and 20 for M>=256, same rule for every arm. Large-M table: the protocol behind thebigMoptcolumn was not supplied and raw logs are not archived here; itsdevandW8A8columns agree with the small-M campaign to within 1.2%.Full detail in
README_BIGMOPT.md.