chenchenghao/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 11 Packages Projects Releases Wiki Activity

Merge CP shared-KV batch owner allocation

Browse Source 
PR #10 provides the W2 owner-lane allocation side needed by the current bs>1 CP shared-KV work: supported multi-request extends now derive page owners per request, preserve flattened request order, and use owner-lane allocation instead of legacy page allocation.\n\nThe merge applies cleanly on top of the W3/W4 target current-reuse work because the touched runtime surface is allocator/mem_cache only. Focused remote tests cover the new owner-lane planner/allocation behavior and the existing W4 target reuse regressions.\n\nConstraint: CP shared-KV page ownership is request-relative and page-granular.\nConstraint: Supported bs>1 CP shared-KV allocation must not silently fall back to legacy allocation.\nRejected: Treat multiple requests as one concatenated extend | would assign owners relative to the wrong request boundary.\nRejected: Keep multi_batch fallback | breaks downstream direct-write owner assumptions.\nConfidence: medium\nScope-risk: moderate\nDirective: Preserve request-order flattening when changing owner allocation; do not reorder by lane for convenience.\nTested: Local py_compile for allocator.py, common.py, cp_shared_kv_compute_owner.py, and test_cp_shared_kv_layout.py.\nTested: Local git diff --check --cached.\nTested: Remote g0034 py_compile for touched mem_cache files and test_cp_shared_kv_layout.py.\nTested: Remote g0034 pytest test_cp_shared_kv_layout.py => 34 passed, 3 warnings.\nTested: Remote g0034 pytest test_nsa_cp_utils.py test_cp_shared_kv_runtime.py => 157 passed, 5 warnings, 2 subtests passed.\nNot-tested: Full ETE/perf run; production HiCache load/back up with bs>1 allocation under sustained traffic.
This commit is contained in:
laoyao0822
2026-06-03 07:22:04 +08:00
parent 50d0008705 520770da4c
commit d36f62a3cd
23 changed files with 1431 additions and 19 deletions
Download Patch File Download Diff File Expand all files Collapse all files

144
docs/rfc/RFC-0001.md Normal file
View File

@@ -0,0 +1,144 @@
<!-- GENERATED: do not edit. Source: RFC-0001 -->
<!-- SIGNATURE: sha256:41623be2ca618d29a789d7b0291260685b076cac42014eb3577e6df62e9f1bae -->
# RFC-0001: CP Shared-KV Batch Owner-Lane Allocation
> **Version:** 0.1.0 | **Status:** normative | **Phase:** stable
---
## 1. Summary
### [RFC-0001:C-SUMMARY] Summary (Informative) <a id="rfc-0001c-summary"></a>
This RFC specifies the W2 allocation contract for NSA prefill CP shared-KV when a real extend batch contains more than one request.
**Scope:** This RFC covers owner-lane page-owner planning, allocation ordering, owner-lane invariants, capacity failure semantics, and compatibility requirements for CP shared-KV extend allocation. It does not specify CP metadata construction, KV direct-write kernels, current reuse, HiCache load/backup, or EAGLE draft behavior except where those later workstreams depend on the allocation contract.
**Rationale:** W2 must make batch-size greater than one enter the same owner-lane allocation contract as batch-size one. Without a normative allocation contract, later direct-write and HiCache workstreams cannot trust that newly allocated logical pages belong to the ranks that compute their contents.
*Since: v0.1.0*
---
## 2. Specification
### [RFC-0001:C-REQUEST-LOCAL-OWNER-PLANNING] Request-local owner planning (Normative) <a id="rfc-0001c-request-local-owner-planning"></a>
For CP shared-KV extend allocation, the system MUST derive page-owner requirements independently for each request in the batch.
The system MUST use each request's own prefix length, extend length, page size, CP size, and in-sequence CP split semantics when deriving page owners.
The system MUST NOT derive owner requirements by treating multiple requests as one concatenated sequence.
The system MUST reject or fail fast for a request whose owner requirements cannot be derived under the CP shared-KV page-aligned contract.
**Rationale:** CP in-sequence ownership is request-relative. Combining requests before deriving owners can assign pages to the wrong compute rank even if the final flattened token count is correct.
*Since: v0.1.0*
### [RFC-0001:C-FLATTENED-OWNER-ORDER] Flattened owner order (Normative) <a id="rfc-0001c-flattened-owner-order"></a>
The system MUST concatenate per-request page-owner lists in request order.
The system MUST preserve the order of pages within each request's owner list.
The system MUST NOT reorder the flattened owner list by owner lane, physical page availability, or CP rank.
The system MUST keep the flattened owner list length equal to the number of newly allocated logical pages for the extend batch.
**Rationale:** The extend output cache location tensor is consumed in flattened request-token order. Reordering owners by lane would make page-level ownership appear locally valid while corrupting request-to-token placement.
*Since: v0.1.0*
### [RFC-0001:C-OWNER-LANE-INVARIANT] Owner-lane allocation invariant (Normative) <a id="rfc-0001c-owner-lane-invariant"></a>
For every newly allocated non-dummy logical page, the system MUST choose a logical page id whose owner lane equals the corresponding flattened expected owner.
The owner lane of a non-dummy logical page MUST be computed as `(logical_page_id - 1) % cp_size`.
The system MUST NOT allocate dummy page zero for request data.
The produced extend cache locations MUST remain in flattened request-token order.
The produced extend cache locations MUST map each newly allocated page to the expected owner lane before later CP direct-write consumers observe the locations.
**Rationale:** W3 direct write can only write local KV/index pages safely if the logical page selected during allocation is owned by the rank that computes that page.
*Since: v0.1.0*
### [RFC-0001:C-NO-LEGACY-FALLBACK] No legacy allocation fallback for supported batches (Normative) <a id="rfc-0001c-no-legacy-fallback"></a>
When CP shared-KV is enabled, context-parallel prefill is enabled, in-sequence CP split mode is selected, and all requests in the extend batch satisfy the page-aligned owner-planning contract, the system MUST use owner-lane allocation for batch-size greater than one.
The system MUST NOT use legacy paged allocation for a supported CP shared-KV batch-size greater than one case.
The system MUST NOT report `multi_batch` as a fallback reason for a supported CP shared-KV batch-size greater than one case.
If a CP shared-KV batch-size greater than one case is unsupported, the system MUST emit an explicit fail-fast or fallback reason using the CP shared-KV diagnostic naming scheme before leaving the owner-lane allocation path.
**Rationale:** Silent legacy allocation breaks the page-owner contract while allowing execution to continue into paths that assume owner-lane placement.
*Since: v0.1.0*
### [RFC-0001:C-CAPACITY-AND-EVICTION] Owner-lane capacity and eviction (Normative) <a id="rfc-0001c-capacity-and-eviction"></a>
If owner-lane allocation cannot be satisfied from immediately available owner-lane pages, the system SHOULD use the existing owner-lane eviction and retry mechanism before reporting capacity failure.
After owner-lane eviction and retry, if the required owner lanes remain unavailable, the system MUST raise a recoverable capacity-wait error rather than falling back to legacy allocation.
The capacity-wait error MUST include required, available, and deficit page counts by owner lane when those diagnostics are available.
The system MUST preserve allocator state when an owner-lane allocation attempt fails without selecting pages.
**Rationale:** Owner-lane pressure is lane-specific, not aggregate-capacity-specific. Reporting lane diagnostics preserves scheduler recoverability and prevents legacy allocation from hiding placement bugs.
*Since: v0.1.0*
### [RFC-0001:C-BS1-COMPATIBILITY] Batch-size one compatibility (Normative) <a id="rfc-0001c-bs1-compatibility"></a>
Batch-size one CP shared-KV allocation MUST continue to satisfy the existing owner-lane allocation contract.
A batch-size one request processed through the batch-aware allocation path MUST produce the same owner sequence as the existing single-request in-sequence owner planner for the same prefix length, extend length, page size, and CP size.
The system SHOULD share owner-lane allocation logic between batch-size one and batch-size greater than one to avoid long-lived divergent behavior.
**Rationale:** W2 is an extension of the current owner-lane allocation path, not a replacement with different semantics for existing single-request prefill.
*Since: v0.1.0*
### [RFC-0001:C-VERIFICATION] Verification requirements (Normative) <a id="rfc-0001c-verification"></a>
Tests MUST cover at least two requests with different prefix lengths and extend lengths.
Tests MUST verify that each newly allocated logical page in the output cache locations has the expected owner lane from the flattened owner list.
Tests MUST verify that owner-lane exhaustion for a batch-size greater than one case raises a recoverable capacity-wait error or returns an allocation failure to the caller that triggers that error path.
Tests MUST verify that batch-size one allocation remains compatible with the existing owner-lane behavior.
Tests SHOULD verify that supported batch-size greater than one CP shared-KV allocation does not call legacy allocation and does not emit a `multi_batch` fallback.
**Rationale:** These checks directly exercise the failure modes that W2 is intended to remove: cross-request owner mixing, wrong logical page owners, capacity fallback, and single-request regression.
*Since: v0.1.0*
### [RFC-0001:C-DOWNSTREAM-CONTRACT] Downstream contract (Normative) <a id="rfc-0001c-downstream-contract"></a>
After successful W2 allocation, downstream CP shared-KV direct-write consumers MAY assume that newly allocated logical pages satisfy the owner-lane invariant for their corresponding request-local page positions.
W2 MUST NOT introduce a new planner collective.
W2 MUST NOT require W3, W4, W6, or W7 consumers to recompute request-local page owners from raw prefix and extend lengths when the allocation result and batch plan metadata are already available.
**Rationale:** W2 exists to establish a stable allocation boundary. Later workstreams should validate and consume the contract instead of re-deriving allocation ownership independently.
*Since: v0.1.0*
---
## Changelog
### v0.1.0 (2026-06-03)
Initial draft

