Repository navigation
[unified-memory] Name the fused KV translate for what it computes (7/7) - #40331
Merged
ch-wan merged 5 commits intoOct 2, 2026
Merged
Conversation
caihuali95
requested review from
BBuf,
DarkSharpness,
Fridge003,
HaiShaw,
HydraQYH,
Qiaolin-Yu,
Ying1123,
alphabetc1,
celve,
hanming-lu,
hebiao064,
hnyls2002,
huangtingwei9988,
hzh0425,
ispobock,
kpham-sgl,
merrymercy,
pyc96,
xiezhq-hermann,
yizhang2077 and
yuan-luo
as code owners
September 19, 2026 09:55
This was referenced Sep 19, 2026
caihuali95
force-pushed
the
mainline/stride-write-loc-kernel
branch
from
September 25, 2026 05:49
5151029 to
47b9fad
Compare
ch-wan
force-pushed
the
mainline/stride-write-loc-kernel
branch
from
September 28, 2026 19:53
47b9fad to
4218b64
Compare
ch-wan
requested review from
JustinTong0323,
sogalin,
wisclmy0611 and
zijiexia
as code owners
September 28, 2026 19:53
github-actions
Bot
added
the
documentation
ch-wan
force-pushed
the
mainline/stride-write-loc-kernel
branch
from
September 28, 2026 21:32
4218b64 to
f326ac4
Compare
ch-wan
changed the title
ch-wan
force-pushed
the
mainline/stride-write-loc-kernel
branch
2 times, most recently
from
September 29, 2026 01:33
bee7c91 to
d509377
Compare
ch-wan
added
run-ci
ch-wan
force-pushed
the
mainline/stride-write-loc-kernel
branch
from
October 1, 2026 06:21
2a6115c to
d794024
Compare
Collaborator
|
/rerun-group unified memory |
Collaborator
|
/rerun-group unified-memory |
Contributor
|
⛔ Known groups: ⛔ Known groups: |
Contributor
|
Results for 🚀 🚀 🚀 🚀 🚀 🚀 |
…putes `write_loc_to_kernel_ids` now serves every KV translate, reads as well as writes, and returns physical ids. Rename it `translate_token_ids` (kernel `translate_token_ids_kernel`, block `TRANSLATE_TOKEN_IDS_BLOCK`), with its kernel-registry entry, callers, test classes and the strided GPU test file. Pure rename.
…as they are The kernel's argument comments still said `loc` is a 1-D int64 tensor of widened ids; it is 1-D or a strided `(rows, cols)` view, int32 or int64, and widened only under DCP. `_row_col_strides` handled a 0-d tensor that the strided path, where both tensors are at least 1-D, never sees. The allocator docstrings that called the translate's result "kernel-facing" or "dense" now call it what it is, the physical id.
…or physical ids `translate_write_loc_for_kernel` returns physical ids and `_assert_kernel_id_bound` bounds them; with the kernel-facing and physical id spaces now the same, call them `translate_write_loc` and `_assert_physical_id_bound`, and say "physical" in the bound's message.
- translate_token_ids_kernel: under out_width, `out` is a packed 1-D buffer of W lanes, not shaped like `loc`. - test_unified_mla_gpu_parity.py still gave the page-major formula kernel_id(t) = (t//ps)*(ps*L) + t%ps for what the unified pool receives; under the token-major views it receives the token's own id. - test_inkling_unified.py said re-running full->swa on the rebound out_cache_loc indexes past the swa v2p table and trips a device assert. A full physical page number is a valid swa virtual page, so the map silently resolves it to an unrelated token's swa slot.
Under the token-major views the id a kernel indexes the per-layer views with is the physical token id, and the identifiers now say so. The comments, docstrings and two unified-memory error messages still called it "kernel-facing"; they now say "physical". Sites that state the equivalence on purpose, that describe a historical bug, or where the word means something else (the Mamba checkpoint pool, the communicator) are left alone.
ch-wan
force-pushed
the
mainline/stride-write-loc-kernel
branch
from
October 2, 2026 00:13
d794024 to
dfb0dac
Compare
This was referenced Oct 2, 2026
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
2 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.
7/7 — name the fused KV translate for what it computes
Base: #40330 (6/7). 5 commits, 28 files, +119/−121.
After 6/7 one kernel serves every KV translate, reads and writes alike, and the kernel-facing
and physical id spaces are the same. The names still said otherwise. This PR renames, with no
behaviour change:
write_loc_to_kernel_ids/write_loc_to_kernel_id_kernel/WRITE_LOC_BLOCKtranslate_token_ids/translate_token_ids_kernel/TRANSLATE_TOKEN_IDS_BLOCKtranslate_write_loc_for_kernel(sub-allocator and composites)translate_write_loc_assert_kernel_id_bound_assert_physical_id_boundtest_write_loc_to_kernel_ids_strided.pytest_translate_token_ids_strided.pyThe kernel registry entry, callers and tests follow. The kernel's argument comments now
describe its real inputs (1-D or strided
(rows, cols), int32 or int64, widened under DCP), anunreachable 0-d branch goes, and the allocator docstrings say "physical" where they said
"kernel-facing".
Comments follow the identifiers: every comment, docstring and unified-memory error message
that called these ids "kernel-facing" now says "physical". Left alone are the sites that
state the equivalence on purpose, one that describes a historical bug, and two where the
word means something else (the Mamba checkpoint pool, the communicator). Three
descriptions that were wrong are corrected: the kernel's
outunderout_width, theparity test's page-major formula, and the inkling test's account of a re-translated loc.
The first commit is a scripted mechanical rename; the next two are comments and the two
identifier renames; the last two are the corrections and the comment sweep, each
droppable on its own.
Validation
The stack is rebased onto main 047be72; every PR's own diff is unchanged file for file.
Unit and GPU tests on the H200, this head against that main, over every non-e2e test file the stack
touches plus
unit/mem_cache,cp,dcp,unit/layers/attention,kernels/ops/memoryandkernels/ops/kvcache: no failure at this head that main does not also have; the extra passes are thestack's own tests. The server-launching multi-GPU files under
cp/anddcp/error the same way onboth sides on a one-GPU run.
End to end at this head on the H200, unified memory against a baseline on the same model and flags
(GSM8K 5-shot, n=200):
Before the rebase, at d509377 — the whole stack — each admitted SM100 backend was run with
unified memory against a baseline on the same model, flags and box (B200, GSM8K 5-shot, n=100):
No write was refused by the physical write-loc check and no device-side assert fired. trtllm_mha
without cuda graphs decodes about 7% slower under unified memory, from the translate launches
an eager step adds.
The commits added since that head fix static-pool write doors, tests and comments; none changes a
unified-memory read or write path.
CI States
Latest PR Test (Base): ⏳ Run #36944919905
Latest PR Test (Extra): ⏳ Run #36944919686
Latest PR Test (AMD ROCm 10): ⏳ Run #36944919884