Repository navigation
[Qwen 3.8] LMCache offload for QSA + GDN + PLE state - #2636
ZLkanyo009 wants to merge 69 commits into
Conversation
…oints Qwen3.8-Flash-Next (qwen4_exp / qwen4_exp_text) refused KV offload: none of the existing layouts carries its state. Its full-attention pages are a QSA arena (k, v, index_raw, index_compressed, plus rope for multimodal), not the MHA pool Qwen3.5 checkpoints into, and its PLE conv / n-gram windows share the GDN state's lifetime but sit outside gdn-paged-state-v1. Reuse the Qwen3.5 `qwen` state tier instead of a new codec: - pool_layout/qsa_page_unit_geometry.py: one contiguous region per QSA field and layer, with the mha_page_unit arithmetic; published bytes exclude the alignment gaps between fields. - Qwen4ExpMetadataBuilder turns PAGE checkpoints on itself, aimed at the QSA pool, and appends the PLE conv and n-gram planes after GDN conv/ssm. The layout id gains `ple_conv=...:ple_ngram=...:ple_order=conv,ngram`. get_kv_transfer_tensors publishes byte views via page_region() and attaches the paged_state_checkpoint_spec. - KVCacheTensor gains index_raw / index_compressed / rope_positions after index_scale; older models keep their byte order. - select_offload_layout maps qwen4_exp / qwen4_exp_text to `qwen`; offload_layout=dense is still refused. Pipeline parallelism and RapidServe stay on fork checkpoints and validate_config keeps refusing offload there. - LMCache MP native saves step one stored image at a time instead of jumping to the last checkpoint, so the image a strictly-shorter load needs is not skipped. qwen4_exp_text is deliberately not added to GDN_PAGE_CHECKPOINT_MODEL_TYPES: that set aims the image at the MHA pool.
…paths Both paths missed a prefix they had stored whenever the state image was not where the lookup could land. In-process lmcache_offload: a full-prompt hit gave back its last token and was floored to a chunk, so a prompt of whole chunks (768 at chunk 256) was looked up as 512 -- below both rungs it keeps (the anchor at 704, the ladder at 768). `_joint_kv_boundary` found no rung, no load was offered, and the request recomputed everything. A hybrid claims the joint boundary, which is already capped one hash block short of the prompt, so the K3 scheduler (and its Qwen subclass) now keeps the hit whole for per-request-state requests (`_withholds_last_prompt_token`). The KV leg moves the covering chunk and the claim stops at the anchor. Dense requests still withhold. lmcache_mp: a state image is stored only at a chunk boundary, but the prompt-end anchor is floored to the hash block, so about three anchors in four never reached LMCache; at the default 8192 interval the ladder cannot stand in, and the next turn or an identical resend found nothing. With images on, the MP scheduler now asks BlockManager to floor the anchor to the chunk grid (`align_checkpoint_anchor`). This costs HBM less than one chunk of reuse and adds no checkpoint and no forward. PAGE-only (replay) layouts are unchanged. Also note why Qwen4Exp's state transfer is not readable midstep: nothing snapshots the PLE windows inside a forward. Measured on MI355X, TP1 (HBM evicted first, then the same prompt resent): - in-process, 768 tokens: done 0 -> 704, 768 tokens retrieved, state from the CPU tier. - mp, default interval: 768 -> resumes at 512, 1000 -> 768, 1000 + 300 -> 768 (all 0 before); a 1050-token prompt whose anchor was already on the grid is unchanged at 1024.
🏷️ CI GuideRuns automatically on every eligible PR before approval:
Heavy model tests:
|
lmcache_engine_id() only carried the DP rank, so every TP1 server on one host was 'atom-offload-dp0'. The lookup socket is derived from it under a fixed host directory (/tmp/vllm_rpc without vLLM); each server's bind takes the path over, and a scheduler started next to other servers connects to whichever bound last. Measured with 5 in-process lmcache_offload servers in one container: one worker held 3 scheduler connections, two had none, and the two orphaned servers ran with joint KV boundaries=0 (no CPU fetches). Add config.port (the model runner's MASTER_PORT, unique per server on a host and identical across its processes). Without a port the id is unchanged. Co-authored-by: Cursor <cursoragent@cursor.com>
Co-authored-by: Cursor <cursoragent@cursor.com>
#2635 has the GDN builder zero the PAGE units a state checkpoint gives back (`clear_page_units`), walking the MHA `kv_pools`. The Qwen3.8 builder keeps its PAGE images in the QSA arena and leaves `kv_pools` empty, so with both PRs in, `--state-checkpoint-interval-tokens` fails in `warmup_per_req_cache` with `ValueError: max() iterable argument is empty`. Override `clear_page_units` for Qwen3.8: upload the ids through a pinned buffer fenced on its previous upload, as the GDN version does, and put those blocks of every QSA field back to its declared fill with the new `LayerMajorArena.reset_entries`. A 1-byte field (fp8 K/V, fp4 / e8m0 kept as uint8) is filled through a uint8 view: `index_fill_` has no fp8 kernel on either device. The id buffer is rebuilt if a later arena has more blocks. Carried in this PR so the merge order does not matter. The override uses nothing from #2635, and until #2635 lands nothing calls it: its only callers (`checkpoint_clears` in the attention builder and the end of `warmup_per_req_cache`) come with #2635. Whichever of the two merges second, Qwen3.8 clears its released units in the right pool. Tests: `reset_entries` on fp8 / fp32 / bf16 / uint8 / int8 fields, zero and non-zero fills, runs on the CPU runner; the builder-level test (empty `kv_pools`, QSA arena reset) needs AITER like the other builder tests here. Co-authored-by: Cursor <cursoragent@cursor.com>
|
Pushed With #2635 also merged, Qwen3.8 with |
|
GPU validation of 484dc30 (MI355X, The warmup crash is fixed. I started servers with GSM8K, full set run twice: 0.9773 / 0.978. But pass 2 barely hit the CPU tier (offload_hit 0.0002), because a cyclic full-set sweep evicts the blocks the next pass needs. So I also ran l2's split method: questions 0-439, pass 1 stores, pass 2 retrieves from CPU. With #2635 merged, retrieval from CPU regresses:
In every regressed run, all 440 pass-2 requests are served from CPU (~480k tokens, 425-431 state-tier joints, 0 load failures), so the counters look healthy. The wrong answers solve a different problem: doc 199 (answer 7500) comes back as "Tomatoes: 2 × $5 = $10, Apples: 5 × $6 = $30, Total: $40". The empty ones run to the 12288-token limit. That looks like restored state that doesn't belong to the prompt. Turning off the QSA clear restores accuracy twice, so the clear added by #2635 + 484dc30 is the cause. One ordering hazard, measured. The Runs, logs and the diagnostic patches (local branches |
…load fills them `BlockPool.take_released_units` assumed the batch whose allocation reuses a given-back unit is the one that drains it, so the forward's `build` clears it before anything writes it. True for the forward, not for a connector load: `EngineCore` dispatches the connector metadata (`start_load_kv`) ahead of the forward, and the load fills its blocks from the connector's own threads and streams, ordered after nothing on the compute stream. A unit is free again the moment it is given back, and under pool pressure a load's blocks are exactly that: allocating them evicts a checkpoint, and its units come straight back out. Two ways the clear then lands after the load: * A pass with a forward: the load and the `build`-time clear race. * A pass with no forward -- every request parked on a load, which is how the second pass of an offload run starts -- builds an empty batch that drains nothing, yet its metadata still goes out and the loads start. The units are cleared by the next non-empty batch, after the loads finished: the retrieved KV is zeroed. `_dispatch_idle_offload_work` and PP's `_dispatch_connector_only_batch` dispatch loads with no batch at all. With #2636 (Qwen3.8 LMCache offload, whose QSA arena now clears too) requests fetched back from CPU lost part of their context: GSM8K on the fetched pass fell to 0.93-0.95 from 0.977-0.984, with 3-19 empty / truncated answers, independent of MTP and of the CPU pool size. Qwen3.5's LMCache offload on main has the same path. Clear before loads instead: whenever connector metadata with work is about to be dispatched, the scheduler moves every unit given back so far onto it (`ConnectorMetadata.checkpoint_clears`, `Scheduler.carry_checkpoint_clears`), at both scheduling branches (an empty pass included), the idle dispatch and PP's connector-only dispatch. The worker zeroes them on the compute stream and waits for every clear issued so far, a forward's included, before `start_load_kv` (`clear_released_units` / `wait_released_unit_clears`). The wait costs nothing once the clears have run, which is the usual case since the engine already waited for the previous forward's output. Metadata with no work is left alone: the engine may drop it, and with no load in it the batch clears those units as before. Units are not held back from allocation instead: KV allocation evicts a checkpoint and reuses its units in the same pass (`_ensure_page_units`), which `can_allocate` already counts on. Tests: the scheduler carries the units on the metadata of a pass where every request is parked on a load (no forward) and of a pass with one, and leaves them to the batch when the metadata has no work; the idle and PP dispatches carry them; the worker clears and waits before `start_load_kv` (needs AITER); the wait waits for the device (needs a GPU). Co-authored-by: Cursor <cursoragent@cursor.com>
Fetched-pass regression with #2635: root cause and fix are in #2635, nothing to change hereFollow-up to the earlier report: GSM8K on the CPU-fetched pass fell to 0.93–0.95 when this branch was combined with #2635. Cause. #2635's clear-on-release zeroed released PAGE units after a connector load had already filled them with retrieved KV. This happened in two ways:
The same thing happens to Qwen3.5 on main's LMCache path without this PR. Details are in the #2635 comment. This PR is fine. The QSA arena's Combined tree:
|
|
@valarLip this one is ready for you. @yhl-amd approved it at Changes today:
CI: Pre Checkin is green. Three Accuracy jobs failed in run 40290659165: DeepSeek-V4.1-Flash DSpark, Kimi-K3 and GLM-5.3-Flash kpool. All three ran on runner |
valarLip
left a comment
valarLip
left a comment
There was a problem hiding this comment.
Reviewed at 484dc30 (read only, no GPU repro). The PAGE-image design and the #2635 interaction write-up are clear. Two things look worth fixing before merge:
- PD disaggregation now starts but silently drops GDN/PLE state;
- the engine-id port fix does not take effect with the default launch.
Details inline, then smaller points.
One test gap outside the hunks: KimiK3OffloadConnector._withholds_last_prompt_token now returns False for every has_per_req_cache request. That also changes Kimi-K3 and Qwen3-Next/3.5 full-prompt hits (need becomes num_prompt - frontier, and correctness then rests on _decide_load_after_alloc loading only at the joint boundary), but the new tests only exercise Qwen3.8.
| raise ValueError("Qwen QSA/PLE metadata do not support TBO slicing") | ||
| if config.kv_transfer_config: | ||
| if config.kv_transfer_config and not _qwen38_offload_ready(config): | ||
| # PP and RapidServe stay on fork slots. Restoring QSA pages there |
There was a problem hiding this comment.
This used to refuse any kv_transfer_config; now it refuses only PP / RapidServe and does not look at the connector. At the same time get_kv_transfer_tensors no longer publishes GDN / PLE as slot_regions. So a PD connector such as {"kv_connector": "mooncake", "kv_role": "kv_producer"} starts fine. mooncake_connector's _has_slot_regions is False, so only QSA pages move, and the decode side continues from zero GDN/PLE state: fluent but wrong output, no error. kimi_mla_gdn_attn.py refuses the same situation through resolve_pd_backend. The same refusal here (offload connectors allowed, PD connectors rejected until the state rides along) would keep this fail-stop.
There was a problem hiding this comment.
Fixed in 12dea23. Qwen4ExpBackend.validate_config now runs resolve_pd_backend(config.kv_transfer_config), as kimi_mla_gdn_attn.py does, and refuses at startup when it names a P/D backend: "Qwen3.8 disaggregated serving is not supported: kv_connector='mooncake' would transfer the QSA pages without the GDN and PLE state ...". Offload connectors (lmcache_offload, lmcache_mp, and multi wrapping only offload leaves) still pass.
Tests (tests/test_qwen4_exp_offload.py): test_validate_config_rejects_pd_connectors (mooncake producer, moriio consumer, no connector named = moriio default, and multi wrapping mooncake + lmcache_offload) and test_validate_config_allows_offload_connectors.
| ) | ||
| or 0 | ||
| ) | ||
| port = getattr(config, "port", None) |
There was a problem hiding this comment.
config.port is the engine's --port (default 8006). model_runner only writes it into MASTER_PORT, which nothing reads; the actual rendezvous uses data_parallel_base_port (model_runner.py:1020). The usual launch (scripts/start_atom_server.sh) passes only --server-port, so two TP1 servers on one host are both atom-offload-p8006-dp0, and the lookup socket still gets taken over: the boundaries=0 / halved-throughput case this is meant to fix. The docstring's "rendezvous port (MASTER_PORT), which no two servers on a host can share" does not hold either. Something genuinely unique per server would fix it, e.g. the server's HTTP port, the resolved data_parallel_base_port, or the parent PID. (Nit: isinstance(port, int) also accepts True.)
There was a problem hiding this comment.
Agreed, config.port only feeds MASTER_PORT. Fixed in 409c46c with the resolved data_parallel_base_port, of the three options the only one that is both unique per live server and identical in every process of one server:
CoreManager._start_distributed_storebinds the runners' TCPStore on it before spawning any engine (0 asks the OS for a free port), holds it for the server's lifetime, and writes the bound port back into the config every engine and worker is spawned with.- The parent PID differs between the EngineCore (scheduler) and the model-runner (worker) processes, so it would split one replica; the HTTP port is not in
Configat all.
The id is now atom-offload-r{port}-dp{rank}; an unresolved port (0, before CoreManager) falls back to atom-offload-dp{rank}. The check is type(port) is int and port > 0, so True and strings are rejected. CacheEngineKey (model_name, world_size, worker_id, chunk_hash, dtype, tags) does not contain the engine id, so a new id per restart does not invalidate stored chunks. Docstring and offload/README.md updated.
Tests (tests/test_lmcache_offload_connector.py): test_two_default_tp1_servers_on_one_host_get_distinct_engine_ids runs two real CoreManager._start_distributed_store calls with the default ParallelConfig() (base port 0) and both config.port=8006, and checks the two ids differ while the deep copy an engine is spawned with keeps its server's id; the updated test_engine_id_with_port_still_pairs_one_replica_and_splits_replicas keeps PP/TP ranks of one replica together; test_engine_id_ignores_an_unresolved_or_non_int_port covers 0 / None / True / "41000". With the previous code these fail.
| # stop on a chunk strictly below the prompt, so skipping to the last | ||
| # checkpoint leaves that chunk with no image. | ||
| entry = self._save_tracker.get(str(seq.id)) | ||
| saved = int(entry[1]) if entry is not None and entry[0] is seq else 0 |
There was a problem hiding this comment.
Saving every image above saved in order, one save in flight per request, costs more than the last-image-only version:
- A 128K prompt at interval 8192 now issues ~16 serial image saves instead of 1, roughly 16x the state bytes written to the CPU tier.
_has_pending_savestays true for the whole sequence, soshould_defer_freeholds the block table for it.- The prompt-end anchor, the resume point nearly every continuation uses, goes last, and it is only pinned when its save is issued, so it can be evicted while queued.
I see the reason in the comment (a load must stop below the prompt). Would saving the anchor first and backfilling the lower rungs afterwards keep that while protecting the most valuable image?
There was a problem hiding this comment.
Changed in 4b52c41: the anchor goes first, and the lower rungs are not backfilled.
_next_image picks the highest image at or below chunk_floor(num_prompt - 1), the highest boundary a resend of this prompt can load (the anchor, already on the chunk grid via align_checkpoint_anchor); only when there is none above the watermark, the highest at or below the frontier (a rung at the prompt end, for a continuation). So a prompt prefilled in one step stores at most two images (anchor, then prompt-end rung) instead of ~16, should_defer_free is released after those, and the anchor is pinned by acquire_checkpoint_source when the first save is issued. Mid-prefill it is the highest computed image, which is what main did before the ordered stepping. _late_save_frontier uses the same order.
I did not add the backfill. A backfilled image sits below the watermark, so it needs a second kind of save that does not move it, and the save path assumes one monotonic line per request in several places: _complete_native_save rolls the watermark back to saved_before on failure, the TP-wide quiescent fence and _retry_or_retire_failed_save restore _save_previous_frontier, the early-release block leases, and the late-save source. Each would need a branch, and the failure paths are where a mistake reads recycled KV. Without it, the skipped rungs are what main skipped too: a request that diverges below the anchor resumes from the highest image stored during chunked prefill, or recomputes. The trade-off is in the _next_image docstring. Happy to add a capped backfill as a follow-up if you think it is worth it.
Tests (tests/test_lmcache_mp_native_scheduler.py): test_the_prompt_end_anchor_is_stored_first_and_pinned (rungs at 8/16/24 for a 24-token prompt: 16 first and pinned, then 24, then nothing pending), test_a_rung_at_the_prompt_end_is_stored_when_no_anchor_formed, test_mid_prefill_saves_the_highest_image_it_has_computed, and test_save_frontier_rescans_only_when_frontier_or_checkpoints_change back to the highest-image expectations. The first and last fail with the previous stepping.
On an MI355X (Qwen3.8 MXFP4, TP1, MTP 3, lmcache_mp with L1 40 GiB, GSM8K first 200, ATOM restarted between the passes with LMCache kept): 0.98 / 0.98, and all 200 pass-2 requests Retrieved from LMCache.
| PLE_STATE_SLOT_CLASS = STATE_SLOT_CLASS | ||
|
|
||
|
|
||
| def _qwen38_offload_ready(config) -> bool: |
There was a problem hiding this comment.
This does not look at kv_transfer_config, and _uses_paged_checkpoints follows it. So every Qwen3.8 deployment with PP=1 and no RapidServe, including servers with no offload at all, moves from fork slots to PAGE checkpoints: GDN+PLE images copied into QSA KV blocks, with a different KV budget, hit behaviour and copy path. The PR description mentions it, but the recipe only covers offload, and it has been validated on TP1 bf16 only. If that is intended, a name like _qwen38_page_checkpoints_supported would say what it decides. If not, gating it on an offload connector keeps the plain server unchanged.
There was a problem hiding this comment.
Option (a), in bb734a4: _uses_paged_checkpoints now also requires _qwen38_offload_configured(config), i.e. a leaf of kv_transfer_config that KVConnectorFactory.is_offload_backend accepts. A plain Qwen3.8 server (and any P/D config, which is refused anyway) keeps fork slots as on main. The recipe says PAGE checkpoints come with an offload connector only.
Test: test_page_checkpoints_follow_an_offload_connector (no config / {} → fork; lmcache_offload, lmcache_mp, multi with an offload leaf → PAGE).
| # Qwen3.8 QSA. Same rule: new tails, never a reorder. Absent | ||
| # tensors stay out, so a model that does not set them keeps | ||
| # the byte layout it already stored. | ||
| getattr(kvt, "index_raw", None), |
There was a problem hiding this comment.
"A model that does not set them keeps the byte layout it already stored" is true across models, but within one model rope_positions exists only when multimodal_config is set (qwen4_exp_attn.py:195). build_page_namespace records neither that switch nor the segment list. So the same checkpoint served with and without image input gets the same namespace but different bytes_per_block. With a shared persistent LMCache (disk or MP server), one reads the other's chunks at the wrong offsets, and the GDN image is overwritten. Adding the segment set (or the multimodal flag) to the namespace document would separate them.
There was a problem hiding this comment.
Fixed in 3dd40a4. build_page_namespace adds image_input: multimodal_config is not None to the namespace document for the model types whose segment set follows it (qwen4_exp, qwen4_exp_text, read through get_hf_text_config). I kept it off other models' documents on purpose: their byte layout does not change with image input, and adding the key unconditionally would change every existing digest.
Tests (tests/test_lmcache_offload_config.py): test_qwen38_page_namespace_separates_image_input, test_image_input_leaves_other_page_namespaces_unchanged.
| ngram = self.ngram_context_len + self.num_spec | ||
| dtype = self.model_runner.config.torch_dtype | ||
| return ( | ||
| f":ple_conv=({state_len},{channels}),{dtype}" |
There was a problem hiding this comment.
The PLE conv state is allocated (slots, channels, state_len) (allocate_per_req_cache), but the fragment records (state_len, channels). Byte counts are unaffected, but this string is part of the persistent layout_id. Correcting it later invalidates every stored image, and a reader that parses it gets the transposed shape. Fixing it now, before anything is stored, is free.
There was a problem hiding this comment.
Fixed in d159907: the fragment is now :ple_conv=(channels,state_len),dtype, matching allocate_per_req_cache. Covered by the layout-id assertion in test_layout_id_names_ple_and_forks_under_pipeline_parallel (":ple_conv=(16,10),torch.bfloat16:").
| raise RuntimeError("Qwen3.8 PAGE checkpoint has no QSA arena") | ||
| return [QsaRegionPool(arena)] | ||
|
|
||
| def _page_unit_regions(self) -> tuple[np.ndarray, np.ndarray]: |
There was a problem hiding this comment.
_page_unit_regions, page_unit_views and page_unit_segments (and checkpoint_page_unit_bytes, get_kv_transfer_tensors's tail) are near line-for-line copies of GDNAttentionMetadataBuilder (gdn_attn.py:1677-1828); the only difference is self._qsa_pools() in place of self.kv_pools.values(). A _page_unit_pools() hook on the parent would let this class inherit them and keep the PAGE contract in one place.
It would also fix a cost: _qsa_pools() builds a new QsaRegionPool (fields × layers as_strided views) on every call, only to compute the cache key. page_unit_views / page_unit_segments build it twice per call, on every checkpoint-copy step and every store-thread segment, all under the GIL. Building it once in allocate_kv_cache_tensors and dropping it in release_kv_pools avoids that.
There was a problem hiding this comment.
Done in 8d6e5b6. GDNAttentionMetadataBuilder gets _page_unit_pools() (its MHA pools) and _published_page_unit_bytes(), and the tail of get_kv_transfer_tensors moved into _with_paged_state_contract. Qwen4ExpMetadataBuilder no longer defines _page_unit_regions / page_unit_views / page_unit_segments / checkpoint_page_unit_bytes; it overrides only the hooks. The QsaRegionPool is built once in allocate_kv_cache_tensors and dropped in release_kv_pools.
Tests (tests/test_qwen4_exp_offload.py): test_qsa_builder_inherits_the_page_contract_through_hooks, test_qsa_region_pool_is_built_once_per_allocation (no construction across views / segments / transfer tensors; None after release). These are AITER-gated and pass on an MI355X with the GPU unit-test set.
| ) | ||
|
|
||
| regions = self._page_unit_regions() | ||
| spec = self.model_runner.state_runtime.checkpoint_spec |
There was a problem hiding this comment.
page_unit_views guards state_runtime / checkpoint_spec being None; this dereferences both directly. On the store thread before they are set, this raises AttributeError instead of taking the same safe path.
There was a problem hiding this comment.
Fixed in 8d6e5b6: both page_unit_views and page_unit_segments read the image size through one _image_bytes() that returns 0 (untrimmed) when state_runtime or its checkpoint_spec is not set, so the store thread takes the same path. Tests: test_page_unit_segments_take_the_untrimmed_path_before_the_spec, test_page_unit_segments_trim_to_the_image.
| (page.region.base_addr, page.region.unit_bytes) | ||
| for page in transfer.pages[:count] | ||
| ] | ||
| if published != list(zip(bases.tolist(), sizes.tolist(), strict=True)): |
There was a problem hiding this comment.
Both sides of this check come from the same _qsa_pools() region tensors (published from region_tensors(), bases/sizes from mha_page_unit_regions over the same pools), so it cannot fail for the mismatch it describes. A change in field order, or rope_positions appearing, would move both together. Comparing against something independent, such as the arena's own byte offsets or spec.page_unit_bytes per field, would let it catch drift.
There was a problem hiding this comment.
Fixed in 274be1f. _with_paged_state_contract now compares the published (base_addr, unit_bytes) with LayerMajorArena.region_layout(), which is byte arithmetic on the arena's declared field offsets and never goes through view() / QsaRegionPool, and requires the regions to sum to spec.page_unit_bytes. The MHA path uses the same check through _page_unit_arenas(); it had the same self-comparison.
Tests: test_transfer_contract_matches_the_arena_byte_layout, test_transfer_contract_catches_published_regions_drifting (k.0 and v.0 swapped in the pool: fails with the old check, caught by the new one), test_transfer_contract_catches_a_missing_field (a dropped field is already caught by the unit-size check), test_arena_region_layout_matches_its_views. The builder tests need AITER, so they skip on the CPU runner; I ran them on an MI355X (d5a9cfe adjusts two of them to the real build), and the swap test fails there with the previous self-comparison.
| "qwen3_5_moe_text", | ||
| # Qwen3.8 publishes its QSA arena as the PAGE image and appends PLE | ||
| # to the same GDN checkpoint. The tier is still `qwen`. | ||
| "qwen4_exp", |
There was a problem hiding this comment.
The header comment above this set still says it is the "Same set as attn_family resolves to Family.GDN", but qwen4_exp(_text) resolves to Family.QSA_GDN. Similarly, offload/README.md and hybrid/qwen/connector.py's module docstring still describe the qwen tier as an MHA page image, and the lmcache_engine_id summary still says it is named per DP rank.
There was a problem hiding this comment.
Fixed in 50f21df: the _GDN_LINEAR_MODEL_TYPES header now says Qwen3-Next/3.5 resolve to Family.GDN and Qwen3.8 to Family.QSA_GDN; offload/README.md and the hybrid/qwen/connector.py docstring describe the QSA arena image (K, V, raw and compressed index keys, rope_positions with image input) plus the PLE windows; the README's engine-id note says per server and DP rank. The lmcache_engine_id docstring was rewritten with the port change (409c46c).
The QSA PAGE image carries GDN and PLE only for the offload tiers. A P/D connector such as mooncake moves the QSA pages alone, so decode continued from a zero recurrent state with no error. Refuse every P/D backend resolve_pd_backend names at startup, as kimi_mla_gdn_attn does; offload connectors (lmcache_offload, lmcache_mp) still start. Co-authored-by: Cursor <cursoragent@cursor.com>
config.port is --port (default 8006); only MASTER_PORT, which nothing reads, comes from it, and the usual launch passes --server-port alone. Two TP1 servers on one host were both atom-offload-p8006-dp0 and still shared one lookup socket. Use parallel_config.data_parallel_base_port instead: CoreManager binds the runners' TCPStore on it before spawning (port 0 asks the OS), holds it for the server's lifetime and writes it into the config every engine and worker is copied from. Reject bools with type(...) is int. Co-authored-by: Cursor <cursoragent@cursor.com>
…them _uses_paged_checkpoints ignored kv_transfer_config, so every Qwen3.8 server with PP=1 and no RapidServe, including one with no offload at all, moved from fork slots to PAGE checkpoints. Gate it on a configured offload connector (lmcache_offload, lmcache_mp, or one inside multi) so a plain server keeps main's fork-slot behaviour, and say so in the recipe. Co-authored-by: Cursor <cursoragent@cursor.com>
…the layout id allocate_per_req_cache lays the PLE conv state out (slots, channels, state_len), but the layout_id fragment said (state_len, channels). Byte counts were unaffected, but the string is part of the persistent PAGE layout id; fix it before anything is stored under it. Co-authored-by: Cursor <cursoragent@cursor.com>
rope_positions joins the QSA arena only when multimodal_config is set, so one checkpoint served with and without image input stored two byte layouts under one namespace; on a shared persistent tier one read the other's chunks at the wrong offsets. Add image_input to the namespace document for the model types whose segment set depends on it, and only for them, so every other namespace is unchanged. Co-authored-by: Cursor <cursoragent@cursor.com>
_page_unit_regions, page_unit_views, page_unit_segments, checkpoint_page_unit_bytes and the tail of get_kv_transfer_tensors were near line-for-line copies of GDNAttentionMetadataBuilder's. Give the parent a _page_unit_pools() hook (plus _published_page_unit_bytes for sizing before allocation) and move the tail into _with_paged_state_contract; the QSA builder overrides only the hooks. Build the QsaRegionPool once in allocate_kv_cache_tensors and drop it in release_kv_pools, instead of fields x layers as_strided views on every copy step and store-thread segment. page_unit_segments now takes the same untrimmed path as page_unit_views when state_runtime or its checkpoint_spec is not set yet, instead of raising AttributeError on the store thread. Co-authored-by: Cursor <cursoragent@cursor.com>
The consistency check in get_kv_transfer_tensors compared the published regions with _page_unit_regions, and both read the same pool views, so a field reorder or a field appearing moved both together and it could never fail. Compare them with LayerMajorArena.region_layout() instead, byte arithmetic on the declared field offsets, and require the regions to sum to spec.page_unit_bytes. Both the MHA pools and the QSA arena answer through a _page_unit_arenas() hook. Co-authored-by: Cursor <cursoragent@cursor.com>
…engine id The _GDN_LINEAR_MODEL_TYPES header, offload/README.md and the hybrid/qwen/connector.py docstring still described the qwen tier as an MHA page image matching Family.GDN; Qwen3.8 resolves to Family.QSA_GDN and publishes its QSA arena plus PLE windows. The README's engine_id note now says the id is per server and DP rank. Co-authored-by: Cursor <cursoragent@cursor.com>
valarLip
left a comment
valarLip
left a comment
There was a problem hiding this comment.
Follow-up on the current head 32848851f (my earlier review was on 9862cc96a). Thanks for the updates — both of the earlier top issues are handled: draft-row-space modules now get their own *_draft index rows, and quantized QSA caches under lmcache_offload are now explicitly refused instead of losing their scale planes. I found no confirmed correctness bug in this round. CPU unit tests for the touched files pass locally (944 passed, 16 skipped); black is clean, and ruff's two remaining hits are on unchanged lines in factory.py.
Main points this round:
- The new refusal (and the recipe) tell users to pass
--kv-cache-dtype, which isn't a registered flag (qwen4_exp_attn.py:149, recipe L91). - Two fixes are only half-applied: multimodal tier keys miss the DSV4/V4.1 scheduler (
_offload_common.py:75), and the whole-chunk give-back lives on the Qwen scheduler while its K3 parent keeps the old behaviour (qwen/connector.py:33). - Four model-independent changes (engine-id namespacing, multimodal key seeding, chunk-grid anchors, whole-chunk give-back) should be split into their own PRs (
offload/config.py:463). kv_transfernow imports a model's attention backend, which is also why the Qwen3.8 offload tests never run in CI (offload/config.py:388,test_qwen4_exp_offload.py:10).
Please remove (marked inline): the second _record_checkpoint_end (block_manager.py:1586), :ple_order (qwen4_exp_attn.py:545), the duplicate region_layout / published_bytes_for (entry_arena.py:593), page_image_readers() (factory.py:189), and the type(port) is int check.
Not inline (file outside the diff): tests/test_fp4_indexer_kv_transfer_gate.py:233's finally pops four registries but not the new _reads_page_images, while test_kv_connector_factory_topology:72 asserts page_image_readers() is exactly two names — the factory's global registry leaks between tests. Harmless today (the stubs register False), but a future True stub would make the topology test order-dependent. Add an unregister to the factory or use monkeypatch.setitem.
| "quantized QSA cache yet (kv_cache_dtype=" | ||
| f"{getattr(config, 'kv_cache_dtype', None)!r}, index_cache_dtype=" | ||
| f"{getattr(config, 'index_cache_dtype', None)!r}). Use " | ||
| "--kv-cache-dtype bf16 --index-cache-dtype bf16, or lmcache_mp, " |
There was a problem hiding this comment.
This error (and the recipe at recipes/Qwen3.8-Flash-Next.md:91) tells users to pass --kv-cache-dtype bf16 --index-cache-dtype bf16, but arg_utils.py:307 registers only --kv_cache_dtype — argparse rejects the hyphenated form as unknown, so following the message verbatim fails to start. The same recipe uses --kv_cache_dtype at L36/L39. Please use the registered spelling in both places (or register the hyphenated alias).
There was a problem hiding this comment.
Resolved by #2655 (folded in at 97af4f1b7): this refusal is gone, and qwen4_exp_attn.py and the recipe now spell --kv_cache_dtype. Correction to the premise, for the record: the hyphenated form does parse. openai_server uses FlexibleArgumentParser (#1686), which registers the dash/underscore counterpart of every long flag, so --kv-cache-dtype fp8 gives kv_cache_dtype='fp8' (checked on main ab4ada40a). The rename was harmless, but it was not needed for startup. An earlier version of this reply pointed to #3024; I closed that issue for the same reason.
| parallelism, DP attention, and TBO are not supported. | ||
| - LMCache offload (`lmcache_offload`, `lmcache_mp`) stores the QSA pages plus | ||
| a checkpoint of the GDN and PLE state. Until ROCm/ATOM#2655, | ||
| `lmcache_offload` needs `--kv-cache-dtype bf16 --index-cache-dtype bf16`. |
There was a problem hiding this comment.
--kv-cache-dtype isn't a registered flag (only --kv_cache_dtype, arg_utils.py:307), and L36/L39 of this recipe use the underscore form — see the comment on qwen4_exp_attn.py:149.
There was a problem hiding this comment.
| _MEDIA_KEY_PERSON = b"atom-lmc-media" | ||
|
|
||
|
|
||
| def tier_token_ids(seq, end: int) -> list[int]: |
There was a problem hiding this comment.
The multimodal-key fix (tier_token_ids) misses the DSV4 scheduler: hybrid/dsv4/connector.py still builds lookup / save / load keys from raw list(seq.token_ids[:n]) at L1971, L2374 and L2441. DeepSeek-V4.1 is multimodal and has compress_ratios, so it takes the hybrid layout; under lmcache_offload, two requests with identical text but different images would hit each other's LMCache KV while the HBM-side state chain already carries cache_seed, so KV and state disagree. The description lists it under "Not changed", but per Fix-then-sweep it belongs in the same change. (I haven't run V4.1 to confirm it reaches this path.)
There was a problem hiding this comment.
Done in #3027, which now holds the multimodal-key fix. DSV4's lookup, load and save all key through tier_token_ids (hybrid/dsv4/connector.py).
I swept the SLOT boundary chain too (_chained_prefix_hashes(seed=…)); it's the same bug, per Fix-then-sweep. That chain started from −1, while BlockManager's chain starts from cache_seed, so two image prompts with the same text still shared a SLOT boundary hash.
Tests: test_dsv4_lookup_is_keyed_the_same_way and test_a_dsv4_save_keys_its_pages_and_slot_by_the_media. Reverting the lookup, the save or the seed each fails one of them.
| sends state stores and loads here. | ||
| """ | ||
|
|
||
| def _withholds_last_prompt_token(self, seq) -> bool: |
There was a problem hiding this comment.
Nothing in _withholds_last_prompt_token is Qwen-specific — it only reads has_per_req_cache, _block_manager and joint_boundary_tokens — yet it lives on QwenOffloadScheduler, while its parent KimiK3OffloadScheduler (kimi_k3/connector.py:590, same _decide_load_after_alloc at L783) keeps the base behaviour. So Kimi-K3 under lmcache_offload with a chunk-aligned prompt still steps a full-prompt hit back a whole chunk and recomputes it. Please move this override up to KimiK3OffloadScheduler (or the state mixin) — it already gates on has_per_req_cache.
There was a problem hiding this comment.
Done in #3027: the override is now KimiK3OffloadScheduler._withholds_last_prompt_token, so Kimi-K3 and its Qwen GDN subclasses all get it, and qwen/connector.py no longer has it. The generic default (always withhold) stays in _offload_common.
test_a_hybrid_full_prompt_hit_keeps_its_last_chunk runs for kimi-k3, qwen3-next and qwen3.5 there, and this PR adds the qwen3.8 row. The Kimi-K3 GPU check (the same chunk-aligned prompt twice, then GSM8K) is listed as pending in #3027's description.
| del seq | ||
| return True | ||
|
|
||
| def _loadable_hit(self, hit: int, num_prompt: int, seq=None) -> int: |
There was a problem hiding this comment.
_loadable_hit(hit, num_prompt, seq=None) now has two modes depending on whether the caller passes seq: the base hook is del seq; return True and only Qwen overrides it, while dsv4/connector.py:2012 and mooncake_store/scheduler.py:539 don't pass seq and get the old semantics. A future caller can silently fall into the old mode. Make seq required at every call site, or have subclasses override _loadable_hit itself.
There was a problem hiding this comment.
Done in #3027: _loadable_hit(hit, num_prompt, seq) has one signature, seq is required, and all three callers pass it (chunked, dsv4/connector.py and mooncake_store/scheduler.py).
| + self._offsets[field.name] * self.entries // itemsize, | ||
| ) | ||
|
|
||
| def region_layout(self) -> list[tuple[int, int]]: |
There was a problem hiding this comment.
LayerMajorArena.region_tensors / region_layout are a third copy of the same walk as MhaKvPool.region_tensors / region_layout and MlaKvPool.region_tensors — the PR notes they're "not merged", and the history shows a35e872 merged them before 12196f1 reverted it. Three copies will drift; please have MhaKvPool delegate to each arena's region_tensors / region_layout. And published_bytes_for is the same sum as the existing mha_published_page_bytes — please keep only one.
There was a problem hiding this comment.
Done in 18ee7ac7b:
MhaKvPool.region_tensorsdelegates to each arena'sregion_tensors;region_layoutalready delegated.LayerMajorArenano longer caches its region list. The caller's_page_unit_region_cachealready caches the walk.mha_published_page_bytesis removed, andpublished_bytes_foris the one sum.GDNAttentionMetadataBuilder._published_page_unit_bytesadds it up over the pool's field groups.
MlaKvPool.region_tensors stays its own walk, because the MLA pool doesn't use LayerMajorArena.
| return cfg | ||
|
|
||
|
|
||
| def lmcache_engine_id(config) -> str: |
There was a problem hiding this comment.
Please split out the model-independent changes. Four changes here aren't specific to Qwen3.8 and affect every model: per-port lmcache_engine_id namespacing (this function), tier_token_ids seeding multimodal keys, native-state MP anchors aligned to the chunk grid, and the whole-chunk give-back. They change the shared scheduling core (block_manager) and every LMCache user's behaviour, buried in a 1912-line model PR, so they can't be reviewed or reverted on their own. Please move them into separate PRs, each with its own validation. Also please remove the type(port) is int check and its comment about bool.
There was a problem hiding this comment.
Split out as #3027: per-port engine id, multimodal keys, MP anchors on the chunk grid, and the whole-chunk give-back, with their tests and docs. This PR's base is now #3027's branch, so the diff here is Qwen3.8 only. Merge order: #3027 → this PR → #2668.
The four are in one PR rather than four, because they share the scheduler files and test modules. Each has its own section in #3027's description, with its own validation: CPU results now, and a GPU check per change listed as pending. I won't ask for a merge until those are in. I'm happy to split it further if you prefer.
The type(port) is int check and its comment are removed. The port is now int(... or 0) and used only if port > 0 (atom/kv_transfer/offload/config.py:476-488 in #3027). valar-cop had asked for that check to be restored in r4439120010, so cc @valar-cop: it's removed per this review.
| return cls._reads_page_images.get(name, False) | ||
|
|
||
| @classmethod | ||
| def page_image_readers(cls) -> list[str]: |
There was a problem hiding this comment.
Please remove page_image_readers(). In production it appears only in one error message, and after _check_kv_transfer passes, _wants_paged_checkpoints' any(reads_page_images) is equivalent to bool(leaves) — a new public API plus two redundant checks (YAGNI). The message can name lmcache_offload / lmcache_mp directly, as the preceding error already does.
There was a problem hiding this comment.
Removed in 18ee7ac7b, together with its assertion in test_kv_connector_factory_topology.py. _wants_paged_checkpoints is bool(kv_transfer_leaves(config)), and the refusal names lmcache_offload / lmcache_mp directly. reads_page_images(name) remains, and check_kv_transfer uses it. The registry leak from the review body is fixed too: both finally blocks in test_fp4_indexer_kv_transfer_gate.py now pop _reads_page_images.
| if published <= 0: | ||
| raise RuntimeError("GDN PAGE checkpoint has no published MHA bytes") | ||
| return published | ||
| return mha_published_page_bytes(self._declare_kv_pools().values()) |
There was a problem hiding this comment.
Name-matches-function: mha_page_unit_regions / views / segments / mha_published_page_bytes, the _page_unit_regions docstring ("in the MHA pool") and the error "no published MHA bytes" now all apply equally to Qwen3.8's QSA arena. The description acknowledges this but leaves it; CLAUDE.md asks to rename when behaviour changes, otherwise a QSA layout error reports MHA and misleads.
There was a problem hiding this comment.
Done in 18ee7ac7b: the module is now pool_layout/layer_major_page_unit_geometry.py, and its functions are layer_major_page_unit_regions / _bases / _stream_sizes / _segments / _views. They serve both layer-major pools, MhaKvPool and the QSA LayerMajorArena. The errors now say "PAGE region …", "no PAGE regions", "no published PAGE bytes" and "the published regions cover". _page_unit_regions names the pools _page_unit_pools returns. In acbd76533, checkpoint_page_unit_bytes' first line says "PAGE block" instead of "MHA page".
| def test_transfer_contract_catches_published_regions_drifting(monkeypatch): | ||
| builder, cls = _paged_builder() | ||
| builder.qsa_arena.region_tensors() | ||
| rows = builder.qsa_arena._regions |
There was a problem hiding this comment.
Tests mutate private internals — swapping/assigning arena._regions (here), clearing _page_unit_region_cache — and several test cache mechanics like "built only once", so any refactor breaks them. In test_entry_arena, half of the region_tensors assertions are the same expression as the implementation, so they'd pass even if it were wrong. Please test observable behaviour (addresses/bytes against an independent expectation) instead.
There was a problem hiding this comment.
Done in 18ee7ac7b. Nothing touches arena._regions or _page_unit_region_cache any more, and the "built once" and walk-count tests are gone. They're replaced by tests of observable results:
test_page_unit_regions_follow_a_reallocated_arena: after reallocating with another block count, the bases and sizes equal hand-computed offsets from the buffer start.test_transfer_contract_catches_published_regions_drifting/_a_missing_field: these hand a swapped or truncated transfer to_with_paged_state_contract.test_a_page_unit_the_arena_does_not_cover_is_refused: the spec's unit is larger than the arena.
test_entry_arena's region test now compares against literal addresses and sizes computed from the field shapes, not against the implementation's own expression.
* perf(qwen4-exp): stack index heads into M for the QSA decode scorer
Decode-sized QSA index scoring (the target's MTP verify and the draft, at
most 256 rows) ran on the CUDA-core kernel, which reads every request's
compressed keys once per query row. The new Triton MFMA kernel stacks 4 rows
x 4 index heads into one M = 16 operand, so a request's keys are read once
per verify and meet every head in a single dot; the per-head ReLU is
applied before the heads are summed back out of M.
On Qwen3.8-Flash-Next (MI355X, TP1, MTP 3) the target's decode scoring
drops from 1130 to 320 us a step and the draft's from 147 to 85 us, -6.2%
of decode kernel time. ATOM_QSA_STACKED_DECODE_SCORER=0 keeps the old
kernel.
Co-authored-by: Cursor <cursoragent@cursor.com>
* feat(qwen4-exp): FP8 KV and FP8 / FP4 index caches for QSA
Qwen QSA refused anything but a BF16 KV cache and ignored
--index-cache-dtype. Serve --kv_cache_dtype fp8 and --index-cache-dtype
{bf16,fp8,fp4} (unset follows the KV dtype, as Config resolves it):
- K/V: E4M3 with one FP32 scale per (token, KV head) = row amax / FP8 max,
written by a new Triton kernel (K and V in one launch). Dynamic scales
need no calibration and never saturate. The sparse GQA kernel reads
E4M3 tiles, multiplies the key scale into its score column and folds the
value scale into the probability before PV; FP8 tiles have their own
tuned block config.
- Compressed index keys: FP8 with a per-row FP32 scale, or MXFP4 (packed
E2M1 pairs, one E8M0 exponent per 32 dims, rounded up so a group never
clips). Every scorer -- the stacked decode one, the CUDA-core one and the
prefill MFMA one -- reads them; FP4 decodes with bit operations. The raw
index keys stay BF16: they are mean-pooled before normalization and a
group can straddle a prefill chunk.
- The MTP draft layer keeps BF16 caches in a `qsa_draft` row space of its
own (ATOM_QSA_DRAFT_KV_BF16, default 1): an FP8 draft KV lowers the real
acceptance length by ~3.5%, a BF16 one not at all.
- `_qsa_fields` declares the scale planes, so the arena sizes the pool by
bytes; the memory profile runs on caches of the served format. Plugin
hosts keep binding their own BF16 caches.
On Qwen3.8-Flash-Next (MI355X, TP1, MTP 3) FP8 KV fits 1.64-1.67x the
tokens, cuts prefill sparse GQA by 17%, and leaves GSM8K and the MTP
acceptance length unchanged. With --block-size 64 and a BF16 index cache
the prefill scorer stays on FlyDSL.
Co-authored-by: Cursor <cursoragent@cursor.com>
* fix(qwen4-exp): register the QSA FP8 scale planes for KV offload
build_kv_cache_tensor bound the FP8 K/V scale planes (and the quantized
compressed-index scale) to the attention module but registered k_scale /
v_scale as None, so lmcache_offload's DenseKVByteCodec stored and restored
the FP8 rows without their scales. A prefix fetched back from CPU then
decoded under whatever scales the reused blocks held: GSM8K pass 2 dropped
from 0.98 to 0.50, every CPU-restored request returning empty or running to
max_tokens. MTP draft rows also registered the target's index fields at the
draft row.
Register every paged field the arena declares, with the row space's own
suffix, and move index_compressed_scale as a new codec tail. BF16 caches
register the same tensors as before.
Co-authored-by: Cursor <cursoragent@cursor.com>
* fix(qwen4-exp): move the QSA FP8 scale planes with lmcache_offload
build_kv_cache_tensor bound the FP8 K/V scale planes and the quantized
compressed-index scale to the attention module but did not register them,
so the in-process codec stored and restored FP8 rows without their
scales: GSM8K pass 2 fell from 0.98 to 0.54. Register k_scale / v_scale
and append index_compressed_scale to `extra_planes`, all through the row
space's own suffix (MTP draft rows registered the target's index rows),
and drop #2636's refusal of lmcache_offload with quantized QSA caches.
Co-authored-by: Cursor <cursoragent@cursor.com>
* refactor(qwen4-exp): bind and register one QSA layer's views from one lookup
Look each field up once and hand the same views to bind_caches and
KVCacheTensor. Fold the codec coverage into the existing draft-row test:
every QSA layer, five cache-format cases, each registered tensor matched by
data_ptr against its own row space and row in the stored order.
Co-authored-by: Cursor <cursoragent@cursor.com>
* refactor(qwen38-fp8-kv): keep the diff to the scale planes
Drop the refusal of quantized QSA caches under in-process offload and
hand the codec the k/v scale rows and the index_compressed scale plane
through extra_planes, with a test that every quantized layer gives the
codec the rows it binds.
Co-authored-by: Cursor <cursoragent@cursor.com>
* refactor(qwen38-fp8-kv): look each QSA field up once again
Bind the layer and describe its KVCacheTensor from one set of views, as
the earlier review asked, keep the recipe's P/D line as it was, and drop
_check_kv_transfer from the readers of the QSA field rules.
Co-authored-by: Cursor <cursoragent@cursor.com>
---------
Co-authored-by: Cursor <cursoragent@cursor.com>
| f"{', '.join(KVConnectorFactory.page_image_readers())}; " | ||
| f"kv_connector={name!r} does not read the QSA/GDN/PLE PAGE " | ||
| "checkpoint image." | ||
| ) |
There was a problem hiding this comment.
description is stale now, #2655 is folded in here, please update it and drop the startup-refusal paragraph..
There was a problem hiding this comment.
@valar-cop Updated the description. The "FP8 KV and #2655" section no longer describes a startup refusal. It now says #2655 is folded in (97af4f1b7): the FP8 K/V scale planes are registered as KVCacheTensor.k_scale / v_scale, and the compressed-index scale plane goes in extra_planes, so in-process lmcache_offload carries them. The rest of the description is unchanged. cc @valarLip
There was a problem hiding this comment.
checked, the fp8 section reads right now: no startup-refusal paragraph left, and the code matches the text — build_kv_cache_tensor registers k_scale / v_scale and puts the index planes, index_compressed_scale included, in extra_planes (qwen4_exp_attn.py:637), with the codec appending them after the tails (kv_byte_codec.py:142).
one leftover: "its validation is in #2655" still points at the folded pr, and the validation section has only the bf16 KV run.. the fp8 pass-2 0.9795 / 0.9818 numbers are in your thread reply, might as well put them there.
There was a problem hiding this comment.
@valar-cop Done. The FP8 section now points to Validation, which has an FP8 bullet with the runs closest to this head:
- FP8 KV + BF16 index on
4aff317a5: 0.9818 / 0.9750. - FP8 KV + FP8 index on
de10df694: 0.9886 / 0.9818. Itsbuild_kv_cache_tensoris the same as97af4f1b7's. - Without the scale planes: 0.9795 / 0.5432.
All are pass 1 / pass 2, with pass 2 from the CPU tier. The 0.9795 / 0.9818 in my 10-29 reply was pass 1 / pass 2 of an older head with #2655 applied locally, so I used these instead. Nothing else in the description changed. cc @valarLip
There was a problem hiding this comment.
checked, both fixed. the fp8 section points at Validation now, and the fp8 bullets with the two heads and the no-scale-plane run are there, no startup-refusal paragraph left.
the two fp8 runs aren't on 97af4f1b7 itself but you say build_kv_cache_tensor is identical and the caveat is in the text, so that's fine. nothing else from me on the description.
…ha/qwen38-lmcache Co-authored-by: Cursor <cursoragent@cursor.com> # Conflicts: # atom/kv_transfer/offload/README.md # atom/kv_transfer/offload/_offload_common.py # atom/kv_transfer/offload/config.py # atom/model_engine/block_manager.py # tests/test_gdn_midstep_checkpoints.py # tests/test_lmcache_offload_connector.py # tests/test_state_checkpoint.py
…ames Per valarLip's review of 3284885: - Move the QSA cache rules (formats, draft rows, row space, rope positions, PAGE segments, KV-transfer check) to `qwen4_exp/qsa_cache_rules.py`, which imports neither Triton nor AITER. `offload/config.py` imports it at module level, and the rule tests run on the CPU runner (`tests/test_qsa_cache_rules.py`). - `state_transfer` keeps one construction path (`replace`) with and without PLE; with PLE a midstep-readable transfer is turned off explicitly and logged. Drop `:ple_order` from the layout id. - `MhaKvPool.region_tensors` delegates to its arenas; `LayerMajorArena` drops its region cache; `published_bytes_for` is the one published sum. - Remove `KVConnectorFactory.page_image_readers()`; pop `_reads_page_images` in the fp4 gate test's cleanup. - Rename `mha_page_unit_geometry` to `layer_major_page_unit_geometry` (it addresses the QSA arena too) and drop "MHA" from its errors. - Tests check addresses and bytes against hand-computed offsets instead of mutating `_regions` / `_page_unit_region_cache`. - Restore `qsa_attention.py`'s `envs` import, which main's FlyDSL QSA path (#2653) uses again. Co-authored-by: Cursor <cursoragent@cursor.com>
The layout id's spelling is a pure function of the SSM storage order, the PAGE fields and the PLE planes, so it moves next to the other QSA rules and its tests run on the CPU runner. Also name PAGE blocks, not MHA pages, in checkpoint_page_unit_bytes' docstring: it serves the QSA arena too. Co-authored-by: Cursor <cursoragent@cursor.com>
|
Update Split. The four model-independent changes are now #3027:
This PR's base is now #3027's branch ( Commits
Self-check with the repo's
Validation (CPU only; no GPU available right now)
Pending
Edit: an earlier version of this comment listed #3024 ( |
…ha/qwen38-lmcache Co-authored-by: Cursor <cursoragent@cursor.com>
…ha/qwen38-lmcache
The six function-level imports of pool_layout.entry_arena and pool_layout.layer_major_page_unit_geometry (rewritten by the rename) move to the module's pool_layout import group; neither module imports back. The Qwen connector docstring is rewrapped to 88 columns. Co-authored-by: Cursor <cursoragent@cursor.com>
…ha/qwen38-lmcache
|
Pushed What changed
Self-check with the repo's
Validation
valar-cop's four GPU checks for the model-independent part are answered on #3027: results. Pending: CI running. #3027 still waits on its Kimi-K3 check. Merge order is unchanged: #3027 → this PR → #2668. The description is updated to this head. |
…ha/qwen38-lmcache
|
Pushed What changed
Validation on
Pending: CI running. Merge order is unchanged: #3027 → this PR → #2668. The description is updated to this head. |
|
@ZLkanyo009 checked 753a552: merge looks clean, only that commit came in and the diff vs base is still qwen3.8 only. ci hasn't reported in this run yet, so hold until it's green; a gpu qwen3.5/lmcache_mp pass on 3f6c4a4 before merge would be nice since nothing has run on this head. |
…ha/qwen38-lmcache
|
Pushed What changed
Validation on
Pending: CI running. Merge order is unchanged: #3027 → this PR → #2668. The description is updated to this head. |
|
@ZLkanyo009 merge reads clean here: no conflict markers, the diff against hold for ci. the only thing after your last qwen3.5 run is the MP dispatch-order bound ( |
|
Qwen3.5 + Head The setup is exactly the one in the description's
CI: this PR's Actions runs only readthedocs right now (pass). |
|
@ZLkanyo009 thanks, that covers it. 0.98 / 0.975 with 200/200 hits and 88.5% reuse matches confirmed pre checkin only runs on |
- Shared files (`atom/config.py` `KVCacheTensor.extra_planes`, the factory's `reads_page_images` capability) no longer name Qwen3.8 in their comments: the hooks are generic. - `qsa_draft_rows` reads the cache formats itself; `qsa_caches_quantized` had no other caller. - Comments keep the constraint and drop the measurements (~120 region tensors per call, ~1.4 GB of paged budget). Co-authored-by: Cursor <cursoragent@cursor.com>
…ze check, tests take the real path - gdn_attn _with_paged_state_contract: the byte-sum test repeats the check _page_unit_regions already makes on the same regions. - Qwen4Exp state_transfer: the PLE midstep downgrade and the plain case share one return. - Flydsl GDN tests give the builder an offload-free config instead of patching _uses_paged_checkpoints. - MhaKvPool._arenas() feeds region_tensors and region_layout; named cache bindings instead of an 8-way map unpack; select_policy imported at module top; comments no longer name models. - Tests: reuse _qwen38_config; drop a layout-id test the qsa rules tests already cover. Co-authored-by: Cursor <cursoragent@cursor.com>
|
Update What changed
Size (this PR's 28 files, against #3027's head): source +554/−256 → +529/−257, tests +1010/−48 → +982/−48, docs unchanged. Validation
Not changed
Pending: Pre Checkin runs only on PRs against main, so it starts once #3027 merges and this PR's base becomes main. cc @valarLip |
Motivation
Qwen3.8-Flash-Next (
qwen4_exp_text) refused KV offload, because no existing layout carries its state. Its full-attention pages are a QSA arena (k,v,index_raw,index_compressed, plusrope_positionswith image input and*_draftrows for a BF16 MTP draft under quantized target caches), not the MHA pool Qwen3.5 checkpoints into. Its PLE conv and n-gram windows live as long as the GDN state does, butgdn-paged-state-v1does not include them.This PR reuses the Qwen3.5
qwenstate tier (#2519) instead of adding a codec.Stacked on #3027. The four model-independent offload fixes this PR used to carry are now in #3027, at valarLip's request (review):
lmcache_mpanchors on the chunk grid;#3027 has since gained two more fixes, both found by the GPU runs, and both are merged into this branch:
lmcache_mprestore is timed from when it is enqueued;lmcache_mpload queued behind a finished one starts its scheduler bound again (numbered in dispatch order since753a55282).This PR's base is #3027's branch, so the diff here is Qwen3.8 only. Merge order: #3027 → this PR → #2668.
Changes
Rules without kernels.
atom/model_ops/qwen4_exp/qsa_cache_rules.pyholds what a Qwen3.8 engine keeps and what may move it:It imports neither Triton nor AITER, so
kv_transferimports it at module level, and its tests run on the CPU runner.qwen4_exp_attn,qsa_attentionandops/qsaread it.Layout.
The QSA arena publishes one region per field and layer through
LayerMajorArena.region_tensors().MhaKvPool.region_tensors()now delegates to its arenas' walk. The page geometry module is renamedpool_layout/layer_major_page_unit_geometry.py, because it addresses both pools; its functions and errors no longer say "MHA".With
lmcache_offloadorlmcache_mpconfigured,Qwen4ExpMetadataBuilderaims PAGE checkpoints at the QSA arena and appends the PLE planes after the GDN conv and SSM state. Without them it keeps fork checkpoints, as onmain. The PAGE contract (_page_unit_pools,_published_page_unit_bytes,_with_paged_state_contract) is inherited fromGDNAttentionMetadataBuilder.published_bytes_foris the one published-bytes sum.The layout id (
qsa_paged_layout_id) adds what the parent's shapes cannot express:ssm_storage=kv|vk(FlyDSL keeps the SSM state transposed);The page namespace adds
image_inputanddraft_kv, because the QSA segment set depends on them. Other models' digests are unchanged.With PLE, the transfer is built the same way as without it (
replace()).ATOM_GDN_STATE_MIDSTEP=1does not apply to Qwen3.8, because nothing snapshots the PLE windows inside a forward. So a midstep-readable parent transfer is turned off explicitly, with zero midstep geometry, and a warning says so.KVCacheTensor.extra_planescarries a layer's additional page planes in stored byte order. The dense codec appends them after the existing tails, so older models keep their byte order.Qwen4ExpBackend.validate_config(check_kv_transfer) accepts only tiers that register withreads_page_images=True, which today means the two LMCache tiers. A P/D connector, ormooncake_store, would move the QSA pages without the GDN/PLE state.Behaviour change without offload
None. Configuring a P/D connector or another offload backend for Qwen3.8 now fails at startup, instead of silently transferring KV without the GDN/PLE state.
Behaviour changes for other models
Qwen3-Next and Qwen3.5, on every deployment that keeps PAGE checkpoints (PP=1, no RapidServe; the default), with or without offload:
MhaKvPool.region_layout()) instead of re-deriving them.MhaKvPool.region_tensors()returns the same roles, order and tensors through its arenas.These run at every startup, because
warmup_per_req_cache()andget_kv_transfer_tensors()are called unconditionally. Covered by the GPU unit tests (test_gdn_page_checkpoint.py,test_mha_kv_pool.py) and the Qwen3.5 runs under Validation. Qwen3-Next was last measured on963d7fa0d; see Validation for why.The fixes that change other models' behaviour are described in #3027.
FP8 KV (#2655)
#2655 is folded into this branch (
97af4f1b7).build_kv_cache_tensorregisters the FP8 K/V scale planes asKVCacheTensor.k_scale/v_scaleand appends the compressed-index scale plane toextra_planes, so in-processlmcache_offloadcarries the scales with quantized QSA KV or index caches. BF16 KV andlmcache_mp(which copies whole arenas) are not affected. The FP8 GSM8K runs are under Validation.Tests and CI coverage
These run on the CPU runner:
tests/test_qsa_cache_rules.py: cache formats, draft rows, row space, page segments, the transfer check's refusals, and the layout id spelling;tests/test_lmcache_offload_config.py;tests/test_kv_connector_factory_topology.py.The ~20 builder tests in
tests/test_qwen4_exp_offload.pystill need AITER._qsa_fieldsusesops/qsa.py'sqsa_cache_row_layout, whose module imports Triton and AITER, andqwen4_exp_attnderives fromgdn_attn. Making those kernel-free means changingops/qsa.pyonmain, which is tracked in #3025.Validation
Each number is labeled with the head it ran on. GPU runs are on MI355X,
rocm/atom-dev:nightly_202610221638. The current head isb7ea8c2c0:076142ed5andb7ea8c2c0tidy this PR's own code without changing behavior, and422906a7emerges #3027's1e6fe588d(its own tidy-up, no behavior change). Sinceacbd76533, the branch has merged #3027 up to45123b668:f1662fe4c(the MP restore deadline),4162fcde5(a CPU test),23d9ee232(the MP scheduler bound),753a55282(that bound's restart in dispatch order), and45123b668(#3027's merge of main30d2ad88f, plus its two-servers test building configs fromConfig).09edcad0dmovedgdn_attn.py's in-function imports to module level.The Qwen3.5 runs below are Qwen3.5-35B-A3B-FP8 +
lmcache_mpGSM8K, questions 0-199. ATOM is restarted between the passes with LMCache kept. This is the deployment the "Behaviour changes for other models" section is about.b7ea8c2c0:factory.py:416-417, as on main.affected_tests.pyCPU arm, against fix(offload): key multimodal prompts by media, one LMCache engine id per server, MP anchors on the chunk grid, keep a hybrid's last chunk #3027's head1e6fe588d: 1078 passed, 150 skipped, nothing failing on head.affected_tests.py's 3 GPU-only files (test_vision,test_gdn_midstep_checkpoints,test_mooncake_store_offload_gpu_e2e): 1535 passed, 147 skipped, 0 failed. One test less than on422906a7e(1536): the layout-id test the QSA rules tests already cover.422906a7e: CPU full suite: 11297 passed, 1165 skipped, the same 3 failures as on main; 2 fewer passed than on22d2035ca, the two tests fix(offload): key multimodal prompts by media, one LMCache engine id per server, MP anchors on the chunk grid, keep a hybrid's last chunk #3027 removed. Qwen3.5: 0.965 / 0.975, all 200 requests hit LMCache in pass 2 (88.5% of prompt tokens).22d2035ca: CPU full suite (run_unit_tests.sh): 11299 passed, 1165 skipped, 3 failed, the same 3 as on main (inductor g++).affected_tests.pyCPU arm, against fix(offload): key multimodal prompts by media, one LMCache engine id per server, MP anchors on the chunk grid, keep a hybrid's last chunk #3027's head45123b668: 1078 passed, 150 skipped, nothing failing on head.3f6c4a40e:affected_tests.pyCPU arm, against fix(offload): key multimodal prompts by media, one LMCache engine id per server, MP anchors on the chunk grid, keep a hybrid's last chunk #3027's head753a55282: 1078 passed, 150 skipped, nothing failing on head.753a55282changes onlymp/scheduler.pyand its test; no GPU run on it yet.59feb21ae:review-pr'saffected_tests.pyCPU arm, against fix(offload): key multimodal prompts by media, one LMCache engine id per server, MP anchors on the chunk grid, keep a hybrid's last chunk #3027's head23d9ee232: 1078 passed, 150 skipped, nothing failing on head.acbd76533: 1502 passed, 138 skipped, 0 failed.09edcad0d:factory.py:416-417, which main has too.affected_tests.pyCPU arm, against4162fcde5: 1078 passed, 150 skipped, nothing failing on head.LoadOperationId(req_id=62) is still not terminal after 51s. That is the failure fix(offload): key multimodal prompts by media, one LMCache engine id per server, MP anchors on the chunk grid, keep a hybrid's last chunk #3027's23d9ee232fixes.recv_async_output,copy_done.synchronize()) on this branch, on fix(offload): key multimodal prompts by media, one LMCache engine id per server, MP anchors on the chunk grid, keep a hybrid's last chunk #3027, and on main69d7931b0without any offload. Main failed 4 times on two GPUs, including runs with a freshTMPDIRand inductor cache. fix(offload): key multimodal prompts by media, one LMCache engine id per server, MP anchors on the chunk grid, keep a hybrid's last chunk #3027's Qwen3-Next results (2caf12f1e,f1662fe4c) were taken before that. On this PR, the last Qwen3-Next run is963d7fa0d, before the split: 0.96 / 0.965 (eager).acbd76533:RUN_H2D_GPU_TESTS=1. With it, the 3 h2d files give 9 failures, all MLA cases intest_h2d_attention_publication.py: the test'sSimpleNamespacelacksATOM_SPARSE_INDEXER_LOGITS_BUDGET_MB. Main69d7931b0has the same 9.test_multiprocess_metrics' server-import subprocess, which timed out at load average 114 and passes on rerun (3 passed).amd/Qwen3.8-Flash-Next-Quark-MXFP4. Setup: TP1, MTP 3,--index-cache-dtype bf16, block 64, in-processlmcache_offload, questions 0-439, interval 256. Pass 1 stores and pass 2 fetches back from the CPU tier.acbd76533: 0.9773 / 0.9750, acceptance 0.714 / 0.729, 484,544 offload tokens, 433 of 440 state images from the tier.7270f5a68(acbd76533plus fix(offload): key multimodal prompts by media, one LMCache engine id per server, MP anchors on the chunk grid, keep a hybrid's last chunk #3027's restore fix): 0.9795 / 0.9795, acceptance 0.718 / 0.721, 483,008 offload tokens, 431 of 440.acbd76533: 0.9841 / 0.9795, acceptance 0.727 / 0.720, 459,712 offload tokens, 412 of 440.lmcache_mp, on7270f5a68, questions 0-199, ATOM restarted between the passes: 0.99 / 0.98, all 200 requests hit LMCache in pass 2 (90.5% of prompt tokens).b7ea8c2c0: 0.9773 / 0.9795, acceptance 0.711 / 0.730, 475,968 offload tokens, 426 of 440.7270f5a68andb7ea8c2c0Qwen3.8 was not rerun: the host's free NUMA memory did not fitoffload_split's pinned memory. The unit tests and Qwen3.5 runs above cover those heads.18ee7ac7b(acbd76533minus the layout-id extraction and one docstring), CPU full suite: 11277 passed, with the same 3 failures as onmain.12196f10b, before the split and the review fixes. GPU unit tests on every test file that touchesqwen4_exp,gdn_attnor QSA, plus the offload / checkpoint files: 1999 passed. The only failures are plugin tests that need vLLM / SGLang, which the container does not have.12196f10b: 0.9773 / 0.9818, 433 of 440 state images from the tier.4aff317a5(fix(qwen4-exp): register the QSA FP8 scale planes for KV offload #2655 on12196f10b): 0.9818 / 0.9750, 427 of 440.de10df694(fix(qwen4-exp): register the QSA FP8 scale planes for KV offload #2655's earlier head): 0.9886 / 0.9818, 430 of 440.Not changed
Seen while making this, left out of this PR:
_GDN_LINEAR_MODEL_TYPESsays the guard readshf_config.model_type; it reads the text config.MlaKvPool.region_tensorsstill walks its own views (the MLA pool does not useLayerMajorArena); delegating it touches every MLA model, so it is left for a follow-up.build_page_namespacestill branches onFamily.QSA_GDNrather than a per-model hook.main's rule, unchanged.kv_connectorby its raw spelling (atom/mono/runtime/deployment.py, main perf(qwen): integrate Mono decode and optimized TP1 MTP3 serving #2610). An alias such asLMCacheConnectorV1, ormulti, passes this PR's check but turns mono off, with an INFO log. Output is correct; only the mono decode speed-up is lost. Tracked in mono:resident_grid_refusalscompares the rawkv_connectorspelling, so connector aliases andmultisilently turn mono decode off #3032.Not covered
_joint_kv_boundary; whichever lands second merges.