44
gov/config.toml Normal file
View File

@@ -0,0 +1,44 @@
[project]
name = "my-project"
# Default owner for new RFCs (uses git user.name if not set)
# default_owner = "@your-handle"
[paths]
docs_output = "docs"
# AI agent directory — target for `govctl init-skills`
# Contains skills/ and agents/ subdirs
# Default: ".claude" (Claude Code / Claude Desktop)
# For Cursor: ".cursor"
# For Windsurf: ".windsurf"
# agent_dir = ".claude"
[schema]
version = 2
# [work_item]
# ID strategy for work items (default: sequential)
# - sequential: WI-YYYY-MM-DD-NNN (solo projects)
# - author-hash: WI-YYYY-MM-DD-{hash}-NNN (multi-person teams, uses git email)
# - random: WI-YYYY-MM-DD-{rand} (simple uniqueness)
# id_strategy = "author-hash"
# [verification]
# Enable project-level default verification guards.
# enabled = true
# default_guards = ["GUARD-GOVCTL-CHECK", "GUARD-CARGO-TEST"]
# [source_scan]
# Scan source files for [[artifact-id]] references during `govctl check`
# enabled = false
# include = ["src/**/*.rs", "crates/**/*.rs", "**/*.md"]
# exclude = []
# [concurrency]
# Maximum seconds to wait for exclusive lock before failing (default: 30)
# Implements [[RFC-0004]] concurrent write safety
# lock_timeout_secs = 30
# [tags]
# Controlled-vocabulary tags for artifact classification — [[RFC-0002:C-RESOURCES]]
# Artifacts may only use tags listed here.
# allowed = ["security", "breaking-change", "performance"]

18
gov/rfc/RFC-0001/clauses/C-BS1-COMPATIBILITY.toml Normal file
View File

@@ -0,0 +1,18 @@
#:schema ../../../schema/clause.schema.json
[govctl]
id = "C-BS1-COMPATIBILITY"
title = "Batch-size one compatibility"
kind = "normative"
status = "active"
since = "0.1.0"
[content]
text = """
Batch-size one CP shared-KV allocation MUST continue to satisfy the existing owner-lane allocation contract.
A batch-size one request processed through the batch-aware allocation path MUST produce the same owner sequence as the existing single-request in-sequence owner planner for the same prefix length, extend length, page size, and CP size.
The system SHOULD share owner-lane allocation logic between batch-size one and batch-size greater than one to avoid long-lived divergent behavior.
**Rationale:** W2 is an extension of the current owner-lane allocation path, not a replacement with different semantics for existing single-request prefill."""

20
gov/rfc/RFC-0001/clauses/C-CAPACITY-AND-EVICTION.toml Normal file
View File

@@ -0,0 +1,20 @@
#:schema ../../../schema/clause.schema.json
[govctl]
id = "C-CAPACITY-AND-EVICTION"
title = "Owner-lane capacity and eviction"
kind = "normative"
status = "active"
since = "0.1.0"
[content]
text = """
If owner-lane allocation cannot be satisfied from immediately available owner-lane pages, the system SHOULD use the existing owner-lane eviction and retry mechanism before reporting capacity failure.
After owner-lane eviction and retry, if the required owner lanes remain unavailable, the system MUST raise a recoverable capacity-wait error rather than falling back to legacy allocation.
The capacity-wait error MUST include required, available, and deficit page counts by owner lane when those diagnostics are available.
The system MUST preserve allocator state when an owner-lane allocation attempt fails without selecting pages.
**Rationale:** Owner-lane pressure is lane-specific, not aggregate-capacity-specific. Reporting lane diagnostics preserves scheduler recoverability and prevents legacy allocation from hiding placement bugs."""

18
gov/rfc/RFC-0001/clauses/C-DOWNSTREAM-CONTRACT.toml Normal file
View File

