Skip to content

[None][fix] Fix fused MHC for DeepSeek-V4-Pro hidden size#13771

Merged
pcastonguay merged 3 commits intoNVIDIA:feat/deepseek_v4from
pcastonguay:fix/dsv4-pro-fused-mhc-hidden-7168
May 5, 2026
Merged

[None][fix] Fix fused MHC for DeepSeek-V4-Pro hidden size#13771
pcastonguay merged 3 commits intoNVIDIA:feat/deepseek_v4from
pcastonguay:fix/dsv4-pro-fused-mhc-hidden-7168

Conversation

@pcastonguay
Copy link
Copy Markdown
Collaborator

@pcastonguay pcastonguay commented May 5, 2026 •
edited
Loading

@coderabbitai summary

Description

Same as #13710 but fixes pre-commit check. Original author: @Oseltamivir

Test Coverage

PR Checklist

Please review the following before submitting your PR:

  • PR description clearly explains what and why. If using CodeRabbit's summary, please make sure it makes sense.

  • PR Follows TRT-LLM CODING GUIDELINES to the best of your knowledge.

  • Test cases are provided for new code paths (see test instructions)

  • Any new dependencies have been scanned for license and vulnerabilities

  • CODEOWNERS updated if ownership changes

  • Documentation updated as needed

  • Update tava architecture diagram if there is a significant design change in PR.

  • The reviewers assigned automatically/manually are appropriate for the PR.

  • Please check this after reviewing the above items as appropriate for this PR.

GitHub Bot Help

To see a list of available CI bot commands, please comment /bot help.

Oseltamivir and others added 3 commits May 4, 2026 20:37
@Oseltamivir
fix: support 7168 fused MHC hidden size
c43326b 
Signed-off-by: Oseltamivir <bryansg2013@gmail.com>
@mingyangHao
test: add DeepSeek-V4 Pro MHC coverage
c7cbaa4 
Signed-off-by: Mingyang Hao <mingyangHao@users.noreply.github.com>
@pcastonguay
Pre-commit fix
70927b2 
Signed-off-by: Patrice Castonguay <55748270+pcastonguay@users.noreply.github.com>
@pcastonguay pcastonguay requested a review from a team as a code owner May 5, 2026 17:35
@pcastonguay pcastonguay requested review from QiJune and removed request for a team May 5, 2026 17:35
@pcastonguay
Copy link
Copy Markdown
Collaborator Author

pcastonguay commented May 5, 2026

/bot skip --comment "Just fixing pre-commit check"

schetlur-nv
schetlur-nv reviewed May 5, 2026
View reviewed changes
Comment thread cpp/tensorrt_llm/kernels/mhcKernels/mhcFusedHcKernel.cu
@tensorrt-cicd
Copy link
Copy Markdown
Collaborator

tensorrt-cicd commented May 5, 2026

PR_Github #46845 [ skip ] triggered by Bot. Commit: 70927b2 Link to invocation

@tensorrt-cicd
Copy link
Copy Markdown
Collaborator

tensorrt-cicd commented May 5, 2026

PR_Github #46845 [ skip ] completed with state SUCCESS. Commit: 70927b2
Skipping testing for commit 70927b2

Link to invocation

