diff --git a/docs/rfc/RFC-0001.md b/docs/rfc/RFC-0001.md new file mode 100644 index 000000000..6c89edada --- /dev/null +++ b/docs/rfc/RFC-0001.md @@ -0,0 +1,144 @@ + + + +# 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) + +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) + +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) + +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) + +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) + +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) + +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) + +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) + +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) + +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 diff --git a/gov/config.toml b/gov/config.toml new file mode 100644 index 000000000..cb2e13027 --- /dev/null +++ b/gov/config.toml @@ -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"] diff --git a/gov/rfc/RFC-0001/clauses/C-BS1-COMPATIBILITY.toml b/gov/rfc/RFC-0001/clauses/C-BS1-COMPATIBILITY.toml new file mode 100644 index 000000000..9269f8d51 --- /dev/null +++ b/gov/rfc/RFC-0001/clauses/C-BS1-COMPATIBILITY.toml @@ -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.""" diff --git a/gov/rfc/RFC-0001/clauses/C-CAPACITY-AND-EVICTION.toml b/gov/rfc/RFC-0001/clauses/C-CAPACITY-AND-EVICTION.toml new file mode 100644 index 000000000..dc24e94f5 --- /dev/null +++ b/gov/rfc/RFC-0001/clauses/C-CAPACITY-AND-EVICTION.toml @@ -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.""" diff --git a/gov/rfc/RFC-0001/clauses/C-DOWNSTREAM-CONTRACT.toml b/gov/rfc/RFC-0001/clauses/C-DOWNSTREAM-CONTRACT.toml new file mode 100644 index 000000000..6174b5e62 --- /dev/null +++ b/gov/rfc/RFC-0001/clauses/C-DOWNSTREAM-CONTRACT.toml @@ -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.""" diff --git a/gov/rfc/RFC-0001/clauses/C-FLATTENED-OWNER-ORDER.toml b/gov/rfc/RFC-0001/clauses/C-FLATTENED-OWNER-ORDER.toml new file mode 100644 index 000000000..eb055afaa --- /dev/null +++ b/gov/rfc/RFC-0001/clauses/C-FLATTENED-OWNER-ORDER.toml @@ -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.""" diff --git a/gov/rfc/RFC-0001/clauses/C-NO-LEGACY-FALLBACK.toml b/gov/rfc/RFC-0001/clauses/C-NO-LEGACY-FALLBACK.toml new file mode 100644 index 000000000..56d49e0ad --- /dev/null +++ b/gov/rfc/RFC-0001/clauses/C-NO-LEGACY-FALLBACK.toml @@ -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.""" diff --git a/gov/rfc/RFC-0001/clauses/C-OWNER-LANE-INVARIANT.toml b/gov/rfc/RFC-0001/clauses/C-OWNER-LANE-INVARIANT.toml new file mode 100644 index 000000000..9d308dc17 --- /dev/null +++ b/gov/rfc/RFC-0001/clauses/C-OWNER-LANE-INVARIANT.toml @@ -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.""" diff --git a/gov/rfc/RFC-0001/clauses/C-REQUEST-LOCAL-OWNER-PLANNING.toml b/gov/rfc/RFC-0001/clauses/C-REQUEST-LOCAL-OWNER-PLANNING.toml new file mode 100644 index 000000000..c37aef293 --- /dev/null +++ b/gov/rfc/RFC-0001/clauses/C-REQUEST-LOCAL-OWNER-PLANNING.toml @@ -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.""" diff --git a/gov/rfc/RFC-0001/clauses/C-SUMMARY.toml b/gov/rfc/RFC-0001/clauses/C-SUMMARY.toml new file mode 100644 index 000000000..51967bb78 --- /dev/null +++ b/gov/rfc/RFC-0001/clauses/C-SUMMARY.toml @@ -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.""" diff --git a/gov/rfc/RFC-0001/clauses/C-VERIFICATION.toml b/gov/rfc/RFC-0001/clauses/C-VERIFICATION.toml new file mode 100644 index 000000000..984ba9192 --- /dev/null +++ b/gov/rfc/RFC-0001/clauses/C-VERIFICATION.toml @@ -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.""" diff --git a/gov/rfc/RFC-0001/rfc.toml b/gov/rfc/RFC-0001/rfc.toml new file mode 100644 index 000000000..17b6c9316 --- /dev/null +++ b/gov/rfc/RFC-0001/rfc.toml @@ -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" diff --git a/gov/schema/adr.schema.json b/gov/schema/adr.schema.json new file mode 100644 index 000000000..cdd376ddd --- /dev/null +++ b/gov/schema/adr.schema.json @@ -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 +} diff --git a/gov/schema/clause.schema.json b/gov/schema/clause.schema.json new file mode 100644 index 000000000..4cece7657 --- /dev/null +++ b/gov/schema/clause.schema.json @@ -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 +} diff --git a/gov/schema/guard.schema.json b/gov/schema/guard.schema.json new file mode 100644 index 000000000..a98c71417 --- /dev/null +++ b/gov/schema/guard.schema.json @@ -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 +} diff --git a/gov/schema/release.schema.json b/gov/schema/release.schema.json new file mode 100644 index 000000000..b93ef2ce9 --- /dev/null +++ b/gov/schema/release.schema.json @@ -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 +} diff --git a/gov/schema/rfc.schema.json b/gov/schema/rfc.schema.json new file mode 100644 index 000000000..741821ea6 --- /dev/null +++ b/gov/schema/rfc.schema.json @@ -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 +} diff --git a/gov/schema/work.schema.json b/gov/schema/work.schema.json new file mode 100644 index 000000000..7f2ffb041 --- /dev/null +++ b/gov/schema/work.schema.json @@ -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 +} diff --git a/gov/work/2026-06-03-implement-cp-shared-kv-batch-owner-lane-allocation.toml b/gov/work/2026-06-03-implement-cp-shared-kv-batch-owner-lane-allocation.toml new file mode 100644 index 000000000..64d72da5e --- /dev/null +++ b/gov/work/2026-06-03-implement-cp-shared-kv-batch-owner-lane-allocation.toml @@ -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" diff --git a/python/sglang/srt/mem_cache/allocator.py b/python/sglang/srt/mem_cache/allocator.py index 7e8c9eefb..a00ab4ea3 100644 --- a/python/sglang/srt/mem_cache/allocator.py +++ b/python/sglang/srt/mem_cache/allocator.py @@ -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, diff --git a/python/sglang/srt/mem_cache/common.py b/python/sglang/srt/mem_cache/common.py index 57d4a4726..b2c8778b7 100644 --- a/python/sglang/srt/mem_cache/common.py +++ b/python/sglang/srt/mem_cache/common.py @@ -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: diff --git a/python/sglang/srt/mem_cache/cp_shared_kv_compute_owner.py b/python/sglang/srt/mem_cache/cp_shared_kv_compute_owner.py index dd0930f60..5eea7a77e 100644 --- a/python/sglang/srt/mem_cache/cp_shared_kv_compute_owner.py +++ b/python/sglang/srt/mem_cache/cp_shared_kv_compute_owner.py @@ -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 diff --git a/test/registered/unit/mem_cache/test_cp_shared_kv_layout.py b/test/registered/unit/mem_cache/test_cp_shared_kv_layout.py index 3b3127279..c5acd9c50 100644 --- a/test/registered/unit/mem_cache/test_cp_shared_kv_layout.py +++ b/test/registered/unit/mem_cache/test_cp_shared_kv_layout.py @@ -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