@@ -0,0 +1,18 @@
#:schema ../../../schema/clause.schema.json
[govctl]
id = "C-DOWNSTREAM-CONTRACT"
title = "Downstream contract"
kind = "normative"
status = "active"
since = "0.1.0"
[content]
text = """
After successful W2 allocation, downstream CP shared-KV direct-write consumers MAY assume that newly allocated logical pages satisfy the owner-lane invariant for their corresponding request-local page positions.
W2 MUST NOT introduce a new planner collective.
W2 MUST NOT require W3, W4, W6, or W7 consumers to recompute request-local page owners from raw prefix and extend lengths when the allocation result and batch plan metadata are already available.
**Rationale:** W2 exists to establish a stable allocation boundary. Later workstreams should validate and consume the contract instead of re-deriving allocation ownership independently."""

20
gov/rfc/RFC-0001/clauses/C-FLATTENED-OWNER-ORDER.toml Normal file
View File

@@ -0,0 +1,20 @@
#:schema ../../../schema/clause.schema.json
[govctl]
id = "C-FLATTENED-OWNER-ORDER"
title = "Flattened owner order"
kind = "normative"
status = "active"
since = "0.1.0"
[content]
text = """
The system MUST concatenate per-request page-owner lists in request order.
The system MUST preserve the order of pages within each request's owner list.
The system MUST NOT reorder the flattened owner list by owner lane, physical page availability, or CP rank.
The system MUST keep the flattened owner list length equal to the number of newly allocated logical pages for the extend batch.
**Rationale:** The extend output cache location tensor is consumed in flattened request-token order. Reordering owners by lane would make page-level ownership appear locally valid while corrupting request-to-token placement."""

20
gov/rfc/RFC-0001/clauses/C-NO-LEGACY-FALLBACK.toml Normal file
View File

@@ -0,0 +1,20 @@
#:schema ../../../schema/clause.schema.json
[govctl]
id = "C-NO-LEGACY-FALLBACK"
title = "No legacy allocation fallback for supported batches"
kind = "normative"
status = "active"
since = "0.1.0"
[content]
text = """
When CP shared-KV is enabled, context-parallel prefill is enabled, in-sequence CP split mode is selected, and all requests in the extend batch satisfy the page-aligned owner-planning contract, the system MUST use owner-lane allocation for batch-size greater than one.
The system MUST NOT use legacy paged allocation for a supported CP shared-KV batch-size greater than one case.
The system MUST NOT report `multi_batch` as a fallback reason for a supported CP shared-KV batch-size greater than one case.
If a CP shared-KV batch-size greater than one case is unsupported, the system MUST emit an explicit fail-fast or fallback reason using the CP shared-KV diagnostic naming scheme before leaving the owner-lane allocation path.
**Rationale:** Silent legacy allocation breaks the page-owner contract while allowing execution to continue into paths that assume owner-lane placement."""

22
gov/rfc/RFC-0001/clauses/C-OWNER-LANE-INVARIANT.toml Normal file
View File

@@ -0,0 +1,22 @@
#:schema ../../../schema/clause.schema.json
[govctl]
id = "C-OWNER-LANE-INVARIANT"
title = "Owner-lane allocation invariant"
kind = "normative"
status = "active"
since = "0.1.0"
[content]
text = """
For every newly allocated non-dummy logical page, the system MUST choose a logical page id whose owner lane equals the corresponding flattened expected owner.
The owner lane of a non-dummy logical page MUST be computed as `(logical_page_id - 1) % cp_size`.
The system MUST NOT allocate dummy page zero for request data.
The produced extend cache locations MUST remain in flattened request-token order.
The produced extend cache locations MUST map each newly allocated page to the expected owner lane before later CP direct-write consumers observe the locations.
**Rationale:** W3 direct write can only write local KV/index pages safely if the logical page selected during allocation is owned by the rank that computes that page."""

20
gov/rfc/RFC-0001/clauses/C-REQUEST-LOCAL-OWNER-PLANNING.toml Normal file
View File

@@ -0,0 +1,20 @@
#:schema ../../../schema/clause.schema.json
[govctl]
id = "C-REQUEST-LOCAL-OWNER-PLANNING"
title = "Request-local owner planning"
kind = "normative"
status = "active"
since = "0.1.0"
[content]
text = """
For CP shared-KV extend allocation, the system MUST derive page-owner requirements independently for each request in the batch.
The system MUST use each request's own prefix length, extend length, page size, CP size, and in-sequence CP split semantics when deriving page owners.
The system MUST NOT derive owner requirements by treating multiple requests as one concatenated sequence.
The system MUST reject or fail fast for a request whose owner requirements cannot be derived under the CP shared-KV page-aligned contract.
**Rationale:** CP in-sequence ownership is request-relative. Combining requests before deriving owners can assign pages to the wrong compute rank even if the final flattened token count is correct."""

16
gov/rfc/RFC-0001/clauses/C-SUMMARY.toml Normal file
View File

@@ -0,0 +1,16 @@
#:schema ../../../schema/clause.schema.json
[govctl]
id = "C-SUMMARY"
title = "Summary"
kind = "informative"
status = "active"
since = "0.1.0"
[content]
text = """
This RFC specifies the W2 allocation contract for NSA prefill CP shared-KV when a real extend batch contains more than one request.
**Scope:** This RFC covers owner-lane page-owner planning, allocation ordering, owner-lane invariants, capacity failure semantics, and compatibility requirements for CP shared-KV extend allocation. It does not specify CP metadata construction, KV direct-write kernels, current reuse, HiCache load/backup, or EAGLE draft behavior except where those later workstreams depend on the allocation contract.
**Rationale:** W2 must make batch-size greater than one enter the same owner-lane allocation contract as batch-size one. Without a normative allocation contract, later direct-write and HiCache workstreams cannot trust that newly allocated logical pages belong to the ranks that compute their contents."""

22
gov/rfc/RFC-0001/clauses/C-VERIFICATION.toml Normal file
View File

@@ -0,0 +1,22 @@
#:schema ../../../schema/clause.schema.json
[govctl]
id = "C-VERIFICATION"
title = "Verification requirements"
kind = "normative"
status = "active"
since = "0.1.0"
[content]
text = """
Tests MUST cover at least two requests with different prefix lengths and extend lengths.
Tests MUST verify that each newly allocated logical page in the output cache locations has the expected owner lane from the flattened owner list.
Tests MUST verify that owner-lane exhaustion for a batch-size greater than one case raises a recoverable capacity-wait error or returns an allocation failure to the caller that triggers that error path.
Tests MUST verify that batch-size one allocation remains compatible with the existing owner-lane behavior.
Tests SHOULD verify that supported batch-size greater than one CP shared-KV allocation does not call legacy allocation and does not emit a `multi_batch` fallback.
**Rationale:** These checks directly exercise the failure modes that W2 is intended to remove: cross-request owner mixing, wrong logical page owners, capacity fallback, and single-request regression."""

33
gov/rfc/RFC-0001/rfc.toml Normal file
View File