@pcastonguay pcastonguay merged commit 9aa3715 into NVIDIA:feat/deepseek_v4 May 5, 2026
11 of 12 checks passed
@pcastonguay pcastonguay mentioned this pull request May 5, 2026
Closed
@functionstackx functionstackx mentioned this pull request May 5, 2026
Merged
lfr-0531 pushed a commit that referenced this pull request May 7, 2026
@pcastonguay @Oseltamivir
@mingyangHao @lfr-0531
[None][fix] Fix fused MHC for DeepSeek-V4-Pro hidden size (#13771)
99f0b2e 
Signed-off-by: Oseltamivir <bryansg2013@gmail.com>
Signed-off-by: Mingyang Hao <mingyangHao@users.noreply.github.com>
Signed-off-by: Patrice Castonguay <55748270+pcastonguay@users.noreply.github.com>
Co-authored-by: Oseltamivir <bryansg2013@gmail.com>
Co-authored-by: Mingyang Hao <mingyangHao@users.noreply.github.com>
(cherry picked from commit 9aa3715)
Signed-off-by: yuhangh <58161490+heyuhhh@users.noreply.github.com>
mingyangHao added a commit to mingyangHao/TensorRT-LLM that referenced this pull request May 8, 2026
@mingyangHao
[None][perf] mHC fused_hc: 4 small kernel-side perf opts (P0-P3)
af5e3fe 
Pipelined micro-optimizations on top of the existing tcgen05 fused_hc
infrastructure (NVIDIA#13771 added hidden=7168 support; this PR builds on that).
Each change is isolated and gated by config / template constexpr.

P0 — Path D KS=112 enable (Phase 4 scalar-vec tail)
  - fused_tf32_pmap_gemm.cuh: replace the
    `HIDDEN % (WARPS_PER_TOK * 32 * BF16_VEC_LI) == 0` static_assert with a
    scalar-vec tail loop that handles the residue past
    `H_VEC_END = (HIDDEN/H_STRIDE)*H_STRIDE`. Tail uses the same uint4
    LDG/STG path so per-thread BW is unchanged; only some lanes/warps in
    the team idle on the last chunk.
  - mhcFusedHcKernel.cu: relax `isSupportedFhcMmaKS<H, KS>` to drop the
    team-stride alignment check (only `Hidden % BF16_VEC_LI == 0` and
    `h_tiles % KS == 0` remain). Add KS ∈ {7, 14, 28, 56, 112} cases to
    pickFhc and pickFhcAllInOne outer switches.
  - mhc_cuda.py: mirror the same trait change in
    `_fused_hc_mma_ks_supported`. Extend `_FUSED_HC_HALF_MMA_KS` /
    `_FUSED_HC_ALL_MMA_KS` to include the new factors.

P1 — TMA descriptor cache
  - cuTensorMapEncodeTiled is a host-side ~1-2 µs call; each launch builds
    4 descriptors, costing 4-8 µs (25-50% of total wall time at small M).
  - Add a thread-local `unordered_map<TmaDescKey, CUtensorMap>` keyed on
    (base ptr, gmem dims, strides, swizzle, dtype, device_id). Both
    launchers call `getCachedTma2D` instead of `makeTma2D`.
  - CUDA-graph compatible: cuTensorMapEncodeTiled records nothing into the
    stream, so cache miss inside capture is safe. Pointer stability across
    calls is already guaranteed by `_FusedHcWorkspaceCache` in mhc_cuda.py.
  - Return-by-value (CUtensorMap is 128 B POD); no rehash invalidation.

P2 — KS=1 direct store + skip workspace zero
  - At KS=1 each (m_block, n) is owned by exactly one CTA, so the GEMM
    epilogue can write D / sqr_sum directly instead of via atomicAdd.
    Guarded by `if constexpr (kNumSplits == 1)` in both Path B mainloop
    and Path D Phase 2.
  - Launchers skip the `fhcZeroWorkspaces` kernel at KS=1 (y_acc / r_acc
    / done_counter all unused under the direct-store path; Phase 3 at
    KS=1 is just `__threadfence_block + __syncthreads`).
  - Effect: -1 kernel launch + -25 atomic ops per token at KS=1.

P3 — Sinkhorn reciprocal multiply
  - Replace per-element `cm[k] / rs` with `cm[k] * inv_rs`
    (`inv_rs = 1.0f / rs`) in row-normalize. fp32 fdiv is multi-cycle on
    B200 vs fmul peak rate; row-normalize runs HC_MULT × sinkhorn_repeat
    times per token (160 fdivs at sinkhorn=20). Math equivalent under
    fp32 round-off; sinkhorn iteration absorbs last-bit drift.
  - Applied in: mhcKernels.cu BigFuse (Path B/E/F shared), Path D Phase 4
    inline bigfuse, Path F Phase 4 inline bigfuse.

Numerical: residual_cur identical, post_mix_cur within 1.2e-7, comb_mix_cur
within 6e-8 (fp32 epsilon), layer_input within 1.5e-5 (bf16 round-off) vs
upstream Path D KS=56 reference.

Measured on idle B300 (sm_103) at V4-Pro shape (n=4, hidden=7168,
sinkhorn=20):
  - half_mma KS=2 M=4096:        171.10 µs → 165.75 µs   (-3%)
  - half_mma KS=1 M=2048:        216.20 µs → 212.19 µs   (-2%)
  - half_mma KS=1 M=512:         188.35 µs → 185.90 µs   (-1.3%)
  - half_mma KS=56 M=64 (P1):     24.52 µs →  23.87 µs   (-2.7%)
  - all_mma  KS=112 M=64 (P0):       FAIL → 21.92 µs    (newly possible)

Autotuner-driven winners are unchanged in M ≤ 384 (still half_mma KS=112);
the visible production gain is on M ≥ 1024 where KS=2/4 wins are touched
by P1+P3.

Signed-off-by: Mingyang Hao <mingyangh@nvidia.com>
Signed-off-by: mingyangh <mingyangh@nvidia.com>
@mingyangHao mingyangHao mentioned this pull request May 8, 2026
Draft
1 task
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Reviewers

@schetlur-nv schetlur-nv schetlur-nv left review comments

@QiJune QiJune Awaiting requested review from QiJune QiJune is a code owner automatically assigned from NVIDIA/trt-llm-torch-modules

Assignees

@pcastonguay pcastonguay

Labels

deepseek-v4

Projects

None yet

Milestone

No milestone

Development

Successfully merging this pull request may close these issues.

6 participants

@pcastonguay @tensorrt-cicd @schetlur-nv @lfr-0531 @Oseltamivir @mingyangHao