Conversation
Profiling a steady-state DeepSeek V4 decode (8 cards, dp8/ep8, W8A8) put 99-105 ms of a
~133 ms step in host-side `decode.prepare_early` against 68 ms of device work, so the host
was setting the pace. Inside it, 94 ms was building 56 block tables per step -- seven KV
groups times eight ranks -- while padding, `block_counts` and `copy_shared` together came to
about 7 ms.
`ring_block_table_from_ids` owned 62 ms of that. Each row is a short cycle of physical block
ids tiled across `max_blocks`, which is 2048 for hca_state and 4096 for the two csa_state
tables, and it built that with
torch.tensor(ids).repeat(ceil(max_blocks / len(ids)))[:max_blocks]
so every row allocated a 2048-4098 element temporary and discarded all but the prefix. The
tiling pattern depends only on (max_blocks, len(ids)), so it is cached once and gathered
straight into the destination row.
`block_table_from_ids` allocated one tensor per row for rows holding a handful of ids; it now
builds one for the call. Smaller: about 2 ms.
Measured on the same tree, same warm compile cache, same prompt, with only these two
functions changed:
| | before | after |
|---|---|---|
| decode step | 130-137 ms | 106 ms |
| host `prepare_early` | 99-105 ms | 75 ms |
| device `l3_dispatch` | 68.2 ms | 67.6 ms |
| 4 requests x 64 tokens | 8.24-8.45 s each | 6.93-7.03 s each |
About 19% off the decode step, with the device side unchanged as expected. Generated text is
byte-identical to the unpatched tree on the same node and the fault counters stay at zero;
both new constructions were checked against the old ones on 20 shape/length combinations
before any of this was timed.
These are this node's numbers -- absolute host times depend on the machine -- but the ratio
is the result, and the device time staying put is what says the saving is real host work
rather than a shifted measurement.
Note for anyone profiling this path next: the host prepare and the device dispatch overlap
under async scheduling, so the step is not their sum; and the compile cache has no
fingerprinting since hw-native-sys#197, so it must be cleared by hand after a pypto or pypto-lib bump or
it will silently reload stale binaries.
|
Important
This repository does not receive automatic reviews because it has fewer than 10 stars. ⚙️ Run configurationConfiguration used: Organization UI Review profile: CHILL Plan: Pro Plus Run ID: 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 |
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
1 participant
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.
What this changes
Two functions that build DeepSeek V4 decode block tables, so they stop allocating a temporary per row. No behaviour change: same values, same dtypes, same validation, same error messages.
Why
Profiling a steady-state decode (8 cards,
--dp 8 --ep 8, W8A8, 32-token completions) says the host sets the pace, not the device:WorkerProcess.execute_stepdecode.prepare_early(host)decode.l3_dispatch(device)Inside
prepare_early, building the block tables is 94 ms — seven KV groups × eight ranks per step — against ~7 ms for padding,block_countsandcopy_sharedcombined. Split seven ways:hca_state_ringcsa_state_ringcsa_inner_ringori_pagedidx/hca_cmp/csa_cmpThe ring tables were built as
torch.tensor(ids).repeat(ceil(max_blocks / len(ids)))[:max_blocks], so each row allocated a 2048–4098 element tensor and threw away all but the prefix. The tiling pattern depends only on(max_blocks, len(ids)), so it is cached and gathered directly into the destination row.paged_ori_block_table_from_idsdelegates to the same function and improves with it.block_table_from_idsallocated one tensor per row for rows holding a handful of ids; now one per call. Worth ~2 ms.Result
Same tree, same warm compile cache, same prompt, only these two functions changed:
prepare_earlyl3_dispatch~19% off the decode step. The device number not moving is the point: the saving is host work removed, not a measurement that shifted.
Verification
max_blocks× row-length combinations, before anything was timed." provides telecommunications equipment and services.\",\n \"headquarters\": \"Shenzhen,").no_fault_in=4;507901,507018,TENSOR_WAIT_TIMEOUTandrunning-stalledall zero.Numbers are this node's — absolute host times are machine-dependent — so treat the ratio as the result.
Still on the table
prepare_earlyis 75 ms against 68 ms of device, so the host still just about sets the pace. What is left is spread across the seven constructions rather than concentrated, and three of the seven groups (the ring compress-state ones) change every few steps while four follow the 128-token block, so memoising by group has a ~32 ms ceiling. I tried it and measured 2.7 ms, so there is something wrong with that approach rather than with the idea — not included here.