@@ -0,0 +1,33 @@
#:schema ../../schema/rfc.schema.json
[govctl]
id = "RFC-0001"
title = "CP Shared-KV Batch Owner-Lane Allocation"
version = "0.1.0"
status = "normative"
phase = "stable"
owners = ["@wxiwnd"]
created = "2026-06-03"
updated = "2026-06-03"
[[sections]]
title = "Summary"
clauses = ["clauses/C-SUMMARY.toml"]
[[sections]]
title = "Specification"
clauses = [
"clauses/C-REQUEST-LOCAL-OWNER-PLANNING.toml",
"clauses/C-FLATTENED-OWNER-ORDER.toml",
"clauses/C-OWNER-LANE-INVARIANT.toml",
"clauses/C-NO-LEGACY-FALLBACK.toml",
"clauses/C-CAPACITY-AND-EVICTION.toml",
"clauses/C-BS1-COMPATIBILITY.toml",
"clauses/C-VERIFICATION.toml",
"clauses/C-DOWNSTREAM-CONTRACT.toml",
]
[[changelog]]
version = "0.1.0"
date = "2026-06-03"
notes = "Initial draft"

103
gov/schema/adr.schema.json Normal file
View File

@@ -0,0 +1,103 @@
{
"$schema": "http://json-schema.org/draft-07/schema#",
"title": "ADR Specification",
"description": "JSON Schema for ADR TOML artifacts after parsing and normalization",
"type": "object",
"required": ["govctl", "content"],
"properties": {
"govctl": {
"type": "object",
"required": ["id", "title", "status", "date"],
"properties": {
"id": {
"type": "string",
"pattern": "^ADR-\\d{4}$"
},
"title": {
"type": "string",
"minLength": 1
},
"status": {
"type": "string",
"enum": ["proposed", "accepted", "rejected", "superseded"]
},
"date": {
"type": "string",
"format": "date"
},
"superseded_by": {
"type": "string",
"pattern": "^ADR-\\d{4}$"
},
"refs": {
"type": "array",
"items": {
"type": "string",
"pattern": "^(RFC-\\d{4}(?::C-[A-Z][A-Z0-9-]*)?|ADR-\\d{4}|WI-\\d{4}-\\d{2}-\\d{2}-(?:[a-f0-9]{4}(?:-\\d{3})?|\\d{3}))$"
}
},
"tags": {
"type": "array",
"items": {
"type": "string",
"pattern": "^[a-z][a-z0-9-]*$"
}
},
"schema": {
"type": "integer"
}
},
"additionalProperties": false
},
"content": {
"type": "object",
"required": ["context", "decision", "consequences"],
"properties": {
"context": {
"type": "string"
},
"decision": {
"type": "string"
},
"consequences": {
"type": "string"
},
"alternatives": {
"type": "array",
"items": {
"type": "object",
"required": ["text"],
"properties": {
"text": {
"type": "string",
"minLength": 1
},
"status": {
"type": "string",
"enum": ["considered", "rejected", "accepted"]
},
"pros": {
"type": "array",
"items": {
"type": "string"
}
},
"cons": {
"type": "array",
"items": {
"type": "string"
}
},
"rejection_reason": {
"type": "string"
}
},
"additionalProperties": false
}
}
},
"additionalProperties": false
}
},
"additionalProperties": false
}

64
gov/schema/clause.schema.json Normal file
View File

@@ -0,0 +1,64 @@
{
"$schema": "http://json-schema.org/draft-07/schema#",
"title": "Clause Specification",
"description": "JSON Schema for clause TOML artifacts with [govctl] metadata + [content] sections",
"type": "object",
"required": ["govctl", "content"],
"properties": {
"govctl": {
"type": "object",
"required": ["id", "title", "kind"],
"properties": {
"id": {
"type": "string",
"pattern": "^C-[A-Z0-9-]+$"
},
"title": {
"type": "string",
"minLength": 1
},
"kind": {
"type": "string",
"enum": ["normative", "informative"]
},
"status": {
"type": "string",
"enum": ["active", "superseded", "deprecated"],
"default": "active"
},
"since": {
"type": "string",
"pattern": "^\\d+\\.\\d+\\.\\d+$"
},
"superseded_by": {
"type": "string",
"pattern": "^(RFC-\\d{4}:)?C-[A-Z0-9-]+$"
},
"anchors": {
"type": "array",
"items": { "type": "string" }
},
"tags": {
"type": "array",
"items": { "type": "string", "pattern": "^[a-z][a-z0-9-]*$" }
},
"schema": {
"type": "integer"
}
},
"additionalProperties": false
},
"content": {
"type": "object",
"required": ["text"],
"properties": {
"text": {
"type": "string",
"minLength": 1
}
},
"additionalProperties": false
}
},
"additionalProperties": false
}

61
gov/schema/guard.schema.json Normal file
View File

@@ -0,0 +1,61 @@
{
"$schema": "http://json-schema.org/draft-07/schema#",
"title": "Verification Guard Specification",
"description": "JSON Schema for Verification Guard TOML artifacts after parsing and normalization",
"type": "object",
"required": ["govctl", "check"],
"properties": {
"govctl": {
"type": "object",
"required": ["id", "title"],
"properties": {
"id": {
"type": "string",
"pattern": "^GUARD-[A-Z][A-Z0-9-]*$"
},
"title": {
"type": "string",
"minLength": 1
},
"refs": {
"type": "array",
"items": {
"type": "string",
"pattern": "^(RFC-\\d{4}(?::C-[A-Z][A-Z0-9-]*)?|ADR-\\d{4}|WI-\\d{4}-\\d{2}-\\d{2}-(?:[a-f0-9]{4}(?:-\\d{3})?|\\d{3}))$"
}
},
"tags": {
"type": "array",
"items": {
"type": "string",
"pattern": "^[a-z][a-z0-9-]*$"
}
},
"schema": {
"type": "integer"
}
},
"additionalProperties": false
},
"check": {
"type": "object",
"required": ["command"],
"properties": {
"command": {
"type": "string",
"minLength": 1
},
"timeout_secs": {
"type": "integer",
"minimum": 1
},
"pattern": {
"type": "string",
"minLength": 1
}
},
"additionalProperties": false
}
},
"additionalProperties": false
}

45
gov/schema/release.schema.json Normal file
View File

@@ -0,0 +1,45 @@
{
"$schema": "http://json-schema.org/draft-07/schema#",
"title": "Release Specification",
"description": "JSON Schema for releases.toml after parsing and normalization",
"type": "object",
"required": ["releases"],
"properties": {
"govctl": {
"type": "object",
"required": [],
"properties": {
"schema": {
"type": "integer"
}
},
"additionalProperties": false
},
"releases": {
"type": "array",
"items": {
"type": "object",
"required": ["version", "date"],
"properties": {
"version": {
"type": "string",
"pattern": "^\\d+\\.\\d+\\.\\d+(?:[-+][0-9A-Za-z.-]+)?$"
},
"date": {
"type": "string",
"format": "date"
},
"refs": {
"type": "array",
"items": {
"type": "string",
"pattern": "^WI-\\d{4}-\\d{2}-\\d{2}-(?:[a-f0-9]{4}(?:-\\d{3})?|\\d{3})$"
}
}
},
"additionalProperties": false
}
}
},
"additionalProperties": false
}

