[None][perf] mHC fused_hc kernel optimizations + DS-V4 entry-boundary RMSNorm fold-in#13892
Merged
lfr-0531 merged 9 commits intoMay 12, 2026
Merged
Conversation
mingyangHao
marked this pull request as ready for review
May 11, 2026 07:40
mingyangHao
requested review from
lfr-0531 and
mikeiovine
and removed request for
a team and
mikeiovine
May 11, 2026 07:40
Collaborator
|
/bot run --add-multi-gpu-test --disable-fail-fast |
Collaborator
|
PR_Github #47697 [ run ] triggered by Bot. Commit: |
mingyangHao
changed the title
Collaborator
|
PR_Github #47697 [ run ] completed with state
|
…lar tail
Phase 4 layer_input loop in fused_allinone_tf32_pmap_gemm_atomic_impl
required HIDDEN to be a multiple of WARPS_PER_TOK * 32 * BF16_VEC_LI.
At KS=112 / HIDDEN=7168 the team stride is 8*32*8=2048 and 7168 % 2048
= 1024 → static_assert blocked the instantiation, and the C++ trait
isSupportedFhcMmaKS<Hidden, KS> rejected the (hidden=7168, ks=112) tuple.
Replace the assert with a vectorized main loop ending at
H_VEC_END = floor(HIDDEN/H_STRIDE)*H_STRIDE, plus a scalar-vec tail
covering [H_VEC_END, HIDDEN). Each tail thread still issues one uint4
LDG/STG (BF16_VEC_LI=8 elements), so per-thread bandwidth matches the
main loop; only the team's lane/warp coverage is shorter, which is fine
since the tail is at most H_STRIDE - BF16_VEC_LI elements wide.
Drop the (Hidden % team-stride) check from isSupportedFhcMmaKS now that
the tail handles the residue. Mirror the relaxed trait in the Python
filter (_fused_hc_mma_ks_supported in mhc_cuda.py).
Add KS ∈ {7, 14, 28, 56, 112} cases to pickFhc and pickFhcAllInOne outer
switches — these are valid divisors of HIDDEN/BLOCK_K = 112 at
hidden=7168. The fhc*InstanceIfSupported<Hidden, KS> wrappers already
guard per-Hidden, so the new cases compile to a TLLM_CHECK_WITH_INFO
on hidden=4096 (where 112 doesn't divide h_tiles=64).
Extend _FUSED_HC_HALF_MMA_KS / _FUSED_HC_ALL_MMA_KS in mhc_cuda.py to
include the new factors. _fused_hc_mma_ks_supported filters per
(hidden, ks) tuple, so 7168-only KS values silently drop out at
hidden=4096.
Numerical: Path D KS=112 at hidden=7168 produces bit-identical
residual_cur to Path D KS=56; post_mix_cur / comb_mix_cur / layer_input
within fp32 rounding noise (max 1.5e-5 bf16 round-off) vs reference.
Unblocks Path D winning M=32-64 buckets (currently capped at half_mma
because Path D could not reach KS=112).
Signed-off-by: Mingyang Hao <mingyangh@nvidia.com>
Signed-off-by: mingyangh <mingyangh@nvidia.com>
cuTensorMapEncodeTiled is a host-side ~1-2 µs call. Each fused_hc launch builds 4 descriptors (residual_in, x_in, W, residual_cur), so the per-call descriptor build is 4-8 µs — 25-50% of total wall time at small M (M ≤ 64). Add a thread-local std::unordered_map<TmaDescKey, CUtensorMap> keyed on all parameters that determine descriptor content (base ptr, gmem dims, strides, swizzle, dtype, device_id). Two launchers (mhcFusedHcLaunch, mhcFusedHcAllInOneLaunch) call into getCachedTma2D() instead of makeTma2D() — the FMA path doesn't use TMA descriptors so it's unaffected. Cache scope: per-host-thread (`thread_local`). Same host thread launching to multiple CUDA streams shares one cache (descriptor content depends only on (device, ptr, shape), not on the stream the kernel runs on). Multi-thread callers each get their own cache (no synchronization overhead). Multi-GPU: device_id is part of the key so a host thread that switches CUDA devices never reuses a descriptor across address spaces. CUDA-graph capture: cuTensorMapEncodeTiled is a pure host function that does not record any stream operation, so cache miss inside capture is safe. The descriptor is passed by value as __grid_constant__; the recorded graph node holds those bytes and replays correctly under workspace-stable replay (already enforced by _FusedHcWorkspaceCache in mhc_cuda.py). Return-by-value: CUtensorMap is a 128 B POD; we copy out before any later cache miss can rehash the underlying unordered_map. Lifetime: entries are never explicitly evicted. Steady-state working set is O(B-bucket × hidden) per thread, typically O(20) entries (~5 KB), never observed to grow under inference workloads. Falls through to makeTma2D() on cache miss, so first-call behavior is identical to before. Signed-off-by: Mingyang Hao <mingyangh@nvidia.com> Signed-off-by: mingyangh <mingyangh@nvidia.com>
…pace zero
At KS=1 each (m_block, n) is owned by exactly one CTA, so the GEMM
epilogue can write D directly instead of going through atomicAdd. The
sqr_sum reduction is similar — single-CTA ownership of each token row.
This removes 24 atomicAdd-per-token (D) + 1 atomicAdd-per-token
(sqr_sum) of GMEM round-trips per launch.
Because the kernel no longer reads the previous value of D / sqr_sum at
KS=1, the launcher can skip the workspace-zero kernel entirely:
- Path B (mhcFusedHcLaunch): skip y_acc + r_acc zero
- Path D (mhcFusedHcAllInOneLaunch): skip y_acc + r_acc + done_counter
zero (Phase 3 at KS=1 uses __threadfence_block + __syncthreads,
never touches done_counter)
KS>1 unchanged: epilogue still atomicAdd, launcher still pre-zeros.
Effect: -1 kernel launch (zero kernel) + -25 atomic ops/token at KS=1.
Most visible at large M where KS=1 is the heuristic choice (M >= 4096).
Verified at hidden=7168, M=2048 / M=4096:
KS=1 (direct, no pre-zero, garbage workspace) vs KS=2 (atomic):
residual_cur diff = 0
post_mix_cur max diff = 2.4e-6
comb_mix_cur max diff = 5.8e-6
layer_input max diff = 6.1e-5 (bf16 round-off)
Signed-off-by: Mingyang Hao <mingyangh@nvidia.com>
Signed-off-by: mingyangh <mingyangh@nvidia.com>
…ultiply
Sinkhorn row-normalize used to do 4 fp32 divisions per iteration:
rs = sum_4(...);
for k in 0..3:
cm[k] = cm[k] / rs + eps
The 4 divisions all share the divisor rs, but NVCC does not always CSE
fdiv in non-fast-math mode. Replace with explicit reciprocal:
inv_rs = 1.0f / rs;
for k in 0..3:
cm[k] = cm[k] * inv_rs + eps
fp32 fdiv on B200 has ~10-cycle latency vs ~4 for fmul, so each sinkhorn
iter saves ~24 cycles on the row-normalize. Column-normalize divisions
use per-k divisors (cs depends on k) so cannot be hoisted — rewritten
in *= 1/(cs+eps) form purely for syntactic consistency, compiler emits
the same PTX as /=.
Applied in three places that run identical sinkhorn:
- mhcKernels.cu :: mhcBigFuseKernel (Path B / E / F bigfuse)
- fused_tf32_pmap_gemm.cuh :: fused_allinone_tf32_pmap_gemm Phase 4
- mhc_fused_fma.cuh :: fused_pmap_gemm_fma_allinone Phase 4
Effect at sinkhorn_repeat=20 (V4-Pro default): ~480 saved cycles per
token (20 iter × 24 cycles), so ~0.25 µs / token at 2 GHz. Most visible
at small M where bigfuse / Phase 4 is a larger fraction of wall time.
Numerical: comb_mix_out max diff vs upstream Path D KS=56 reference is
5.96e-8 (below fp32 epsilon ~1e-7). residual_cur and layer_input
unchanged within bf16 round-off (1e-7 / 2.4e-7).
Signed-off-by: Mingyang Hao <mingyangh@nvidia.com>
Signed-off-by: mingyangh <mingyangh@nvidia.com>
… → 7/3) SASS PC-sampling on M=4096 KS=2 showed ~25% of all stalls landed on a single NANOSLEEP.SYNCS — the pmap warp was input-buffer-starved with only 2 TMA stages for 56 h_tiles. The MMA-side B-stage queue (12 deep) was over- provisioned per the existing P4 wash signal. Trade 5 B-stages (-40 KiB) for +1 input stage (+40 KiB), keeping total dynamic SMEM at 226 KiB. N_B=7 still satisfies the static_assert(N_B_STAGES >= HC_MULT) slack for the MMA pipeline. Measured on B200, M=4096 KS=2 (production-relevant winner): half_mma KS=2 176.5 → 173.9 µs (-1.5%) Path D KS=2 196.2 → 192.8 µs (-1.7%) M=2048 winners (-1.0% to -1.4%) Small-M (M ≤ 64) unchanged Bit-exact output vs baseline at M=4096 KS=2. Signed-off-by: Mingyang Hao <mingyangHao@users.noreply.github.com>
…logue
Inline the next-layer RMSNorm (input_layernorm in modeling_deepseekv4.py's
attention boundary) into Path D's Phase 4 layer_input write. Saves the
separate flashinfer.rmsnorm kernel launch + its HBM round-trip on the
layer_input tensor (~14 KB read + ~14 KB write per token at hidden=7168).
Approach is sglang's `mhc_pre_big_fuse_with_norm_tilelang` trick, adapted
to Path D's tighter SMEM budget. SMEM staging of the un-normalized
layer_input would need 112 KB at TOKS_PER_PASS=8 (KS=2) and won't fit; we
instead let pass 1 write un-normalized bf16 to HBM (same as today), then
pass 2 re-LDGs from L2 (hot from the just-issued STGs) and applies
rsqrt(mean(li²)+norm_eps) * norm_weight before re-STG-ing the normalized
bf16 back to the same address. No FMA recompute (a naive recompute path
costs +12 µs vs the +8 µs L2-reread path on M=4096).
Cross-warp sum_sq reduction (only needed when WARPS_PER_TOK > 1, i.e.
KS ≥ 16) uses a per-team PTX named barrier (bar.sync %team_id, %nthreads)
rather than __syncthreads so inactive teams (which `continue`'d past the
work due to shape_m masking) don't deadlock. Each team gets its own
barrier ID in [1, TOKS_PER_PASS].
API:
- `mhcFusedHcAllInOneLaunch` gains `norm_weight` (bf16 [hidden],
nullable) + `norm_eps` (float). When norm_weight is null, behavior is
identical to the un-fused path (kFuseNorm=false template instance).
- `torch.ops.trtllm.mhc_fused_hc` gains `Tensor? norm_weight=None,
float norm_eps=0.0` (only valid with backend=2 / Path D).
- `mhc_fused_hc` Python entry passes norm_weight through to the
autotuner; `get_valid_tactics` restricts to fused_all_mma tactics when
norm_weight is set.
Measured on B200 with M=4096 hidden=7168 KS=2:
Path D alone (un-fused): 193.2 µs
Path D + flashinfer.rmsnorm (production un-fused): 250.2 µs
Path D w/ fused RMSNorm: 202.0 µs (-19.3%
vs un-fused)
Crossover at M≈1024; M=2048 saves 7.8%; M=512 has +5.6% overhead (caller
should only enable for large-M batches). Un-fused mode (norm_weight=None)
is bit-for-bit identical to the previous binary on all metrics.
Correctness: residual_cur / post_mix / comb_mix are bit-exact vs the
un-fused path; layer_input differs by at most 1 ULP of bf16 (sum_sq is
computed from the bf16-rounded values to match a separate RMSNorm
kernel's behavior).
Signed-off-by: Mingyang Hao <mingyangHao@users.noreply.github.com>
…backends)
Builds on the previous Path D fusion. Now all four mHC backends accept an
optional `norm_weight` and emit normalized layer_input inline:
- Path B (fused_half_mma) / Path E (fused_half_fma):
The shared `mhcBigFuseKernel` (2-kernel split) gains a `kFuseNorm`
template. Phase 2's bigfuse warps (1..N) accumulate sum_sq while writing
un-normalized bf16 layer_input; after a __syncthreads we reduce across
the bigfuse warps via shared mem, compute rsqrt once per token in warp
0, and run a pass 2 that re-LDGs layer_input from L2 and writes the
normalized bf16 back.
- Path F (fused_all_fma):
`fused_pmap_gemm_fma_allinone`'s Phase 4 layer_input loop gains the
same kFuseNorm pass 1 (sum_sq) + pass 2 (re-LDG + normalize) pattern,
using cross-warp SMEM reduction (whole block syncs per token).
- Path D (fused_all_mma): already covered in the previous commit.
The thop op no longer restricts `norm_weight` to backend=2; the autotuner
also drops its Path-D-only restriction in get_valid_tactics — all four
backends are valid tactics when norm_weight is set.
Measured on B200 hidden=7168 (probe_p6c_allpaths.py, fused vs
un-fused-Path + flashinfer.rmsnorm baseline):
Path B (half_mma, M=4096, KS=2): -4.3%
Path E (half_fma, M=512, tn=3 KS=4): -6.9%
Path D (all_mma, M=4096, KS=2): -4.9%
Path F (all_fma, M=64, tn=12): -8.4%
Un-fused mode unchanged (microbench Path D KS=2: 192.91 µs vs 192.87 µs
on previous binary — within noise). residual_cur / post_mix / comb_mix
remain bit-exact; layer_input differs by ≤1 ULP of bf16 (sum_sq is taken
post-bf16-round to match the reference flashinfer.rmsnorm).
Signed-off-by: Mingyang Hao <mingyangHao@users.noreply.github.com>
… note
Addresses three blocking items from PR review:
1. **TMA descriptor cache now bounded with LRU eviction (cap=128).**
Previously the thread_local `unordered_map` grew unbounded. Under CUDA
Graph + workspace-stable replay (production decode path) the cache
reaches a fixed working set of ~20 entries — but in eager mode with the
PyTorch caching allocator handing out fresh `base` pointers on every
workspace miss (dynamic shapes, reset_peak_memory_stats, etc.) the cache
could grow without bound. Adds a standard intrusive-list LRU: O(1)
move-to-front on hit, evict-back on miss when size == cap. 128 entries
× ~256 B ≈ 32 KB per host thread — fits in L1 with headroom.
Also drops an unused `#include <cstring>` noticed in passing.
2. **Removed stale "Only Path D supports the fused norm" docstrings.**
`mhc_fused_hc` Python docstring and the inline comment in
`MhcFusedHcRunner.forward` both still claimed Path D was the only
backend with fused next-layer RMSNorm. P6c extended the fusion to
Path B/E (via `mhcBigFuseKernel`) and Path F (via
`fused_pmap_gemm_fma_allinone`); the autotuner now explores all four
backends regardless of `norm_weight`. Updated comments to match.
3. **Documented why Path F skips the KS=1 workspace-zero optimization.**
P2 added a `ks > 1` guard around `fhcZeroWorkspaces` in the MMA all-in-
one launcher because its kernel epilogue switches to direct-store at
KS=1 (no atomicAdd). Path F's kernel
(`fused_pmap_gemm_fma_allinone`) does not have a parallel direct-store
branch — it always uses `atomicAdd` into y_acc/r_acc, so the workspace
zero is still required. Added a comment in `mhcFusedHcFmaAllInOneLaunch`
explaining the asymmetry; implementing the direct-store fast path on
the FMA side is deferred (Path F is the small-M, M ≤ 32 winner where
the ~1 µs workspace zero is a small fraction of total cost).
Updates `test_mhc_fused_hc_mma_tactic_filter_hidden_sizes` to the post-P0
support set (hidden=7168 now accepts all divisors of h_tiles=112:
{1, 2, 4, 7, 8, 14, 16, 28, 56, 112}). The previous test still encoded
the pre-P0 expectation and should have been updated in the P0 commit.
Unit tests (`tests/unittest/_torch/modules/test_mhc.py`): 46/46 passed
on B200 after the fixes.
Signed-off-by: Mingyang Hao <mingyangHao@users.noreply.github.com>
mingyangHao
force-pushed
the
mingyangh/mhc-v4pro-perf
branch
from
May 11, 2026 11:14
9251945 to
bdff29d
Compare
…layernorm
Plumb norm_weight/norm_eps through mHC.fused_hc into the autotuned kernel,
and pass input_layernorm at the layer-entry boundary in modeling_deepseekv4.
Drops the standalone self.input_layernorm(layer_input) call before attention
at every fused-entry layer (every non-first layer whose previous block
deferred its post_mapping, which is the common case in DS-V4).
The non-fused entry branch (layer 0 / engram / resolved state) still applies
input_layernorm inside _entry_boundary so the caller invariant holds:
layer_input is always RMSNorm-normalized on return.
Tests: add test_mhc_fused_hc_fused_norm comparing mHC.fused_hc(..., norm_weight=,
norm_eps=) against fused_hc(...) + standalone RMSNorm; 4 parametrizations
(M={128,2048}, hidden={4096,7168}) pass within bf16/fp32 ULP tolerance.
Existing test_mhc_fused_hc and test_mhc_fused_hc_backends remain green.
Signed-off-by: Mingyang Hao <mingyangHao@users.noreply.github.com>
mingyangHao
force-pushed
the
mingyangh/mhc-v4pro-perf
branch
from
May 11, 2026 11:22
bdff29d to
ba5418f
Compare
mingyangHao
changed the title
mingyangHao
changed the title
mingyangHao
changed the title
mingyangHao
changed the title
mingyangHao
changed the title
Collaborator
|
/bot run --add-multi-gpu-test --disable-fail-fast |
Collaborator
|
PR_Github #47742 [ run ] triggered by Bot. Commit: |
Collaborator
|
PR_Github #47742 [ run ] completed with state |
lfr-0531
pushed a commit
that referenced
this pull request
May 14, 2026
… RMSNorm fold-in (#13892) Signed-off-by: Mingyang Hao <mingyangh@nvidia.com> Signed-off-by: mingyangh <mingyangh@nvidia.com> Signed-off-by: Mingyang Hao <mingyangHao@users.noreply.github.com> Co-authored-by: Mingyang Hao <mingyangHao@users.noreply.github.com>
lfr-0531
pushed a commit
to lfr-0531/TensorRT-LLM
that referenced
this pull request
May 29, 2026
… RMSNorm fold-in (NVIDIA#13892) Signed-off-by: Mingyang Hao <mingyangh@nvidia.com> Signed-off-by: mingyangh <mingyangh@nvidia.com> Signed-off-by: Mingyang Hao <mingyangHao@users.noreply.github.com> Co-authored-by: Mingyang Hao <mingyangHao@users.noreply.github.com> (cherry picked from commit 50aa581) Signed-off-by: Fanrong Li <lfr-0531@users.noreply.github.com>
lfr-0531
pushed a commit
to lfr-0531/TensorRT-LLM
that referenced
this pull request
Jun 1, 2026
… RMSNorm fold-in (NVIDIA#13892) Signed-off-by: Mingyang Hao <mingyangh@nvidia.com> Signed-off-by: mingyangh <mingyangh@nvidia.com> Signed-off-by: Mingyang Hao <mingyangHao@users.noreply.github.com> Co-authored-by: Mingyang Hao <mingyangHao@users.noreply.github.com> (cherry picked from commit 50aa581) Signed-off-by: Fanrong Li <lfr-0531@users.noreply.github.com>
lfr-0531
pushed a commit
to lfr-0531/TensorRT-LLM
that referenced
this pull request
Jun 3, 2026
… RMSNorm fold-in (NVIDIA#13892) Signed-off-by: Mingyang Hao <mingyangh@nvidia.com> Signed-off-by: mingyangh <mingyangh@nvidia.com> Signed-off-by: Mingyang Hao <mingyangHao@users.noreply.github.com> Co-authored-by: Mingyang Hao <mingyangHao@users.noreply.github.com> (cherry picked from commit 50aa581) Signed-off-by: Fanrong Li <lfr-0531@users.noreply.github.com>
lfr-0531
pushed a commit
to lfr-0531/TensorRT-LLM
that referenced
this pull request
Jun 7, 2026
… RMSNorm fold-in (NVIDIA#13892) Signed-off-by: Mingyang Hao <mingyangh@nvidia.com> Signed-off-by: mingyangh <mingyangh@nvidia.com> Signed-off-by: Mingyang Hao <mingyangHao@users.noreply.github.com> Co-authored-by: Mingyang Hao <mingyangHao@users.noreply.github.com> (cherry picked from commit 50aa581) Signed-off-by: Fanrong Li <lfr-0531@users.noreply.github.com>
lfr-0531
pushed a commit
to lfr-0531/TensorRT-LLM
that referenced
this pull request
Jun 10, 2026
… RMSNorm fold-in (NVIDIA#13892) Signed-off-by: Mingyang Hao <mingyangh@nvidia.com> Signed-off-by: mingyangh <mingyangh@nvidia.com> Signed-off-by: Mingyang Hao <mingyangHao@users.noreply.github.com> Co-authored-by: Mingyang Hao <mingyangHao@users.noreply.github.com>
lfr-0531
pushed a commit
to lfr-0531/TensorRT-LLM
that referenced
this pull request
Jun 11, 2026
… RMSNorm fold-in (NVIDIA#13892) Signed-off-by: Mingyang Hao <mingyangh@nvidia.com> Signed-off-by: mingyangh <mingyangh@nvidia.com> Signed-off-by: Mingyang Hao <mingyangHao@users.noreply.github.com> Co-authored-by: Mingyang Hao <mingyangHao@users.noreply.github.com>
lfr-0531
pushed a commit
to lfr-0531/TensorRT-LLM
that referenced
this pull request
Jun 12, 2026
… RMSNorm fold-in (NVIDIA#13892) Signed-off-by: Mingyang Hao <mingyangh@nvidia.com> Signed-off-by: mingyangh <mingyangh@nvidia.com> Signed-off-by: Mingyang Hao <mingyangHao@users.noreply.github.com> Co-authored-by: Mingyang Hao <mingyangHao@users.noreply.github.com>
lfr-0531
pushed a commit
to lfr-0531/TensorRT-LLM
that referenced
this pull request
Jun 12, 2026
… RMSNorm fold-in (NVIDIA#13892) Signed-off-by: Mingyang Hao <mingyangh@nvidia.com> Signed-off-by: mingyangh <mingyangh@nvidia.com> Signed-off-by: Mingyang Hao <mingyangHao@users.noreply.github.com> Co-authored-by: Mingyang Hao <mingyangHao@users.noreply.github.com> (cherry picked from commit 50aa581) Signed-off-by: Fanrong Li <lfr-0531@users.noreply.github.com>
lfr-0531
pushed a commit
to lfr-0531/TensorRT-LLM
that referenced
this pull request
Jun 13, 2026
… RMSNorm fold-in (NVIDIA#13892) Signed-off-by: Mingyang Hao <mingyangh@nvidia.com> Signed-off-by: mingyangh <mingyangh@nvidia.com> Signed-off-by: Mingyang Hao <mingyangHao@users.noreply.github.com> Co-authored-by: Mingyang Hao <mingyangHao@users.noreply.github.com> (cherry picked from commit 50aa581) Signed-off-by: Fanrong Li <lfr-0531@users.noreply.github.com>
lfr-0531
pushed a commit
to lfr-0531/TensorRT-LLM
that referenced
this pull request
Jun 16, 2026
… RMSNorm fold-in (NVIDIA#13892) Signed-off-by: Mingyang Hao <mingyangh@nvidia.com> Signed-off-by: mingyangh <mingyangh@nvidia.com> Signed-off-by: Mingyang Hao <mingyangHao@users.noreply.github.com> Co-authored-by: Mingyang Hao <mingyangHao@users.noreply.github.com> (cherry picked from commit 50aa581) Signed-off-by: Fanrong Li <lfr-0531@users.noreply.github.com>
lfr-0531
pushed a commit
to lfr-0531/TensorRT-LLM
that referenced
this pull request
Jun 17, 2026
… RMSNorm fold-in (NVIDIA#13892) Signed-off-by: Mingyang Hao <mingyangh@nvidia.com> Signed-off-by: mingyangh <mingyangh@nvidia.com> Signed-off-by: Mingyang Hao <mingyangHao@users.noreply.github.com> Co-authored-by: Mingyang Hao <mingyangHao@users.noreply.github.com> (cherry picked from commit 50aa581) Signed-off-by: Fanrong Li <lfr-0531@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.
@coderabbitai summary
● ## Summary
mHC (Multi-Head Hyper-Connection) kernel perf optimization for DeepSeek-V4
(Pro: hidden=7168, Flash: hidden=4096). Targets
mhc_fused_hc— thecross-layer post+pre boundary op called every layer (attention and MoE
sides).
Changes (7 commits)
[perf] Path D: enable KS=112 at hidden=7168[perf] cache TMA descriptorscuTensorMapEncodeTiledkeyed on (ptr, dims, swizzle, dtype, device). Saves ~1-2 µs/call host-side.[perf] KS=1 direct-store + skip workspace zeroatomicAddwith direct store and skip the y_acc/r_acc memset.[perf] sinkhorn rcpfdivwith 1fdiv+ 4fmulper row-normalize.[perf] pipeline rebalance N_B=12→7, N_INPUT=2→3NANOSLEEP.SYNCS— pmap warp input-buffer-starved. Re-balanced same total SMEM from over-provisioned B-stages to TMA-input depth.[perf] Path D: fuse next-layer RMSNormnorm_weightto Path D. Phase 4 layer_input epilogue accumulatessum_sqduring pass 1 (write bf16), pass 2 re-LDGs from L2 + normalize + STG. Saves the separate flashinfer.rmsnorm kernel launch + HBM round-trip. Per-team named PTX barrier for cross-warp sum_sq reduce.[perf] extend fused RMSNorm to Path B/E/FmhcBigFuseKernel(covers Path B + E) andfused_pmap_gemm_fma_allinone(Path F). All four backends now acceptnorm_weight; autotuner is no longer restricted to Path D when fusing.Perf (B300, V4-Pro hidden=7168, autotuner-best per M)
mHC + RMSNorm chain wall time, vs initial baseline (pre-P0 + flashinfer.rmsnorm):
Fused mode wins at every M. Small-M wins are dominated by eliminating the
separate flashinfer.rmsnorm launch overhead (~7-10 µs); large-M wins come
from skipping the layer_input HBM round-trip.
Correctness
All commits validated for
residual_cur / post_mix / comb_mixbit-exactvs the baseline at multiple M/KS configurations.
layer_inputdiffers by ≤1 ULP of bf16 (P6 computessum_sqfrom thebf16-rounded values to match a separate flashinfer.rmsnorm).
Un-fused mode (
norm_weight=None) is unchanged byte-for-byte vs theprevious binary on microbench.
Test plan
/bot run)