Repository navigation
[Feat] [NVIDIA] megakernels: Expose hybrid runtime and DSpark verification views - #45406
Conversation
|
Companion provider and kernel draft: sgl-project/sgl-nvidia-mega#17. It depends on the typed hybrid runtime views proposed here. |
Oasis-Git
left a comment
Oasis-Git
left a comment
There was a problem hiding this comment.
[Codex]: Two non-blocking code-style suggestions for the new binding surface.
| "token_to_kv_pool", | ||
| "req_to_token_pool", | ||
| ) | ||
| if not all(hasattr(backend, name) for name in required): |
There was a problem hiding this comment.
[Codex]: Could we use an explicit backend capability/interface to identify the supported hybrid layout, then access its fields directly? The string-based hasattr checklist here, followed by getattr probes for the pools and mappings, makes the supported contract implicit and harder to check when backend interfaces change. Keeping this detection at one explicit boundary would fit the typed-view design better.
There was a problem hiding this comment.
Addressed in aa76be9. _prepare_hybrid is now the explicit supported-layout boundary: it checks HybridLinearAttnBackend, its TRTLLMMLABackend/KDAAttnBackend children, HybridLinearKVPool, HybridReqToTokenPool, and use_mla. Pools and mappings are then read directly into typed _HybridSources; the string-based attribute checklist and fallback probes are gone. This uses the existing concrete interfaces rather than introducing a new generic capability API.
Rechecked against the current PR head. Tests cover a lookalike backend, unsupported child backend/pool types, and a non-MLA pool without exposing a hybrid view. The latest framework run passed 26 tests and 48 subtests.
| block_table = getattr(mla_metadata, "block_kv_indices", None) | ||
| sequence_lengths = getattr(mla_metadata, "seq_lens_k", None) | ||
| write_slots = getattr(forward_batch, "out_cache_loc", None) | ||
| state_indices = getattr(kda_metadata, "mamba_cache_indices", None) |
There was a problem hiding this comment.
[Codex]: Once the supported metadata layout has been resolved, could we type the metadata/ForwardBatch parameters and read these declared fields directly? The getattr(..., None) calls turn a missing or renamed required field into the same generic error as a legitimately unset tensor. Direct access plus the existing None-value validation would keep the required interface explicit; the same applies to out_cache_loc_is_physical below.
There was a problem hiding this comment.
Addressed in aa76be9 and c26520b. The batch parameters now use ForwardBatch, MLA metadata is typed as TRTLLMMLADecodeMetadata, and KDA metadata uses its declared ForwardMetadata type. Required tensor fields and out_cache_loc_is_physical are read directly; _mla_metadata also reads the validated backend's forward_decode_metadata directly, without the optional batch-attribute probe.
The existing unset-value checks remain. Tests distinguish a missing field (AttributeError) from a present-but-None tensor (the explicit incomplete-metadata error), and cover uninitialized metadata and nonphysical write slots. These checks pass at the current PR head; the latest framework run passed 26 tests and 48 subtests.
|
Ready for review. The companion Kimi provider/kernel PR is sgl-project/sgl-nvidia-mega#17, with correctness and initial B1/B16/B64 performance results now included. This PR overlaps the binding/spec/runtime surfaces being moved in #45389. @BBuf, could we review and land this semantic API first, then have #45389 rebase and preserve it? /tag-and-rerun-ci |
nvpohanh
left a comment
nvpohanh
left a comment
There was a problem hiding this comment.
[by Claude Code] This review is done by agents. Please feel free to disagree if they do not make sense.
| f"megakernel binding: KDA layer {layer_id} has {len(conv)} " | ||
| "convolution states, expected one" | ||
| ) | ||
| states[layer_id] = KdaStateView(conv=conv, temporal=state.temporal) |
There was a problem hiding this comment.
[by Claude Code] Severity: functional | Confidence: Medium
Issue
KdaStateView exposes only conv and temporal. With --enable-linear-replayssm, mamba2_layer_cache() also returns the replayssm_d/k/g ring, and KDA decode writes the ring and flushes temporal only every --linear-replayssm-cache-len steps. _prepare_hybrid does not check for this, so the plugin gets an incomplete recurrent state. The Kimi K3 spec in sgl-nvidia-mega#17 does not refuse this flag either.
Example
A batch size outside the plugin widths runs stock ReplaySSM decode. The next covered step reads a temporal that lags the ring, and the next stock step replays the ring on top of a state the megakernel already advanced. Outputs drift silently.
Fix
In _prepare_hybrid, raise when state.replayssm_d is not None, or add the ring to the view.
There was a problem hiding this comment.
Addressed in 662e7d9. _prepare_hybrid now rejects any KDA layer with state.replayssm_d is not None before exposing the hybrid view, with an error identifying both ReplaySSM flags. KdaStateView is explicitly documented as non-ReplaySSM; this rejects the unsupported configuration rather than adding partial ring support.
Added regression coverage for ReplaySSM on either mapped KDA layer. The focused megakernel suite passes 23 tests and 41 subtests with CUDA unavailable; pre-commit checks also pass.
| ) | ||
|
|
||
| def _mla_metadata(self, sources: _HybridSources, forward_batch: Any) -> Any: | ||
| metadata = getattr(forward_batch, "decode_trtllm_mla_metadata", None) |
There was a problem hiding this comment.
[by Claude Code] Severity: style | Confidence: High
Issue
_mla_metadata probes with getattr(forward_batch, "decode_trtllm_mla_metadata", None) and getattr(sources.full_backend, "forward_decode_metadata", None). SGLang does not allow defensive getattr. Here it also hides an unsupported full-attention backend until the first covered step, which then fails with "MLA decode metadata is unavailable". Line 304 has the same problem: forward_mode is a declared ForwardBatch field, so getattr(forward_batch, "forward_mode", "unknown") is not needed.
Fix
In _prepare_hybrid, check isinstance(full_backend, TRTLLMMLABackend) (the CuteDSL and Tokenspeed MLA backends subclass it). Then read sources.full_backend.forward_decode_metadata directly; init_forward_metadata stores the same object in forward_batch.decode_trtllm_mla_metadata. On line 304, use forward_batch.forward_mode.
There was a problem hiding this comment.
Addressed in c26520b. _mla_metadata now reads sources.full_backend.forward_decode_metadata directly and no longer takes a batch argument. The TRTLLMMLABackend type check and direct forward_batch.forward_mode access are already in place.
The explicit uninitialized-metadata error remains, and capture getters still retain each graph's selected metadata object. Updated the tests to remove batch-override precedence while preserving multiple-bucket ownership and tensor-replacement checks. The focused suite passes 23 tests and 41 subtests with CUDA unavailable; pre-commit checks pass.
| *, model: torch.nn.Module, built: BuiltMegakernel, binding: MegakernelBinding | ||
| ) -> None: | ||
| """Chain a built object's optional post-pool preparation into the model hook.""" | ||
| prepare = getattr(built, "prepare_before_cuda_graph_capture", None) |
There was a problem hiding this comment.
[by Claude Code] Severity: style | Confidence: High
Issue
_install_preparation_bridge finds the new hook with getattr(built, "prepare_before_cuda_graph_capture", None) and getattr(model, "prepare_before_cuda_graph_capture", None). SGLang does not allow getattr probes. The hook is also not declared on BuiltMegakernel, so plugin authors and type checkers see it only in the spec.py docstring.
Fix
Declare the hook as data. For example, add an optional MegakernelSpec field such as prepare: Optional[Callable[[BuiltMegakernel, PreparationView], None]] = None; existing API v1 specs keep working, and install_megakernel reads it without getattr. For the model hook, consider calling the megakernel preparation from BaseRunner._warmup right after prepare_model_resources(mr) instead of patching model.__dict__.
There was a problem hiding this comment.
Implemented in 4611394. MegakernelSpec.prepare(built, view) is now an explicit optional callback, invoked from BaseRunner._warmup after model preparation and before autotune/capture. Removed built-object hook discovery and model-hook replacement; preparation uses the same binding passed to the build and dispatch. Specs without a callback remain unchanged. The companion Kimi registration is updated in sgl-project/sgl-nvidia-mega@756ad77.
Validation: 26 SGLang unit tests + 44 subtests passed, including ordering, warmup reset, absent callbacks, error propagation, and zero-width builds. A TP8 dummy-weight smoke passed startup, B1 megakernel graph capture, and a five-token generation request. The expanded plugin suite passed 124 tests + 28 subtests; its eight remaining failures reproduced on the unchanged parent pair, with no new regressions. No kernels changed.
| return self.pool.set_swa_key_buffer_radix_fused_norm_rope | ||
|
|
||
|
|
||
| @dataclass(frozen=True) |
There was a problem hiding this comment.
[by Claude Code] Severity: style | Confidence: High
Issue
This PR adds seven new @dataclass(frozen=True) classes: KdaStateView, HybridPoolView, PreparationView, MlaStepView, HybridStepView, _SealedSource, and _HybridSources. The repo rule .claude/rules/no-dataclasses.md says new data containers must use msgspec.Struct. The sibling contract types in megakernels/spec.py (Step, BuildContext, Requirements) already do.
Fix
Define them as msgspec.Struct, frozen=True (add kw_only=True where useful) and drop the dataclasses import.
There was a problem hiding this comment.
Addressed in 749b255: all seven records now use msgspec.Struct, frozen=True, and the dataclasses import is removed. Fields, order, defaults, and positional construction are preserved; no kw_only or API-version change.
The companion test-only update constructs malformed frozen views with msgspec.structs.replace instead of bypassing immutability with object.__setattr__; the same rejection cases and assertions are retained. No plugin runtime or kernel changes.
Validation: framework suite passed 26 tests and 44 subtests. Candidate and unchanged-parent plugin suites each passed 124 tests and 28 subtests, with the same eight pre-existing failures and no missing cases or new regressions. Pre-commit checks pass in both repositories.
| runner.attn_backend.req_to_token_pool.mamba_map | ||
| ), | ||
| } | ||
| values = { |
There was a problem hiding this comment.
[by Claude Code] Severity: style | Confidence: High
Issue
_seal_hybrid_sources keeps two 14-entry dicts, getters and values, that must list the same names. values was read from the same runner a moment earlier, so each getter already returns its value. If the dicts drift, a name is either not sealed or raises KeyError.
Fix
Keep only getters and seal with for name, get in getters.items(): self._seal(name, get(), get). seal_capture_view repeats the pattern: _hybrid_step and _capture_getters read the same four fields, and the "changed before seal" check compares two reads taken back to back. One read path that returns both the tensors and their getters would remove it.
There was a problem hiding this comment.
Addressed in 190da2e. _seal_hybrid_sources now uses only the ordered getters and calls _seal(name, getter(), getter). Capture sealing uses one table of (name, view tensor, owner-bound getter); the separate _capture_getters method and parallel maps are removed. Root-to-leaf ordering and per-graph metadata ownership are preserved.
I retained the changed before seal check: view creation and sealing remain separate calls, and dropping the check could seal a replacement buffer while the plugin still receives the old view. The existing replacement test now covers all four fields before the first seal as well as after sealing. This keeps the dispatcher and API unchanged.
Validation: 26 framework tests and 48 subtests passed; pre-commit passed. The plugin suite was unchanged at 124 passed and 28 subtests, with the same eight pre-existing failures reproduced on the unchanged parent and no new regressions. No plugin or kernel changes.
|
/tag-and-rerun-ci |
|
/rerun-failed-ci |
|
/tag-and-rerun-ci extra |
github-actions
Bot
added
the
run-ci-extra
|
Heads-up: #45389 is merging now and moves
On the plugin side, |
1b88d59 to
d8543ba
Compare
|
@BBuf Thanks for the migration map. Rebased this onto the new specification and consolidated binding paths, including the test locations and patch targets. The isolated Kimi compatibility commit prefers the new namespace and falls back only when that namespace is absent; it does not require DSpark. Native old/new-layout checks passed on GB300 for ordinary decode and the separately restacked DSpark extension: 34 focused test-file runs, with no failures or skips, plus verified module/type/registration identities and matching wheels across layouts. Kernel payloads are unchanged. DSpark remains in #46479 / sgl-project/sgl-nvidia-mega#25; prior full-model measurements remain attributed to their original revisions. |
leejnau
changed the title
|
Consolidated the existing DSpark commits from #46479 into this PR, keeping the commit boundaries intact:
The source tree is identical to the former #46479 tip; this is a fast-forward, not a rewrite of the reviewed foundation. The DSpark-only diff can be reviewed separately. sgl-nvidia-mega #25 remains the companion implementation. Validation attribution and limitations are unchanged in the updated description. |
|
cc @YAMY1234 |
| return self.original( | ||
| input_ids, positions, forward_batch, *args, **kwargs | ||
| ) | ||
| view = self.binding.forward_view(forward_batch) |
There was a problem hiding this comment.
Would it make sense to reconcile these calls with the architecture-specific bindings already on main before rebasing? Main selects EagleBinding for Qwen3.5, whose forward_view(self) takes no batch argument and which does not implement seal_capture_view or seal_capture_buffers (also called unconditionally at L1381). Keeping main's binding selection, a covered Qwen3.5 forward fails with TypeError: EagleBinding.forward_view() takes 1 positional argument but 2 were given; adapting only that signature still leaves the capture hooks failing. The same EagleBinding passes main's existing dispatch/capture-tail path. Could we use a common lifecycle contract, or scope the new hooks to bindings that implement them, and add a regression that runs the real EagleBinding through the covered forward and the capture tail? The generic binding mocks don't reach this boundary.
There was a problem hiding this comment.
Agreed—fixed in 27687018. The rebase preserves main's architecture-specific binding selection. Only MegakernelBinding receives the batch argument and the new capture-sealing hooks; EagleBinding keeps its no-argument forward_view() and existing assert_live() capture-tail check.
Added a regression using the real EagleBinding selected for Qwen target and MTP draft architectures, covering eager/capture dispatch, the capture tail, and detection of a replaced KV buffer. The CPU-only isolated-import regression reproduced the reported TypeError before the fix and passes afterward. 3e41ef87 also aligns the runtime test fixtures with API v2; companion sgl-nvidia-mega#25 now declares v2 and tests the optional entry/variant contract.
| legacy: bool = True, | ||
| ): | ||
| self.hybrid = hybrid | ||
| if legacy: |
There was a problem hiding this comment.
Non-blocking: would it make sense to avoid allocating the six accessor closures for every legacy ForwardView, while keeping lazy access to the optional backend methods? This runs on every covered eager forward, including non-hybrid models (CUDA graph replay avoids it). I haven't measured a regression; just keeping the eager path's host work small.
There was a problem hiding this comment.
Addressed in 5f31ea56. Legacy views now retain one private backend reference and resolve optional operations through class properties, avoiding per-view accessor closures while preserving lazy access. Hybrid views still retain only their typed surface, without a backend reference.
Extended the CPU test to check that no callable accessors are stored on the instance, missing optional methods remain lazy, later backend updates are visible, and hybrid views do not expose legacy metadata/pool/backend references. The isolated-import checks pass; no performance improvement is claimed.
d3d6d8c to
3e41ef8
Compare
BBuf
left a comment
BBuf
left a comment
There was a problem hiding this comment.
Please address the [P1] items, then re-ping.
Agent-assisted pre-review, findings verified against 3e41ef876401.
|
|
||
| states: dict[int, KdaStateView] = {} | ||
| for layer_id in mamba_map: | ||
| state = request_pool.mamba2_layer_cache(layer_id) |
There was a problem hiding this comment.
[P1] HybridReqToTokenPool has no mamba2_layer_cache() at this head. The hybrid preparation callback reaches this line during warmup and raises AttributeError before capture, for both decode and DSpark. Use the KDA backend's state_layer_cache(layer_id)/state router and update the four sealed getters too. The test-only method currently hides this API mismatch.
|
|
||
| class RequestPool(HybridReqToTokenPool): | ||
| def __init__(self, *, verify=False, q=4): | ||
| self.mamba_pool = MambaPool(verify=verify, q=q) |
There was a problem hiding this comment.
[P1] This assignment invokes the inherited setter, which writes self.state.pool, but the fixture leaves self.state as None. Every hybrid_backend() call therefore raises before the new assertions run under native imports. Initialize StateStorage as the real pool does and validate this registered file without substituting framework classes.
Use the KDA backend state API and initialize native state storage in the regression fixture.
|
@leejnau I pushed Automated pre-review for |
BBuf
left a comment
BBuf
left a comment
There was a problem hiding this comment.
Please address the [P1] items, then re-ping.
Agent-assisted pre-review, findings verified against cc01927631b2.
| with ( | ||
| patch(BACKEND, return_value=fixture.backend), | ||
| patch.object( | ||
| torch.cuda, "is_current_stream_capturing", return_value=True |
There was a problem hiding this comment.
[P1] This patches capture to True on CPU tensors, so it never exercises graph replay or the KDA verify/commit path. The supplied GPU results are from older revisions, and the existing Kimi GPU suite does not enable this plugin. Could we add a native GPU regression covering partial acceptance, changed inputs/slots, and replay, then run it on this head with a pinned plugin revision?
|
/rerun-failed-ci |
1 similar comment
|
/rerun-failed-ci |
…ation views (sgl-project#45406) Co-authored-by: BBuf <35585791+BBuf@users.noreply.github.com>
Motivation
Hybrid MLA/KDA megakernel plugins need typed access to SGLang-owned cache and step state after model resources are prepared and before CUDA graph capture. This includes ordinary decode and fixed-width Kimi K3 DSpark target verification. The companion implementation is sgl-nvidia-mega #25.
Modifications
9b00e04f: explicit post-resource preparation, typed per-step views, graph-static ownership checks, and existing non-hybrid behavior, rebased onto the current NVIDIA kernel-plugin layout and API v2.ca36ffa8: request-local KDA verification state, MLA query-width metadata, auxiliary hidden-state buffers, and capture-lifetime checks. SGLang retains ownership of acceptance, recurrent-state commit, and stock draft-model execution.27687018: preserve architecture-specific EagleBinding dispatch/capture behavior, with a real-binding regression.5f31ea56: avoid per-view legacy accessor closures while retaining lazy access and typed-only hybrid views.3e41ef87: align runtime test fixtures with the API-v2 installed-object and optional-entry contracts.The foundation and DSpark commits formerly proposed in #46479 remain separate reviewable commits. No kernel arithmetic or default execution path is changed. The focused DSpark configuration is GB300, TP8, five draft proposals plus one anchor (q6), and DCP off. Uniform q4–q8 metadata support retains fail-closed layout and ownership checks. The original full-model evidence uses an 8,192-token context; bounded long-context qualification is described separately in the companion plugin PR.
The isolated Kimi plugin compatibility commit in the companion PR supports ordinary decode independently of DSpark. It prefers the new interfaces and falls back only when the new namespace is absent.
Accuracy Tests
Current rebased pair: SGLang
3e41ef8764018321d612d15f69da8a9f4c35796aand plugin90e1b1b48a79d556dbddc093a3358b1d9b7912e2. CPU-only Python checks passed 55 tests and 150 subtests, covering the actual binding, registry, runtime/preparation functions and plugin contracts, including real EagleBinding target/draft dispatch and capture-tail checks. Server/GPU dependencies unavailable on the host were isolated with test stubs; this is not a full native-framework import or GPU requalification. Focused repository lint, import ordering and formatting checks pass. The GPU measurements below remain attributed to their original revisions.The native old/new-layout matrix passed on GB300 for ordinary-decode and DSpark bindings: all 34 isolated test-file runs completed without failures or skips. Checks covered actual module/type/registration/preparation identities and audited installed wheel payloads, identical across layouts within each pair. The ordinary-decode pair is SGLang
d8543baa/ plugin32262d4; the native-tested DSpark pair is0e9e2aa/e64e4d0. The subsequentd3d6d8cchanges only one test import's formatting, with identical AST and no production changes. This is focused compatibility validation, not a new full-model accuracy run.Historical full-model measurements apply specifically to SGLang
02406d4ef05680cbf3d77287354d887b039141d9+ sgl-nvidia-mega194d693d15a51f74a727889531fd6d9cd01bf355, not the rebased pair. GSM8K was 1,283/1,314 specialized versus 1,284/1,314 stock; accepted draft tokens per verification were 3.7966 versus 3.8264, excluding the bonus token. The known numerical canary remains FAIL, separately from quality and timing; these results do not establish general numerical equivalence.Speed Tests and Profiling
On that same historical DSpark pair, matched end-to-end output throughput was 1.43–2.09x stock for B1/B4 and 128/6,144-token inputs; see the measurements and limitations. These are bounded end-to-end results, not isolated kernel speedups or new measurements on the consolidated head. Earlier ordinary-decode measurements are recorded in merged sgl-nvidia-mega #17.
Checklist
CI States
Latest PR Test (Base): ✅ Run #41152729305
Latest PR Test (Extra): ❌ Run #41152729359
Latest PR Test (AMD ROCm 10): ❌ Run #41152729331