120
gov/schema/rfc.schema.json Normal file
View File

@@ -0,0 +1,120 @@
{
"$schema": "http://json-schema.org/draft-07/schema#",
"title": "RFC Specification",
"description": "JSON Schema for RFC TOML artifacts with [govctl] metadata section",
"type": "object",
"required": ["govctl", "sections"],
"properties": {
"govctl": {
"type": "object",
"required": [
"id",
"title",
"version",
"status",
"phase",
"owners",
"created"
],
"properties": {
"id": {
"type": "string",
"pattern": "^RFC-\\d{4}$"
},
"title": {
"type": "string",
"minLength": 1
},
"version": {
"type": "string",
"pattern": "^\\d+\\.\\d+\\.\\d+$"
},
"status": {
"type": "string",
"enum": ["draft", "normative", "deprecated"]
},
"phase": {
"type": "string",
"enum": ["spec", "impl", "test", "stable"]
},
"owners": {
"type": "array",
"items": { "type": "string" },
"minItems": 1
},
"created": {
"type": "string",
"format": "date"
},
"updated": {
"type": "string",
"format": "date"
},
"supersedes": {
"type": "string",
"pattern": "^RFC-\\d{4}$"
},
"refs": {
"type": "array",
"items": {
"type": "string",
"pattern": "^(RFC-\\d{4}(?::C-[A-Z][A-Z0-9-]*)?|ADR-\\d{4}|WI-\\d{4}-\\d{2}-\\d{2}-(?:[a-f0-9]{4}(?:-\\d{3})?|\\d{3}))$"
}
},
"tags": {
"type": "array",
"items": {
"type": "string",
"pattern": "^[a-z][a-z0-9-]*$"
}
},
"signature": {
"type": "string",
"pattern": "^[0-9a-f]{64}$"
},
"schema": {
"type": "integer"
}
},
"additionalProperties": false
},
"sections": {
"type": "array",
"items": {
"type": "object",
"required": ["title"],
"properties": {
"title": { "type": "string" },
"clauses": {
"type": "array",
"items": { "type": "string" }
}
},
"additionalProperties": false
}
},
"changelog": {
"type": "array",
"items": {
"type": "object",
"required": ["version", "date"],
"properties": {
"version": {
"type": "string",
"pattern": "^\\d+\\.\\d+\\.\\d+$"
},
"date": { "type": "string", "format": "date" },
"notes": { "type": "string" },
"added": { "type": "array", "items": { "type": "string" } },
"changed": { "type": "array", "items": { "type": "string" } },
"deprecated": { "type": "array", "items": { "type": "string" } },
"removed": { "type": "array", "items": { "type": "string" } },
"fixed": { "type": "array", "items": { "type": "string" } },
"security": { "type": "array", "items": { "type": "string" } }
},
"additionalProperties": false
}
}
},
"additionalProperties": false
}

157
gov/schema/work.schema.json Normal file
View File

@@ -0,0 +1,157 @@
{
"$schema": "http://json-schema.org/draft-07/schema#",
"title": "Work Item Specification",
"description": "JSON Schema for Work Item TOML artifacts after parsing and normalization",
"type": "object",
"required": ["govctl", "content"],
"properties": {
"govctl": {
"type": "object",
"required": ["id", "title", "status"],
"properties": {
"id": {
"type": "string",
"pattern": "^WI-\\d{4}-\\d{2}-\\d{2}-(?:[a-f0-9]{4}(?:-\\d{3})?|\\d{3})$"
},
"title": {
"type": "string",
"minLength": 1
},
"status": {
"type": "string",
"enum": ["queue", "active", "done", "cancelled"]
},
"created": {
"type": "string",
"format": "date"
},
"started": {
"type": "string",
"format": "date"
},
"completed": {
"type": "string",
"format": "date"
},
"refs": {
"type": "array",
"items": {
"type": "string",
"pattern": "^(RFC-\\d{4}(?::C-[A-Z][A-Z0-9-]*)?|ADR-\\d{4}|WI-\\d{4}-\\d{2}-\\d{2}-(?:[a-f0-9]{4}(?:-\\d{3})?|\\d{3}))$"
}
},
"tags": {
"type": "array",
"items": {
"type": "string",
"pattern": "^[a-z][a-z0-9-]*$"
}
},
"schema": {
"type": "integer"
}
},
"additionalProperties": false
},
"content": {
"type": "object",
"required": ["description"],
"properties": {
"description": {
"type": "string"
},
"journal": {
"type": "array",
"items": {
"type": "object",
"required": ["date", "content"],
"properties": {
"date": {
"type": "string",
"format": "date"
},
"scope": {
"type": "string"
},
"content": {
"type": "string",
"minLength": 1
}
},
"additionalProperties": false
}
},
"acceptance_criteria": {
"type": "array",
"items": {
"type": "object",
"required": ["text"],
"properties": {
"text": {
"type": "string",
"minLength": 1
},
"status": {
"type": "string",
"enum": ["pending", "done", "cancelled"]
},
"category": {
"type": "string",
"enum": [
"added",
"changed",
"deprecated",
"removed",
"fixed",
"security",
"chore"
]
}
},
"additionalProperties": false
}
},
"notes": {
"type": "array",
"items": {
"type": "string"
}
}
},
"additionalProperties": false
},
"verification": {
"type": "object",
"properties": {
"required_guards": {
"type": "array",
"items": {
"type": "string",
"pattern": "^GUARD-[A-Z][A-Z0-9-]*$"
},
"uniqueItems": true
},
"waivers": {
"type": "array",
"items": {
"type": "object",
"required": ["guard", "reason"],
"properties": {
"guard": {
"type": "string",
"pattern": "^GUARD-[A-Z][A-Z0-9-]*$"
},
"reason": {
"type": "string",
"minLength": 1
}
},
"additionalProperties": false
}
}
},
"additionalProperties": false
}
},
"additionalProperties": false
}

68
gov/work/2026-06-03-implement-cp-shared-kv-batch-owner-lane-allocation.toml Normal file
View File

