Repository navigation
[None][feat] Add shared KV runtime assembly - #20266
chienchunhung wants to merge 2 commits into
Conversation
|
/bot run |
|
PR_Github #77703 [ run ] triggered by Bot. Commit: |
|
PR_Github #77703 [ run ] completed with state
|
8079d06 to
13eb253
Compare
|
/bot run --disable-fail-fast |
|
PR_Github #77928 [ run ] triggered by Bot. Commit: |
13eb253 to
f01ed34
Compare
|
/bot run --disable-fail-fast |
|
PR_Github #77935 [ run ] triggered by Bot. Commit: |
|
PR_Github/20266-13eb253 #77928 was force-killed by a newer pipeline run. |
|
PR_Github #77935 [ run ] completed with state
|
f01ed34 to
a9854b9
Compare
chienchunhung
changed the title
a9854b9 to
cea5825
Compare
|
/bot run --disable-fail-fast |
cea5825 to
27ce45d
Compare
|
/bot run --disable-fail-fast |
|
PR_Github #79381 [ run ] completed with state
|
c318349 to
9e91263
Compare
|
/bot run --disable-fail-fast |
|
PR_Github #79460 [ run ] triggered by Bot. Commit: |
|
PR_Github #79460 [ run ] completed with state
|
| "compatibility_id": profile.compatibility_id, | ||
| "writer_layout": profile.model_layout, | ||
| "cache_dtype": profile.dtype, | ||
| "attention": {"full": True, "sparse": False, "sink_tokens": 0}, |
There was a problem hiding this comment.
The lender leaves windows, sinks, sparsity, LoRA and draft facts to this scope, and its layout id deliberately excludes them, yet here they are constants. Widening the profile later without touching this function would silently share bytes between different configurations. Suggest deriving each from its field, with a test that each one changes the scope.
| retain their exact access/registration meaning. | ||
| """ | ||
|
|
||
| BackendFactory = Callable[[tuple[SlotRegion, ...], WakeFn, object, OwnerCleanupFn], Backend] |
There was a problem hiding this comment.
With this signature each factory captures lender access privately. Suggest a data-only factory, __call__(inputs: BackendInputs) -> BuiltBackend, so no factory can reach lend_*: inputs carry regions, page layout, meaning digest, device, hold entry, wake, trip, lease limit, trusted registry and config, and BuiltBackend returns the backend with its hint factory.
| if TYPE_CHECKING: | ||
| from tensorrt_llm._torch.pyexecutor.kv_cache.kv_cache_manager_v2 import KVCacheManagerV2 | ||
|
|
||
| OwnerCleanupFn = Callable[[Callable[[], None]], bool] |
There was a problem hiding this comment.
With this dispatcher every provider decides when its hold may be released. This supersedes the earlier finalizer suggestion: the factory takes the hold through inputs.take_hold() before registering anything, the assembly keeps it, and the holder calls release_hold() only once that binding's shutdown() returns empty with no lease unreturned.
| """ | ||
| if not isinstance(profile, ProfileSpec): | ||
| raise TypeError("profile must be ProfileSpec") | ||
| lender = attach_lender(manager, scope=profile.scope, staging=profile.staging) |
There was a problem hiding this comment.
This attaches the lender for one store profile on this rank alone, and a failed profile check leaves it attached. Suggest a small private attach module, called from the assembly entry, whose outcome is agreed across ranks before any binding is built: a failure on any rank fails startup everywhere.
|
|
||
|
|
||
| @dataclass(frozen=True) | ||
| class RegionNamespace: |
There was a problem hiding this comment.
This fixes one memory kind per lender, yet one lender serves every binding: native peers take DEVICE/ALL leases, stores HOST/NAMED. Suggest making the kind a binding property: a routing table maps direction, lending method, row selection and kind to a binding, and each factory registers only the kinds it declares.
| object.__setattr__(self, "scope", derive_kv_scope(self)) | ||
|
|
||
|
|
||
| def derive_kv_scope(profile: ProfileSpec) -> bytes: |
There was a problem hiding this comment.
Every attach needs a scope, native-only ones included. Superseding the earlier placement, suggest moving derive_kv_scope to tensorrt_llm/_torch/disaggregation/kv_scope.py, one rule for P and D, over the KV-meaning parts of SourceIdentity (not GEMM, fusion or allreduce choices), checkpoint revision or commit, writer layout, cache dtype, speculative mode, draft pool and lora_enabled.
| normal shutdown, without retrying attachment or detaching the lender. | ||
|
|
||
| Args: | ||
| manager: Serving manager, never a sizing probe. |
There was a problem hiding this comment.
Once native peers use the lender, the capacity-estimation executor is assembled too, registering backends during sizing. Suggest two rules: its shutdown report must be clean, every binding's shutdown() empty and hold released, or startup fails; and a single_instance factory refuses a new instance while the previous one lives.
| "draft_pool": None, | ||
| } | ||
| encoded = json.dumps(payload, sort_keys=True, separators=(",", ":"), allow_nan=False) | ||
| return b"trtllm-kv-scope-v1\0" + hashlib.sha256(encoded.encode("utf-8")).digest() |
There was a problem hiding this comment.
The scope reaches only the lender's object names, so rows lent in place between peers are never checked against it. Suggest a meaning digest, the scope plus layout declarations that change what bytes mean, that every backend compares with each peer at first contact, refusing a mismatch with a reason.
| return DomainBinding(lender, profile, namespace) | ||
|
|
||
|
|
||
| def check_view_compatibility(lease_view: LeaseView, namespace: RegionNamespace) -> None: |
There was a problem hiding this comment.
This per-lease check refuses every row without keys or addresses, which is every in-place DEVICE row. Suggest removing it: routing matches memory kinds, factories register only accessible regions, and the next #20534 update has the lender validate region geometry and overlap at attach, so RegionNamespace can go too.
| scope: bytes = field(init=False) | ||
|
|
||
| def __post_init__(self) -> None: | ||
| """Reject unsupported facts before any memory or backend exposure.""" |
There was a problem hiding this comment.
Configurations not yet served are better refused in the assembly, before any attach, naming every reason in one error: a manager other than KVCacheManagerV2, CSA2 until generation computes the remote tail, and a remote cache beside a KV connector that loads KV. This supersedes failing at config load.
| The initial profile consumes HOST/NAMED leases from the lender's complete mixed | ||
| memory namespace. Model provenance and writer facts derive its canonical scope; | ||
| public layout and region facts validate the serving profile. Neither | ||
| a provider name nor a caller-supplied hint establishes compatibility or isolation. |
There was a problem hiding this comment.
Nothing here builds hints or vets their endpoints. Suggest each backend family hand over a hint factory, which orchestration calls with a pairing key and the peer's entry, and one trusted endpoint registry owned by the assembly: orchestration checks hint endpoints against it, and each backend checks again.
| raise TypeError("sink_token_length must be an integer") | ||
| if self.sink_token_length != 0: | ||
| raise ValueError("initial profile excludes attention sinks") | ||
| if self.lora_enabled is not False: |
There was a problem hiding this comment.
Refusing lora_enabled disables remote reuse for every LoRA deployment, though many requests lack adapters. Suggest a per-request bypass: orchestration neither lends, fetches nor publishes for requests with a lora_task_id, so no request reads another adapter's bytes. lora_enabled still enters the scope, so deployments with and without LoRA never share names.
| # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. | ||
| # See the License for the specific language governing permissions and | ||
| # limitations under the License. | ||
| """Private runtime assembly shared by transfer and Store integrations. |
There was a problem hiding this comment.
Suggest making this module the single construction point: one entry per executor attaches the lender, agrees the outcome across ranks, and builds every backend and routing once. Reason: with construction split across modules, run-time code can create or swap backends, and ownership and shutdown order lose one source of truth.
|
|
||
|
|
||
| @dataclass(frozen=True) | ||
| class DomainBinding: |
There was a problem hiding this comment.
Suggest dropping the per-profile binding: keep one lender record per KV manager, and put records shared by construction, holding and request handling in a small private module that none of them owns. Reason: tying the lender to one store profile breaks once another backend kind shares the same pages.
| """ | ||
|
|
||
| id: str | ||
| lifecycle_profile_id: str |
There was a problem hiding this comment.
Suggest dropping lifecycle_profile_id and choosing containment per backend from what that backend can prove: a backend that can isolate a failed peer keeps serving, and one that cannot is stopped as a whole. Reason: a per-profile policy ties safety to deployment shape rather than to the evidence a backend provides.
| draft_pool: SourceIdentity | None | ||
| joint_reuse: bool | ||
| staging: StagingOptions | ||
| backend_config: object = None |
There was a problem hiding this comment.
Suggest configuring each backend instance separately, with its own settings and lease limit read once at construction. Reason: one lender serves several backends at once, such as peers plus one or more stores, and a single opaque profile-level configuration cannot express their different limits or failure behavior.
Signed-off-by: Shixiaowei02 <39303645+Shixiaowei02@users.noreply.github.com>
9e91263 to
1a5fc31
Compare
chienchunhung
changed the title
Signed-off-by: Chien-Chun Hung <2679986+chienchunhung@users.noreply.github.com>
1a5fc31 to
dce241c
Compare
|
/bot run --disable-fail-fast |
|
PR_Github #80286 [ run ] triggered by Bot. Commit: |
|
PR_Github #80286 [ run ] completed with state
|
Summary
Add private runtime assembly for shared KV transfers. Use canonical lender views and bind consuming write returns to the originating lender.
Scope
No provider is activated. Request orchestration, DEVICE activation, provider qualification, and qualified nonfatal containment remain separate integration work.
Verification
Revision:
dce241c32ab9b8b61fc40589274680d5ab734dbc.python3 -m pytest -c tests/unittest/pytest.ini -v -ra. The four skips are inherited Blackwell-only foundation cases.bc07138a5; candidate native/build sources match that foundation. This is not provider or multi-rank qualification.Notes
Based on #20534 at
bc07138a5f8ca79088a9a6c1aece3eebdd9d95ed; #20117 is already merged in that base. #20534 must land first.One implementation commit follows the lender prerequisite. Implementation-only diff.