[None][feat] Port DFlash weight-load checks, DSpark norm_dim and disagg tests from the Rubin branch - #19185
Draft
reasonsolo wants to merge 3 commits into
Draft
[None][feat] Port DFlash weight-load checks, DSpark norm_dim and disagg tests from the Rubin branch#19185reasonsolo wants to merge 3 commits into
reasonsolo wants to merge 3 commits into
Conversation
Collaborator
|
DFlash/DSpark related changes will be merged in this PR: #19040 |
reasonsolo
changed the title
…agg tests from the Rubin branch Merge-back of the portable subset of the Rubin runtime work. Most of the originally scoped feature (KV cache manager v2 uGPU localization, the DSpark MLA drafter, the disagg benchmark fill gate) is NOT included: its production half is absent from the source branch as well as from main, so porting the corresponding tests would have added red tests rather than coverage. See the PR description for the per-file disposition. Production: * modeling_dflash: fail loudly when a draft checkpoint omits backbone weights the drafter does not share with the target. `_assert_backbone_complete` walks the constructed module tree rather than a hand-kept list, and `WRAPPER_OWNED_WEIGHTS` gates the wrapper-owned `fc` / `hidden_norm` that a module walk cannot see. Previously `allow_partial_loading=True`, needed for the target-shared `embed_tokens` / `lm_head`, let a truncated checkpoint load clean and leave parameters at `torch.empty`. Also restructures the attention-backend init so every op handle starts at None and adds the `_uses_worker_attention_backend` hook for drafters that bring their own block decode; the FA4 backend added on main is preserved. * modeling_utils: hoist the fused-module component table out of `_load_weights_impl` into `FUSED_MODULE_COMPONENTS` so the check above reads the same table the loader uses instead of keeping a third copy. No behavior change. * modeling_speculative: bound an external drafter's position tables with the target's `max_seq_len`. Without it the drafter falls back to the checkpoint's advertised `max_position_embeddings` and allocates a table sized for a context the runtime never serves. * speculative/dflash: clamp the running context length before deriving the block-decode query positions. `_ctx_len` is clamped to `_max_ctx` only after the step's accepted tokens are folded in, so a request already at the ceiling indexed `num_accepted` positions past the end of the sequence. * dspark_rmsnorm_rope: add an optional `norm_dim` so the fused RMSNorm/RoPE kernel can normalize a prefix of the row and pass the remaining rope lanes through raw, as DeepSeek-style MLA needs. Defaults to the whole row and is bit-identical to the previous behavior there. `norm_dim` is now also forwarded to the support predicate, which the source branch omitted. * kv_cache_manager_v2/_block_radix_tree: document why the stale-tail prune requires every life cycle to be pageless, mirroring the C++ implementation. Tests: * New `test_block_radix_tree_stale_prune` covers that predicate, including the negative control that dead tails still get pruned. * New `test_disaggregated_multinode` is a two-node DSpark disagg accuracy harness. It skips unless run under a 2-node, 1-task-per-node Slurm allocation; the exact srun invocation is in the file header. It is deliberately left out of the test lists, as no current stage uses that layout. * `test_disagg_index_mapper_early_release` gains a case asserting that `release_index_slot` detaches every page-index view before the slot is reused. The import is adjusted for the `pyexecutor.kv_cache` package move. * `test_cache_reuse_adapter` gains a case for the DSpark gen-init draft reserve being removed before the SWA trim. Signed-off-by: Lizhi Zhou <1432185+reasonsolo@users.noreply.github.com>
…dtype An external drafter has its own KV cache, sized and allocated separately from the target's. It was inheriting the target's quantization config, so an fp8-KV target silently forced the drafter's pool to fp8 as well -- charging the budget at a dtype the drafter never asked for and does not necessarily support. Neutralize the target's KV quant config when building the draft cache manager for an external drafter, and fall the draft KV dtype back to "auto" when the inherited value is "fp8". Ported from the Rubin branch, where this landed as 33624323a6 and was then lost when a later replay commit (2ba428f26f) re-took pyexecutor/_util.py from main. Restored on the internal branch as 8626bd6577; content survival of the original measured 0%. NOT RUN: no GPU here. Needs an external-drafter plus fp8-KV-target speculative decoding run to validate. Signed-off-by: Lizhi Zhou <1432185+reasonsolo@users.noreply.github.com>
…the rebase On user/lizhiz/rubin-advance, configure_kv_cache_capacity re-derives max_tokens from the final estimated budget once estimation finishes whenever the V2 KV cache manager is in use. The rebase onto main dropped that block; main never had it either. Without it, a one-model speculative-decoding draft KV cache that shares the config reads the same max_gpu_total_bytes as the target. That cap does not scale with the draft's much smaller per-token footprint, which depends on num_local_layers, so both managers claim the whole budget and the draft OOMs. Deriving max_tokens restores V1's behaviour: V2's quota becomes min(max_gpu_total_bytes, max_tokens * bytes_per_token), which picks the layer-scaled draft budget when the draft manager reads the shared config. An explicit user-provided max_tokens still wins. Restored as it stands on rubin-advance, using this branch's existing self._is_kv_cache_manager_v2 attribute rather than recomputing the issubclass check locally. Signed-off-by: Lizhi Zhou <1432185+reasonsolo@users.noreply.github.com>
reasonsolo
force-pushed
the
user/lizhiz/mb-runtime
branch
from
September 18, 2026 02:59
0de2ee5 to
b493415
Compare
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
2 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Description
Ports the portable subset of the Rubin branch's runtime work — DFlash weight-load checks, the DSpark
norm_dimfused path, and disaggregated-serving test coverage.models/modeling_dflash.py—_assert_backbone_completeweight-load validation, folded aroundmain's FA4 backend so the_uses_worker_attention_backendgate covers all three backends rather than two.modeling_utils.pyhoistsFUSED_MODULE_COMPONENTSfrom an inline dict literal to a module constant so the check can use it (no behaviour change).speculative/dflash.py—ctx_len_nowclamp fix.cute_dsl_kernels/blackwell/dspark_rmsnorm_rope.py+ its custom op —norm_dimsupport. Also fixes a Rubin omission:cute_dsl_dspark_rmsnorm_ropedid not forwardnorm_dimto its own support predicate, so the fused path silently never engaged for the non-default case.test_disaggregated_multinode.py(new),test_disagg_index_mapper_early_release.py, part oftest_cache_reuse_adapter.py.Scope: 11 of 31 candidate files
Two thirds of the candidate set was dropped after verification, because the production half is missing from both
mainand the Rubin branch — these are test-only cherry-picks that are already red at their source:test_kv_cache_localization.py(+848)test_dspark_attention.py(+905)_output_projection,_qkv_projections— module-level import, so the whole file fails to collect, including testsmaincurrently passestest_kimi_k3_dspark_semantics.py(+718)MLADSparkForCausalLM+ yarn-rope helpers (~700 of 718 lines)test_benchmark_disagg.py_benchmark_*attrs onPyExecutorspeculative/interface.py,eagle3.pyattn_metadata.ugpu_enabled,draft_kv_cache_manager.fork_join_attnspeculative/model_drafter.pymainby #18721Also excluded:
scheduler_v2.py, which contains a duplicate_try_recompute_pause_for_gendefinition (lines 1477 and 1560) where the second shadows the first and bypassesmain's newer_try_allocate_generation, losing the joint draft-pool check; and its equal-cost context chunking is env-gated off and incomplete (commit_ctx_costis never called, so the budget never binds).The
K3DsparkForCausalLMarchitecture registration was dropped:_arch_index.pymaps it tomodeling_dspark, but no such class exists there onmain. It is also an unreleasedInferact/Kimi-K3-DSparkname, so it would needarchitecture_allowlist.pyclearance it cannot currently get.Test Coverage
No new test-list entries needed — both ported unit tests land in directories already covered by directory-level entries.
test_disaggregated_multinode.pyis not registered: it requires a 2-node / 1-task-per-node Slurm allocation that no current stage provides, so registering it into an existing GB300 2-node stage would only ever skip. The exactsruninvocation is in the file header; this needs a CI owner's decision.No test has been executed and nothing has been built. Static verification only: hunk-by-hunk review,
python3 -m compileall, and pre-commit.The
norm_dimproduction change lands untested onmain— its test is blocked on_rmsnorm_rope_batched(norm_dim=)inmodeling_dspark.py, which is outside this PR. It is additive and bit-identical at the default; if carrying an untested knob is unwelcome, those two files can be dropped.PR Checklist
🤖 Generated with Claude Code