@@ -0,0 +1,68 @@
#:schema ../schema/work.schema.json
[govctl]
id = "WI-2026-06-03-001"
title = "Implement CP shared-KV batch owner-lane allocation"
status = "done"
created = "2026-06-03"
started = "2026-06-03"
completed = "2026-06-03"
refs = ["RFC-0001"]
[content]
description = "Implement RFC-0001 W2 so CP shared-KV extend allocation supports real batch-size greater than one through owner-lane allocation instead of legacy multi_batch fallback. Scope includes request-local owner planning at allocation time, batch-aware alloc_extend_compute_owner support, owner-lane capacity behavior, and focused unit tests."
[[content.journal]]
date = "2026-06-03"
scope = "gov"
content = "Created active work item for RFC-0001 W2 implementation and defined acceptance criteria."
[[content.journal]]
date = "2026-06-03"
scope = "plan"
content = "Wrote implementation plan docs/superpowers/plans/2026-06-03-cp-shared-kv-batch-owner-lane-allocation.md covering RFC-0001 clauses and tests."
[[content.journal]]
date = "2026-06-03"
scope = "implementation"
content = "Implemented W2 batch owner-lane allocation: added request-local batch owner planner, routed supported bs>1 CP shared-KV extend allocation through alloc_extend_compute_owner, removed allocator bs=1 guard, added bs>1 no-legacy-fallback, owner-lane invariant, capacity wait, unsupported fail-fast, and bs1 compatibility tests. Compliance checker returned COMPLIANT. Local py_compile and govctl render/check passed before remote verification."
[[content.journal]]
date = "2026-06-03"
scope = "verification"
content = "Remote uv verification passed on ubuntu@10.20.32.33:/mnt/beegfs/qjs-work/sglang after rsync source-only sync excluding binaries and caches: test_nsa_cp_utils.py ran 39 tests OK; test_cp_shared_kv_layout.py ran 34 tests OK; test_alloc_pages_with_owners.py ran 10 tests OK. Non-fatal warnings observed: missing Python.h during CUDA utils probe and Triton unsupported CPU fallback."
[[content.acceptance_criteria]]
text = "CP shared-KV batch-size greater than one derives request-local page owners and enters owner-lane allocation"
status = "done"
category = "added"
[[content.acceptance_criteria]]
text = "Batch-aware alloc_extend_compute_owner returns flattened request-token order locations whose new pages match expected owner lanes"
status = "done"
category = "added"
[[content.acceptance_criteria]]
text = "Supported CP shared-KV batch-size greater than one no longer uses legacy multi_batch fallback"
status = "done"
category = "fixed"
[[content.acceptance_criteria]]
text = "Owner-lane exhaustion for batch-size greater than one raises KVCapacityWaitError with lane diagnostics instead of falling back"
status = "done"
category = "fixed"
[[content.acceptance_criteria]]
text = "Batch-size one owner-lane allocation remains compatible with existing behavior"
status = "done"
category = "chore"
[[content.acceptance_criteria]]
text = "Focused W1/W2 unit tests pass in the remote uv environment"
status = "done"
category = "chore"
[[content.acceptance_criteria]]
text = "govctl check and govctl render pass"
status = "done"
category = "chore"

3
python/sglang/srt/mem_cache/allocator.py
View File

