Perf: fan out DeepSeek V4-flash hc_head over idle cores - #822
Conversation
At decode (T=8, one token-tile) both the rms and the reduce scopes ran as a single task, so one core streamed all 512KB of x_flat while the other 23 idled. - Fan the reduce over (token-tile x D-slice) and fold the rsqrt/sigmoid gate into it, deleting the hc_head_pre_fused scope and the transposed pre_t GM buffer - Split the rms sum-of-squares over RMS_OK K-slices into per-slice sq_part rows; the consumer sums them and applies rsqrt inline, so the inv_rms buffer is gone too - Raise LINEAR_OK 8 -> 16 to fill the idle cubes - Zero mixes_raw with create_tensor(init_value=0) on the AICPU instead of a hc_head_seed kernel, which cost ~5us of wall on the cube's critical path to zero 1KB - Size sq_part / mixes_raw from the dynamic t_dim rather than a static T_MAX, so decode (8) and prefill (128) each allocate what they use Five scopes drop to three. Isolated hc_head span 58.8 -> 42.6 us (median of 5, decode T=8, a2a3 single card), and run-to-run spread narrows from 45-65 us to 39-43 us. Constants renamed *_K_CHUNK -> *_K_TILE; prefill_mtp's aliased imports follow.
📝 WalkthroughWalkthroughThe DeepSeek-V4 HC-head implementation now uses tile-based RMS and linear split-K computation, intermediate accumulation buffers, and direct final reduction. Golden HC-head and packed-prefill reference paths use the updated tile constants. ChangesHC-head tiled compute rewrite
Estimated code review effort: 4 (Complex) | ~45 minutes Sequence Diagram(s)sequenceDiagram
participant hc_head
participant RMS_partitions
participant LINEAR_partitions
participant final_reduce
hc_head->>RMS_partitions: Partition sum-of-squares into sq_part
hc_head->>LINEAR_partitions: Accumulate projection blocks into mixes_raw
RMS_partitions->>final_reduce: Supply sq_part
LINEAR_partitions->>final_reduce: Supply mixes_raw
final_reduce->>hc_head: Write BF16 y_flat
Possibly related PRs
Poem
🚥 Pre-merge checks | ✅ 5✅ Passed checks (5 passed)
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
🧹 Nitpick comments (1)
models/deepseek/v4-flash/hc_head.py (1)
17-42: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winReinstate divisibility assertions for tile constants.
The prior
T_MAX/assertions that validated tiling constants againstHC_DIM/Dare removed here. Without them, a future change tohc_mult/hidden_size(which driveHC_DIM/D) could silently violateHC_DIM % (RMS_OK * RMS_K_TILE) == 0,HC_DIM % (LINEAR_OK * LINEAR_K_TILE) == 0, orD % D_SPMD == 0/D_SPMD % D_TILE == 0, causing the reduction loops to skip trailing K/D ranges with no error surfaced.🛡️ Suggested guard
LINEAR_OK = 16 RMS_OK = 16 + +assert HC_DIM % (RMS_OK * RMS_K_TILE) == 0 +assert HC_DIM % (LINEAR_OK * LINEAR_K_TILE) == 0 +assert D % D_SPMD == 0 and D_SPMD % D_TILE == 0Based on the provided line-range change details noting "the earlier derived T_MAX and several transpose/chunk-derived constants and assertions are no longer part of this configuration section."
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@models/deepseek/v4-flash/hc_head.py` around lines 17 - 42, Reinstate validation near the tiling constants in the model configuration: assert that HC_DIM is divisible by RMS_OK * RMS_K_TILE and LINEAR_OK * LINEAR_K_TILE, and that D is divisible by D_SPMD and D_SPMD is divisible by D_TILE. Restore any required derived T_MAX or related tiling guards used by these checks, without changing the existing tile values.
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Nitpick comments:
In `@models/deepseek/v4-flash/hc_head.py`:
- Around line 17-42: Reinstate validation near the tiling constants in the model
configuration: assert that HC_DIM is divisible by RMS_OK * RMS_K_TILE and
LINEAR_OK * LINEAR_K_TILE, and that D is divisible by D_SPMD and D_SPMD is
divisible by D_TILE. Restore any required derived T_MAX or related tiling guards
used by these checks, without changing the existing tile values.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Organization UI
Review profile: CHILL
Plan: Pro Plus
Run ID: a7733cfa-c26f-4c61-afa6-b91c550faa46
📒 Files selected for processing (2)
models/deepseek/v4-flash/hc_head.pymodels/deepseek/v4-flash/prefill_mtp.py
1 participant
Summary
hc_head_reduceover (token-tile x D-slice) and fold the rsqrt/sigmoid gate into it, deleting thehc_head_pre_fusedscope and the transposedpre_tGM buffer. At decode (T=8, one token-tile) the reduce previously ran as a single task, so one core streamed all 512KB ofx_flatwhile the other 23 idled.RMS_OKK-slices, each writing its ownsq_partrow; the consumer sums the rows and applies rsqrt inline, so theinv_rmsbuffer is gone as well.LINEAR_OK8 -> 16 to fill the idle cubes.mixes_rawviapl.create_tensor(init_value=0)on the AICPU instead of ahc_head_seedkernel, which cost ~5us of wall on the cube's critical path to zero 1KB.sq_part/mixes_rawfrom the dynamict_dim/t_linearinstead of a staticT_MAX, so decode (8 rows) and prefill (128 rows) each allocate what they use.RMS_K_CHUNK/LINEAR_K_CHUNK/D_CHUNKto*_K_TILE/D_TILE;prefill_mtp.py's aliased imports follow.Five scopes drop to three. Isolated
hc_headspan 58.8 -> 42.6 us (median of 5, decode T=8, a2a3 single card), with run-to-run spread narrowing from 45-65 us to 39-43 us.