vibe-coded BMM sharding - #1
Closed
suyoggupta wants to merge 1 commit into
Closed
Conversation
Signed-off-by: Suyog Gupta <suyogg@nvidia.com>
nvkgoyal
reviewed
May 9, 2025
| ad_logger.debug("After sharding: " + str(gm)) | ||
| return gm | ||
|
|
||
| def bmm_shard(gm: GraphModule, rank: int, world_size: int) -> GraphModule: |
|
see NVIDIA#4163 |
lucaslie
pushed a commit
that referenced
this pull request
Apr 21, 2026
Three issues fixed: 1. past_key_value (singular) renamed to past_key_values (plural) for LlamaAttention and GPT2Attention in transformers 5.x. 2. use_cache parameter removed from attention forward calls in transformers 5.x (cache is always updated in-place). 3. rope_theta/rope_scaling config attributes replaced with unified rope_parameters dict (same fix as test_gpt_attention.py). Without fix #1-2, the DynamicCache was never populated because the kwarg was silently ignored, causing an empty tuple from dynamic_cache_to_legacy and an IndexError at line 1026. Signed-off-by: Jonas Li <6110159+longlee0622@users.noreply.github.com>
nvchenghaoz
added a commit
that referenced
this pull request
Jun 30, 2026
…llm-gen W4A16 runner (idea idea_0023) On SM100, the DSV4 routed MXFP4 MLP (torch_mxfp4_moe_from_routing(_ep), up_gate / deepseek-SwiGLU) now runs on the trtllm-gen W4A16 bf16-act runner bf16_mxe2m1_block_scale_moe_runner instead of the fp32 dequant+bmm reference (the #1 decode op post-idea_0002). Weights are re-interleaved ([up|gate] split-half -> gpt-oss interleaved gate@0::2/up@1::2 -- the cos 0.012->0.977 fix proven in idea_0008) and pad/shard/shuffled once via prepare_trtllm_gen_moe_mxfp4_weights, cached by raw-weight identity so the prep runs pre-CUDA-graph-capture and captured decode steps reuse it (mirrors _prepare_weights_scales_cached). Precomputed routing is forwarded as topk_ids/topk_weights (used verbatim; routed_scaling_factor already folded upstream); off-rank routes are tagged with the invalid sentinel expert_id == local_experts so the kernel SKIPS them -- clamping them onto a valid expert floods its histogram and drops real routes (the masking fix beyond idea_0008's all-local microbench). The sharding transform's existing all_reduce sums the per-rank local partials. Gated to is_sm_100f() + up_gate + deepseek (DSV4 only); gpt-oss and non-SM100 keep the torch reference. DeepSeek-V4-Flash proxy(10L) on B200, EP=8, decode window 40-50: gather_scatter share: pre 27.71% -> post 7.34% (instances 6740->5940; -20.37pp, 20x noise floor 1.0%) dequant swarm (gather_scatter+copy_cast) abs GPU time: 232633us -> 88131us (-62.1%) total decode GPU time: 382152us -> 228124us (-40.31%; -46.47% excluding NCCL-run-to-run comm variance). copy_cast share is flat (33.17->31.29%) only because the denominator shrank 40%; its absolute time dropped -55362us. Cross-checked absolute total_us (the matmul_ogs false-positive guard, idea_0005): a genuine GPU-time collapse, not denominator inflation. Full-model tpot confirmation deferred to Tier-2 batched validation. Accuracy: numdiff = justified pure-perf rewrite (proxy uses random init -> greedy outputs differ for both arms; byte-equality is moot). Unit test test_mxfp4_moe_trtllm_gen_from_routing: cos~1.0 (masked, nonzero-bias) and 0.92-0.97 (all-local) vs the torch reference at DSV4 shapes. The kernel SwiGLU clamp(gate,max=L)*sigmoid(a*clamp(gate))*clamp(up,+/-L) is bit-equivalent to the reference's deepseek SwiGLU; W4A16 (mxfp4 weight + bf16 act) is the model's native precision -- the same trtllm-gen kernel gpt-oss uses in production -- so this is faithful, not a degradation. PASSED (unit cos; Tier-2 GSM8K is the final gate). Signed-off-by: Chenghao Zhang <211069071+nvchenghaoz@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
3 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.
Simple BMM sharding with all_gathers