@@ -996,9 +996,6 @@ class CPSharedPagedTokenToKVPoolAllocator(PagedTokenToKVPoolAllocator):
and directly persist that page.
"""
if len(prefix_lens_cpu) != 1 or len(seq_lens_cpu) != 1:
raise ValueError("compute-owner allocation supports batch size 1 only")
num_new_pages = get_num_new_pages(
seq_lens=seq_lens_cpu,
page_size=self.page_size,

79
python/sglang/srt/mem_cache/common.py
View File

@@ -14,6 +14,7 @@ from sglang.srt.mem_cache.base_prefix_cache import (
)
from sglang.srt.mem_cache.allocator import compute_owner_lane_free_room_deficits
from sglang.srt.mem_cache.cp_shared_kv_compute_owner import (
build_batch_in_seq_page_compute_owners,
build_in_seq_page_compute_owners,
get_in_seq_page_compute_owner_unavailable_reason,
)
@@ -66,6 +67,27 @@ def _log_cp_shared_kv_alloc_fallback(
)
def _raise_cp_shared_kv_alloc_fail_fast(
*,
reason: str,
batch_size: int,
extend_num_tokens: int,
page_size: int,
) -> None:
message = (
"CP shared-KV compute-owner page assignment is unavailable; "
"rejecting extend allocation instead of falling back to legacy page "
f"allocation. batch_size={batch_size} extend_num_tokens={extend_num_tokens} "
f"page_size={page_size} reason={reason}"
)
logger.error(
"[CP_SHARED_KV_FAIL_FAST][compute_owner_alloc] reason=%s %s",
reason,
message,
)
raise RuntimeError(message)
def _compute_owner_lane_stats_for_eviction(
*,
tree_cache: BasePrefixCache,
@@ -453,7 +475,8 @@ def alloc_paged_token_slots_extend(
)
page_compute_owners = None
compute_owner_unavailable_reason = None
if alloc_extend_compute_owner is not None and len(prefix_lens_cpu) == 1:
compute_owner_planning_required = False
if alloc_extend_compute_owner is not None:
try:
server_args = get_global_server_args()
except ValueError:
@@ -464,28 +487,52 @@ def alloc_paged_token_slots_extend(
and server_args.enable_nsa_prefill_context_parallel
and server_args.nsa_prefill_cp_mode == "in-seq-split"
):
extend_len = int(seq_lens_cpu[0].item() - prefix_lens_cpu[0].item())
page_compute_owners = build_in_seq_page_compute_owners(
extend_len=extend_len,
extend_prefix_len=int(prefix_lens_cpu[0].item()),
page_size=int(allocator.page_size),
cp_size=int(allocator.cp_size),
)
compute_owner_planning_required = True
extend_lens = [
int(seq_lens_cpu[i].item() - prefix_lens_cpu[i].item())
for i in range(len(prefix_lens_cpu))
]
extend_prefix_lens = [
int(prefix_lens_cpu[i].item()) for i in range(len(prefix_lens_cpu))
]
if len(prefix_lens_cpu) == 1:
page_compute_owners = build_in_seq_page_compute_owners(
extend_len=extend_lens[0],
extend_prefix_len=extend_prefix_lens[0],
page_size=int(allocator.page_size),
cp_size=int(allocator.cp_size),
)
else:
page_compute_owners = build_batch_in_seq_page_compute_owners(
extend_lens=extend_lens,
extend_prefix_lens=extend_prefix_lens,
page_size=int(allocator.page_size),
cp_size=int(allocator.cp_size),
)
if page_compute_owners is None:
compute_owner_unavailable_reason = (
get_in_seq_page_compute_owner_unavailable_reason(
for extend_len, extend_prefix_len in zip(
extend_lens, extend_prefix_lens
):
reason = get_in_seq_page_compute_owner_unavailable_reason(
extend_len=extend_len,
extend_prefix_len=int(prefix_lens_cpu[0].item()),
extend_prefix_len=extend_prefix_len,
page_size=int(allocator.page_size),
cp_size=int(allocator.cp_size),
)
or "unknown"
)
if reason is not None:
compute_owner_unavailable_reason = reason
break
if compute_owner_unavailable_reason is None:
compute_owner_unavailable_reason = "unknown"
else:
compute_owner_unavailable_reason = "server_args_not_enabled"
elif alloc_extend_compute_owner is not None:
compute_owner_unavailable_reason = (
"multi_batch" if len(prefix_lens_cpu) != 1 else "unknown"
if page_compute_owners is None and compute_owner_planning_required:
_raise_cp_shared_kv_alloc_fail_fast(
reason=compute_owner_unavailable_reason or "compute_owner_not_available",
batch_size=len(prefix_lens_cpu),
extend_num_tokens=extend_num_tokens,
page_size=int(allocator.page_size),
)
if page_compute_owners is not None:

28
python/sglang/srt/mem_cache/cp_shared_kv_compute_owner.py
View File

@@ -75,3 +75,31 @@ def build_in_seq_page_compute_owners(
owners.extend([owner] * unit_count)
return owners
def build_batch_in_seq_page_compute_owners(
*,
extend_lens: list[int],
extend_prefix_lens: list[int],
page_size: int,
cp_size: int,
) -> Optional[List[int]]:
if len(extend_lens) != len(extend_prefix_lens):
raise ValueError(
"extend_lens and extend_prefix_lens must have the same length, "
f"got {len(extend_lens)} and {len(extend_prefix_lens)}"
)
batch_owners: List[int] = []
for extend_len, extend_prefix_len in zip(extend_lens, extend_prefix_lens):
owners = build_in_seq_page_compute_owners(
extend_len=int(extend_len),
extend_prefix_len=int(extend_prefix_len),
page_size=page_size,
cp_size=cp_size,
)
if owners is None:
return None
batch_owners.extend(owners)
return batch_owners

325
test/registered/unit/mem_cache/test_cp_shared_kv_layout.py
View File

@@ -209,6 +209,57 @@ class TestCPSharedPagedAllocator(CustomTestCase):
[0, 0, 1, 1, 2, 2, 3, 3, 3, 3, 2, 2, 1, 1, 0, 0],
)
def test_batch_compute_owner_page_assignment_flattens_request_order(self):
from sglang.srt.mem_cache.cp_shared_kv_compute_owner import (
build_batch_in_seq_page_compute_owners,
)
owners = build_batch_in_seq_page_compute_owners(
extend_lens=[64 * 3, 64 * 5],
extend_prefix_lens=[0, 64 * 2],
page_size=64,
cp_size=4,
)
self.assertEqual(owners, [0, 1, 2, 0, 1, 2, 3, 3])
def test_batch_compute_owner_page_assignment_matches_single_request(self):
from sglang.srt.mem_cache.cp_shared_kv_compute_owner import (
build_batch_in_seq_page_compute_owners,
build_in_seq_page_compute_owners,
)
batch_owners = build_batch_in_seq_page_compute_owners(
extend_lens=[64 * 7],
extend_prefix_lens=[64 * 2],
page_size=64,
cp_size=4,
)
single_owners = build_in_seq_page_compute_owners(
extend_len=64 * 7,
extend_prefix_len=64 * 2,
page_size=64,
cp_size=4,
)
self.assertEqual(batch_owners, single_owners)
def test_batch_compute_owner_page_assignment_returns_none_for_misaligned_prefix(
self,
):
from sglang.srt.mem_cache.cp_shared_kv_compute_owner import (
build_batch_in_seq_page_compute_owners,
)
owners = build_batch_in_seq_page_compute_owners(
extend_lens=[64, 64],
extend_prefix_lens=[0, 1],
page_size=64,
cp_size=4,
)
self.assertIsNone(owners)
def test_compute_owner_page_assignment_keeps_page_aligned_short_extend(self):
from sglang.srt.mem_cache.cp_shared_kv_compute_owner import (
build_in_seq_page_compute_owners,
@@ -355,6 +406,54 @@ class TestCPSharedPagedAllocator(CustomTestCase):
)
self.assertEqual(allocator.available_size(), page_size * 16)
def test_shared_allocator_compute_owner_alloc_supports_multi_request_flattened_order(
self,
):
from sglang.srt.mem_cache.allocator import CPSharedPagedTokenToKVPoolAllocator
from sglang.srt.mem_cache.cp_shared_kv_compute_owner import (
build_batch_in_seq_page_compute_owners,
)
page_size = 64
cp_size = 4
owners = build_batch_in_seq_page_compute_owners(
extend_lens=[page_size * 3, page_size * 5],
extend_prefix_lens=[0, page_size * 2],
page_size=page_size,
cp_size=cp_size,
)
allocator = CPSharedPagedTokenToKVPoolAllocator(
logical_size=page_size * 32,
physical_size=page_size * 8,
page_size=page_size,
dtype=torch.bfloat16,
device="cpu",
kvcache=None,
need_sort=False,
cp_size=cp_size,
cp_rank=0,
)
locs = allocator.alloc_extend_compute_owner(
prefix_lens=torch.tensor([0, page_size * 2], dtype=torch.int64),
prefix_lens_cpu=torch.tensor([0, page_size * 2], dtype=torch.int64),
seq_lens=torch.tensor(
[page_size * 3, page_size * 7], dtype=torch.int64
),
seq_lens_cpu=torch.tensor(
[page_size * 3, page_size * 7], dtype=torch.int64
),
last_loc=torch.tensor([-1, page_size * 2 - 1], dtype=torch.int64),
extend_num_tokens=page_size * 8,
page_compute_owners=owners,
)
self.assertIsNotNone(locs)
req0_pages = locs[: page_size * 3].view(-1, page_size)[:, 0] // page_size
req1_pages = locs[page_size * 3 :].view(-1, page_size)[:, 0] // page_size
logical_pages = torch.cat([req0_pages, req1_pages])
self.assertEqual(((logical_pages - 1) % cp_size).tolist(), owners)
def test_compute_owner_alloc_does_not_use_torch_isin_for_page_removal(self):
from sglang.srt.mem_cache.allocator import CPSharedPagedTokenToKVPoolAllocator
@@ -555,6 +654,135 @@ class TestCPSharedPagedAllocator(CustomTestCase):
self.assertEqual(allocator.calls, 1)
self.assertEqual(out.numel(), page_size * 8)
def test_compute_owner_alloc_uses_batch_owner_plan_for_multi_request(self):
from types import SimpleNamespace
from sglang.srt.mem_cache import common
page_size = 64
class FakeAllocator:
def __init__(self):
self.page_size = page_size
self.cp_size = 4
self.calls = []
def available_size(self):
return page_size * 100
def alloc_extend_compute_owner(
self,
_prefix_lens,
_prefix_lens_cpu,
_seq_lens,
_seq_lens_cpu,
_last_loc,
extend_num_tokens,
page_compute_owners,
):
self.calls.append(list(page_compute_owners))
return torch.arange(extend_num_tokens, dtype=torch.int64)
def alloc_extend(self, *_args, **_kwargs):
raise AssertionError("legacy allocation should not be used for bs>1")
class FakeTreeCache:
def __init__(self, allocator):
self.token_to_kv_pool_allocator = allocator
def is_chunk_cache(self):
return False
def evict(self, _params):
raise AssertionError("tree eviction should not be used on first success")
allocator = FakeAllocator()
server_args = SimpleNamespace(
enable_nsa_prefill_cp_shared_kv=True,
enable_nsa_prefill_context_parallel=True,
nsa_prefill_cp_mode="in-seq-split",
)
with patch.object(common, "get_global_server_args", return_value=server_args):
out = common.alloc_paged_token_slots_extend(
tree_cache=FakeTreeCache(allocator),
prefix_lens=torch.tensor([0, page_size * 2], dtype=torch.int64),
prefix_lens_cpu=torch.tensor([0, page_size * 2], dtype=torch.int64),
seq_lens=torch.tensor(
[page_size * 3, page_size * 7], dtype=torch.int64
),
seq_lens_cpu=torch.tensor(
[page_size * 3, page_size * 7], dtype=torch.int64
),
last_loc=torch.tensor([-1, page_size * 2 - 1], dtype=torch.int64),
extend_num_tokens=page_size * 8,
)
self.assertEqual(allocator.calls, [[0, 1, 2, 0, 1, 2, 3, 3]])
self.assertEqual(out.numel(), page_size * 8)
def test_compute_owner_alloc_fail_fast_for_multi_request_unsupported_owner_plan(
self,
):
from types import SimpleNamespace
from sglang.srt.mem_cache import common
page_size = 64
class FakeAllocator:
def __init__(self):
self.page_size = page_size
self.cp_size = 4
self.owner_calls = 0
self.legacy_calls = 0
def available_size(self):
return page_size * 100
def alloc_extend_compute_owner(self, *_args, **_kwargs):
self.owner_calls += 1
raise AssertionError("owner allocation should not run without plan")
def alloc_extend(self, *_args, **_kwargs):
self.legacy_calls += 1
raise AssertionError("legacy allocation should not be used for bs>1")
class FakeTreeCache:
def __init__(self, allocator):
self.token_to_kv_pool_allocator = allocator
def is_chunk_cache(self):
return False
def evict(self, _params):
raise AssertionError("tree eviction should not be used")
allocator = FakeAllocator()
server_args = SimpleNamespace(
enable_nsa_prefill_cp_shared_kv=True,
enable_nsa_prefill_context_parallel=True,
nsa_prefill_cp_mode="in-seq-split",
)
with patch.object(common, "get_global_server_args", return_value=server_args):
with self.assertRaises(RuntimeError) as cm:
common.alloc_paged_token_slots_extend(
tree_cache=FakeTreeCache(allocator),
prefix_lens=torch.tensor([0, 1], dtype=torch.int64),
prefix_lens_cpu=torch.tensor([0, 1], dtype=torch.int64),
seq_lens=torch.tensor([page_size, page_size + 1], dtype=torch.int64),
seq_lens_cpu=torch.tensor(
[page_size, page_size + 1], dtype=torch.int64
),
last_loc=torch.tensor([-1, 0], dtype=torch.int64),
extend_num_tokens=page_size * 2,
)
self.assertIn("prefix_not_page_aligned", str(cm.exception))
self.assertEqual(allocator.owner_calls, 0)
self.assertEqual(allocator.legacy_calls, 0)
def test_compute_owner_alloc_skips_aggregate_evict_before_owner_attempt(self):
from types import SimpleNamespace
@@ -880,6 +1108,103 @@ class TestCPSharedPagedAllocator(CustomTestCase):
[params.owner_lane_deficits for params in tree_cache.evict_params],
)
def test_compute_owner_capacity_wait_reports_owner_lane_deficits_for_multi_request(
self,
):
from types import SimpleNamespace
from sglang.srt.mem_cache import common
from sglang.srt.mem_cache.base_prefix_cache import EvictResult
page_size = 64
class FakeAllocator:
def __init__(self):
self.page_size = page_size
self.cp_size = 4
self.owner_calls = []
def available_size(self):
return self.page_size * 4
def allocator_state_str(self):
return "allocator_state_for_test"
def alloc_extend_compute_owner(
self,
_prefix_lens,
_prefix_lens_cpu,
_seq_lens,
_seq_lens_cpu,
_last_loc,
_extend_num_tokens,
page_compute_owners,
):
self.owner_calls.append(list(page_compute_owners))
return None
def compute_owner_lane_stats(self, _page_compute_owners):
return [2, 2, 2, 2], [2, 2, 2, 0], [0, 0, 0, 2]
def alloc_extend(self, *_args, **_kwargs):
raise AssertionError("legacy allocation should not be used for bs>1")
class FakeTreeCache:
def __init__(self):
self.token_to_kv_pool_allocator = FakeAllocator()
self.evict_params = []
def is_chunk_cache(self):
return False
def evictable_size(self):
return page_size * 8
def evict(self, params):
self.evict_params.append(params)
return EvictResult(num_tokens_evicted=0)
def pretty_print(self):
raise AssertionError("recoverable capacity wait should not dump tree")
server_args = SimpleNamespace(
enable_nsa_prefill_cp_shared_kv=True,
enable_nsa_prefill_context_parallel=True,
nsa_prefill_cp_mode="in-seq-split",
)
tree_cache = FakeTreeCache()
with patch.object(common, "get_global_server_args", return_value=server_args):
with self.assertRaises(common.KVCapacityWaitError) as cm:
common.alloc_paged_token_slots_extend(
tree_cache=tree_cache,
prefix_lens=torch.tensor([0, page_size * 2], dtype=torch.int64),
prefix_lens_cpu=torch.tensor(
[0, page_size * 2], dtype=torch.int64
),
seq_lens=torch.tensor(
[page_size * 3, page_size * 7], dtype=torch.int64
),
seq_lens_cpu=torch.tensor(
[page_size * 3, page_size * 7], dtype=torch.int64
),
last_loc=torch.tensor([-1, page_size * 2 - 1], dtype=torch.int64),
extend_num_tokens=page_size * 8,
)
err = cm.exception
self.assertEqual(
tree_cache.token_to_kv_pool_allocator.owner_calls,
[[0, 1, 2, 0, 1, 2, 3, 3], [0, 1, 2, 0, 1, 2, 3, 3]],
)
self.assertEqual(err.required_by_owner, [2, 2, 2, 2])
self.assertEqual(err.available_by_owner, [2, 2, 2, 0])
self.assertEqual(err.deficit_by_owner, [0, 0, 0, 2])
self.assertIn(
[0, 0, 0, 2],
[params.owner_lane_deficits for params in tree_cache.evict_params],
)
def test_alloc_for_extend_releases_req_slots_on_recoverable_capacity_wait(self):
from types import SimpleNamespace