Repository navigation
Conversation
delock
force-pushed
the
gma/qwen_ki_expr
branch
from
September 25, 2026 07:43
f1f9beb to
b67e55c
Compare
Add a segment-KI framework that accelerates greedy decode in the
hybrid-engine rollout by injecting native CUDA kernels into
comm-free segments of the model, plus a C++ decode loop that runs
the full autoregressive generation with near-zero per-step Python
overhead.
Segment detection is structural (attribute-name based): any HF
family whose gated MLP uses gate_proj/up_proj/down_proj or whose
GatedDeltaNet block uses the in_proj_{qkv,z,b,a} set is picked up
without per-model code. Projections carrying TP collectives are
never fused. The GLU replacement reads the original weight
Parameters directly through an autograd.Function with exact
backward, so train/generate share one forward path with no weight
copies and no inject/eject switching (rollout-training
co-location).
Components:
- fused_glu.cu: dual-weight GEMV+silu (MLP), GDN gates, decode
attention (warp-per-head, KV split-K, online softmax), fused
residual+RMSNorm, triple QKV GEMV, graph-capturable decode step
- decode_loop.cu: C++ generation loop taking a backend-agnostic
std::function replay callable (CUDA graph, eager forward, or
future SYCL); argmax and buffer updates run as a private kernel,
EOS checked via amortized D2H sync every 16 steps
- kernel_reference.py: pure-PyTorch executable specifications for
every kernel (test oracle, runtime fallback, porting reference)
- DeepSpeedStaticCache: hybrid-slot static cache with a
pass-through slot for linear-attention (GDN) layers that reuses
HF's cudagraph-safe state buffers by reference
- rollout integration: prefill -> cache setup -> graph capture ->
C++ decode loop, layered fallback (full-step graph for b=1,
graph+argmax for b>1, C++ loop, Python loop)
Verified on Qwen3.5-4B (RTX 4080 SUPER, greedy, b=1): correct text,
69.5 tok/s vs HF eager 22.5 (~3.1x), 92% of vLLM 74.9; co-location
checks (gradient flow, weight freshness, zero copies) all pass.
Signed-off-by: Guokai Ma <guokai.ma@intel.com>
delock
force-pushed
the
gma/qwen_ki_expr
branch
from
September 26, 2026 07:00
11231a7 to
d41c914
Compare
delock
changed the title
Signed-off-by: Guokai Ma <guokai.ma@intel.com>
use_segki=False previously still injected the decode attention and fused norm kernels in the graph-capture path (they were gated on op availability only), so disabling segki did not restore the fully native forward. Now the flag disables all kernel injection. Signed-off-by: Guokai Ma <guokai.ma@intel.com>
The GDN segment previously materialized a concatenated weight copy (_ki_gdn_fused_weight) at install time — stale after any optimizer step, so the fusion was effectively inference-only. Replace it with a quad_gemv kernel (b=1) that reads the four original weight matrices directly through a GDNInputProj autograd.Function with exact backward, mirroring the dual-weight GLU pattern: - b=1 decode: quad_gemv kernel, warp-per-row across qkv|z|b|a - b>1: on-the-fly concat of four GEMMs (kernel_reference oracle) - gradients flow to the original in_proj Parameters The fused weight buffer is gone; weight freshness after training is now guaranteed by construction for both GLU and GDN segments. Signed-off-by: Guokai Ma <guokai.ma@intel.com>
Signed-off-by: Guokai Ma <guokai.ma@intel.com>
Three groups covering the segment-KI contracts: 1. Kernel vs reference (GPU): every fused_glu kernel (dual_gemv, quad_gemv, gdn_gates, triple_gemv, fused_add_norm, decode_attn) compared against its kernel_reference.py oracle on identical inputs. 2. Injection invariants (CPU): apply_segment_ki forward equivalence, no weight copies installed, composite autograd path and backward gradient scattering to original Parameters. 3. Rollout-train-rollout integration (GPU, Qwen3.5-0.8B-Base): the co-location contract end-to-end — generate with graph capture + segKI, one SGD step through the same injected model, generate again with changed output, no inject/eject/sync anywhere in between. Executed on RTX 4080 SUPER: 11 passed (6 kernel-consistency, 4 CPU invariants, 1 integration). Signed-off-by: Guokai Ma <guokai.ma@intel.com>
…rch.cuda The PR moved DeepSpeedStaticLayer to the 'write_position = next slot to write == cached-token count' convention (the graph-path kernels write at the position then increment, so get_seq_length must return the position itself, not position + 1). Continuous batching is unaffected: it passes cache_position explicitly and never reads get_seq_length. The three master tests pinned the old +1 convention and cleared-row 0; updated to the new contract (-1 is the empty-row marker, matching cache initialization). Also replace torch.cuda.is_available() with get_accelerator() in the new kernel test file (check-torchcuda hook). Local: static_cache 6/6, segment_ki CPU 4/4 pass. Signed-off-by: Guokai Ma <guokai.ma@intel.com>
check-torchcuda also rejects tensor.cuda() calls; use get_accelerator().device_name() consistently (module-level DEV) so the tests are accelerator-portable alongside the GPU_AVAILABLE gate. Signed-off-by: Guokai Ma <guokai.ma@intel.com>
Signed-off-by: Ma, Guokai <guokai.ma@intel.com>
delock
marked this pull request as ready for review
September 28, 2026 06:27
delock
requested review from
hwchen2017,
loadams,
tjruwase and
tohtana
as code owners
September 28, 2026 06:27
![chatgpt-codex-connector[bot]](https://avatars.githubusercontent.com/in/1144995?s=60&v=4){kind=small}
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: b31faddfc4
ℹ️ About Codex in GitHub
Codex has been enabled to automatically review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
When you sign up for Codex through ChatGPT, Codex can also answer questions or update the PR, like "@codex address that feedback".
Comment on lines
+184
to
+185
| from deepspeed.module_inject.kernel_reference import dual_gemv_silu_mul | ||
| return self.down_proj(dual_gemv_silu_mul(input, self.gate_proj.weight, self.up_proj.weight)) |
There was a problem hiding this comment.
The structural selector accepts every bias-free module with these three projection names, but this replacement always computes SiLU. A GELU-gated MLP such as Gemma satisfies that selector and will therefore silently produce different logits whenever use_segki=True; constrain detection to SiLU MLPs or retain/call the module's configured activation.
Useful? React with 👍 / 👎.
Comment on lines
+123
to
+124
| grad_gate_w = torch.matmul(grad_gate_coef.t(), hidden) | ||
| grad_up_w = torch.matmul(grad_up_coef.t(), hidden) |
There was a problem hiding this comment.
For normal decoder MLP inputs shaped [batch, sequence, hidden], this branch is used during backward, but Tensor.t() only accepts tensors with at most two dimensions. Consequently, a training step through any injected GLU with a 3-D activation fails at this expression; flatten the leading dimensions as the GDN backward path does before forming the weight gradients.
Useful? React with 👍 / 👎.
Comment on lines
+233
to
+236
| if batch_size == 1 and seq_len == 1: | ||
| fused = GDNInputProj.apply( | ||
| hidden_states.squeeze(0).squeeze(0), self.in_proj_qkv.weight, self.in_proj_z.weight, self.in_proj_b.weight, | ||
| self.in_proj_a.weight, gdn_kernel_op).view(1, 1, -1) |
There was a problem hiding this comment.
A CUDA use_segki=True rollout installs the native GDN op regardless of model dtype, and the b=1 decode path passes it here. quad_gemv unconditionally obtains BFloat16 data pointers, so ordinary fp16 Qwen3.5 inference raises instead of using the composite path; gate native dispatch on bf16 (and contiguous compatible weights) or provide an fp16 kernel.
Useful? React with 👍 / 👎.
Comment on lines
+590
to
+594
| if getattr(module, "_ki_orig_layer_forward", None) is None: | ||
| module._ki_orig_layer_forward = module.forward | ||
| module._ki_norm_op = kernel_op | ||
| module._ki_postattn_scale = torch.add(1.0, norms[1].weight.detach()).contiguous() | ||
| module.forward = _fused_norm_layer_forward.__get__(module, type(module)) |
There was a problem hiding this comment.
This creates a detached, independent copy of the RMSNorm scale at injection time. After a co-located training step updates post_attention_layernorm.weight, graph decode still multiplies by the stale copy, so rollout no longer reflects the trained model; compute the scale from the live parameter in the forward path or refresh it after updates.
Useful? React with 👍 / 👎.
Comment on lines
+916
to
+917
| graph_op.decode_step_graph(static_logits[:, -1, :].contiguous(), static_token.view(batch_size, 1), | ||
| write_pos, static_attn, full_token_buf) |
There was a problem hiding this comment.
The graph-step op is loaded and used for every batch-one graph capture, even when use_segki is false, but decode_step_graph rejects logits that are not CUDA bf16. Existing fp16 graph-capture rollouts therefore fail during capture rather than reaching a fallback; only add this node for bf16 logits or retain the PyTorch argmax/update path.
Useful? React with 👍 / 👎.
Comment on lines
+198
to
+201
| int64_t new_pos = write_pos[0] + 1; | ||
| write_pos[0] = new_pos; | ||
| out_buf[new_pos] = best; | ||
| if (new_pos + 1 < max_len) { mask[new_pos + 1] = true; } |
There was a problem hiding this comment.
After the first captured decode forward, new_pos is the slot that will hold the next input token, but this reveals new_pos + 1. Thus any full-attention layer that remains on native SDPA sees its current KV slot masked and an unwritten future slot unmasked on the next replay, corrupting generation for models not patched by install_decode_attention; reveal new_pos instead.
Useful? React with 👍 / 👎.
Comment on lines
+273
to
+276
| layer_types = getattr(text_config, "layer_types", None) or ["full_attention"] * num_layers | ||
| self._layers = [ | ||
| DSStaticGDNSlot(max_cache_len) if t == "linear_attention" else DeepSpeedStaticLayer(max_cache_len) | ||
| for t in layer_types |
There was a problem hiding this comment.
For hybrid configurations this now substitutes DSStaticGDNSlot into the cache layer list, but continuous batching still iterates every layer and unconditionally copies target_layer.keys and target_layer.values in _continuous_prefill. The first admitted request on a model with a linear_attention layer therefore raises AttributeError; skip/pass through linear slots in that prefill and compaction path.
Useful? React with 👍 / 👎.
Comment on lines
+125
to
+127
| } | ||
|
|
||
| return steps_done; |
There was a problem hiding this comment.
When an EOS is found at a periodic check, buf_ptr is a CUDA device pointer, yet this host-side loop dereferences it directly. On normal CUDA allocations this is an invalid host memory access (and can crash) precisely when EOS is encountered; launch a padding kernel or use a device tensor operation instead.
Useful? React with 👍 / 👎.
…ed dispatch, live norm scale Codex review follow-ups: - only fuse gated MLPs whose activation is positively SiLU; the fused GLU hardcodes SiLU and would silently change logits for GELU-gated families - flatten leading dims in the GLU backward weight-gradient GEMMs; .t() rejects the [batch, seq, hidden] training activations - gate native dual_gemv/quad_gemv dispatch on bf16 + contiguous inputs so fp16 models fall back to the composite reference instead of raising - recompute the fused RMSNorm 1+w scale from the live Parameter in the forward path; the injected snapshot went stale after optimizer steps Signed-off-by: Guokai Ma <guokai.ma@intel.com>
…ous batching - only mount decode_step_graph for CUDA bf16 logits; fp16 rollouts keep the PyTorch argmax path instead of failing during graph capture - _continuous_prefill skips pass-through GDN slots (no keys/values to copy) - DeepSpeedStaticCache.compact only compacts KV layers; GDN slots keep no per-row rows, and layer[0] may be a GDN slot Signed-off-by: Guokai Ma <guokai.ma@intel.com>
- reveal new_pos (not new_pos+1) after a decode step: new_pos is the slot the written token occupies on the next replay; the off-by-one masked the current slot and unmasked an unwritten one for native-SDPA layers - decode_loop pads the post-EOS tail with a device kernel; writing the device pointer from host on EOS was an illegal host memory access Signed-off-by: Guokai Ma <guokai.ma@intel.com>
Qwen3.5 MLPs carry transformers.activations.SiLUActivation, not torch.nn.SiLU, so the newly added SiLU gate silently skipped every GLU/GDN segment on Qwen3.5 (caught by the rollout-train-rollout integration test asserting segments_replaced > 0). Verified on RTX 4080 SUPER: 12/12 tests pass (incl. integration), E2E Qwen3.5-4B text correct at 69.0 tok/s with segment fusion restored. Signed-off-by: Guokai Ma <guokai.ma@intel.com>
There was a problem hiding this comment.
Copilot review overview
🟡 Changes recommended
Several paths currently break gradients, padded prompts, fallback decoding, hybrid continuous batching, and native-op compilation or packaging.
Review effort: Balanced
Findings: 10
Open (18)
What changed in this PR
Adds segment-based kernel injection and native decode acceleration for hybrid-engine greedy rollout.
Changes:
- Adds fused CUDA kernels, reference implementations, and op builders.
- Extends static-cache support for hybrid attention models.
- Integrates graph capture, native decode loops, documentation, and tests.
| File | Description |
|---|---|
csrc/module_inject/fused_glu.cu |
Implements fused CUDA kernels. |
csrc/module_inject/decode_loop.cu |
Implements the native decode loop. |
deepspeed/module_inject/segment_ki.py |
Detects and injects fused model segments. |
deepspeed/module_inject/kernel_reference.py |
Provides PyTorch kernel references. |
deepspeed/ops/module_inject/__init__.py |
Exports native-op builders. |
deepspeed/ops/module_inject/fused_glu.py |
Adds the fused-kernel builder. |
deepspeed/ops/module_inject/decode_loop.py |
Adds the decode-loop builder. |
deepspeed/utils/static_cache.py |
Adds hybrid cache slots and position semantics. |
deepspeed/runtime/rollout/hybrid_engine_rollout.py |
Integrates injection and accelerated decoding. |
tests/unit/module_inject/test_segment_ki_kernels.py |
Tests kernels and injection behavior. |
tests/unit/utils/test_static_cache.py |
Updates static-cache expectations. |
docs/code-docs/source/inference-engine.rst |
Documents segment-based injection. |
💡 Add a code-review agent skill or configure MCP servers for context-aware, tailored reviews. Learn more in the docs.
…citly __nv_bfloat16 and the std::function<void()> pybind type caster were only reachable through transitive includes; declare both directly so the extension does not depend on include order. Signed-off-by: Guokai Ma <guokai.ma@intel.com>
The native gdn_gates and fused_add_norm ops have no autograd registration, so a training step through their fused branches silently cut gradients (gates: a/b projections, A_log, dt_bias; fused_add_norm: the residual/norm path). Route grad-enabled forwards through the composite/native ops; the kernels still serve the launch-bound no-grad decode path. Signed-off-by: Guokai Ma <guokai.ma@intel.com>
kwargs.get('cu_seq_lens_q') or kwargs.get('cu_seqlens') evaluates the
tensor's boolean value, which raises for any multi-element packed-sequence
input. Use dict.get's default argument instead.
Signed-off-by: Guokai Ma <guokai.ma@intel.com>
The decode_attn kernel attends every cached slot up to write_pos unconditionally and never sees attention_mask, so a padded prompt leaks pad positions into the softmax and yields logits that differ from the original path. Rollout flags padded prompts once per generate call and the injected forward rides the SDPA path for them. Signed-off-by: Guokai Ma <guokai.ma@intel.com>
The capture forward consumes static_token and advances the GDN conv/recurrent buffers whether or not decode_step_graph was mounted, so gating the restore on graph_op left the fallback graph (native op unavailable) starting from a state that already consumed the first decode token. Signed-off-by: Guokai Ma <guokai.ma@intel.com>
Collaborator
Author
Hi @nathon-lee , kindly reminder in case you didn't submit your review comments. Thanks! |
The conv and scan closures popped cu_seq_lens_q/cu_seqlens before calling causal_conv1d_fn and the FLA scan, letting state bleed across packed sequences; upstream passes them through. Filter only unrelated generation kwargs (use_cache). Signed-off-by: Guokai Ma <guokai.ma@intel.com>
triple_gemv consumes weights only, so enabling it for projections with biases (Qwen3.5 attention_bias) silently dropped all three biases and produced different logits. Require bias-free q/k/v for the fused path; the composite nn.Linear forward keeps honoring them. Signed-off-by: Guokai Ma <guokai.ma@intel.com>
The graph path passed the per-type attention-mask dict unconditionally, but the dict form is a hybrid-model contract: standard full-attention models (LLaMA and families) expect a tensor/None and regress under use_graph_capture even with use_segki=False. Build the per-type mapping only when layer_types contains linear_attention and keep the tensor form (2-D prompt mask at prefill, 4-D static mask at warmup/capture) otherwise. Signed-off-by: Guokai Ma <guokai.ma@intel.com>
The warp reduction kept the lower thread ID on equal bf16 maxima, but thread strides do not follow vocabulary order, so tied logits (common at bf16 with large vocabularies) could argmax to a different token than torch.argmax. Break ties on the smaller vocabulary index. Signed-off-by: Guokai Ma <guokai.ma@intel.com>
The b=1 full-step kernel had no direct coverage: verify argmax against torch.argmax including a tied-logits case (first-index rule), write_pos advancement, one-slot-per-step mask revelation, and output-buffer updates across repeated calls. Signed-off-by: Guokai Ma <guokai.ma@intel.com>
The rollout-train-rollout test asserted the greedy sequence changes after an SGD step, but changed logits need not change any argmax, so a correct live-weight implementation could fail it. Probe the injected projection directly: its output on fixed input must differ once the optimizer updates the weights it reads. The post-training generate still runs for path coverage. Signed-off-by: Guokai Ma <guokai.ma@intel.com>
Both tiers default to enabled; a second escape hatch beyond the config flags adds env lookups without a user. Kernel injection is already gated by use_segki and the per-dispatch checks. Signed-off-by: Guokai Ma <guokai.ma@intel.com>
Contributor
|
Hi, @delock Thanks for the PR and for addressing the earlier comments. I recommend addressing these two points before merging:
A few additional points to check; I am not treating these as confirmed merge blockers:
The position-handling regression is reproduced with the versions above. The other points are requests for testing or clarification, not confirmed runtime failures. |
Signed-off-by: Guokai Ma <guokai.ma@intel.com>
Sampling requests (temperature > 0) previously left the graph path for module.generate (~60% slower decode). The graph now captures the forward for them too and sampling runs outside it — temperature + top-k + top-p + multinomial on the replayed logits, measured at ~0.9% decode cost on Qwen3.5-4B. The full-step argmax kernel stays greedy-only; the Python step loop handles sampling and keeps the existing EOS semantics. Signed-off-by: Guokai Ma <guokai.ma@intel.com>
decode_attn was hardcoded to batch 1 (a single shared write_pos, flat [heads, dim] query). One block now covers a (row, query-head) pair with the grid spanning batch * heads, K/V index by row, and write_pos accepts a shared scalar or one entry per row — so batched decode rides the kernel instead of falling back to SDPA against the full-width static cache (the fallback measured 54% of vLLM at b=8, 2K context). The reference implementation and kernel test cover b=1, b=4 shared-pos, and b=4 per-row-pos. Signed-off-by: Guokai Ma <guokai.ma@intel.com>
decode_step_graph was b=1-only (single argmax, scalar write_pos, flat buffers). One block now covers a row: per-row argmax with the first-index tie rule, per-row write_pos advance, row-strided mask and token buffers; the rollout rides it at any batch size for greedy, with per-row EOS detection and tail padding. The python per-step loop now serves sampling only. GEMV kernels stay b=1: batched rows run the composite cuBLAS path where weights amortize across rows anyway. Signed-off-by: Guokai Ma <guokai.ma@intel.com>
The kernel is row-parallel (one block per row) but the guard kept it at b=1; batched decode paid eager RMSNorm in the profiled elementwise soup (~5 ms/step at b=8, 2K). Keep the seq-len, dtype, contiguity, and no-grad gates. Signed-off-by: Guokai Ma <guokai.ma@intel.com>
The eager q_norm/RoPE/output-gate were top elementwise contributors at b=8 (attribution-profiled). The kernel now normalizes q in-register (per-head RMSNorm with the (1+w) scale), applies rotate-half RoPE at the row's GPU-resident position (register pairing via shfl with lane^16, cos/sin tables indexed on-device so replays stay graph-safe), and multiplies the sigmoid gate in the epilogue. K keeps its eager norm + RoPE: the KV cache stores post-rope keys from prefill, so the kernel cannot re-derive them. Contract test covers b=4 with per-row positions. Signed-off-by: Guokai Ma <guokai.ma@intel.com>
The b>1 composite ran four GEMMs and an eager cat; the batched kernel (one warp per output feature per row, grid.y over rows) computes the concatenated projection directly — eliminating the cat and its repack from the GDN input path. Signed-off-by: Guokai Ma <guokai.ma@intel.com>
The composite reference ran two GEMMs plus the silu/mul activation as separate eager kernels; the batched kernel computes silu(h*Wg)*(h*Wu) in one launch per row block, removing the eager activation kernels (72 launches/step at b=8) from the MLP path. Signed-off-by: Guokai Ma <guokai.ma@intel.com>
Replaces the eager norm(x, z) (rsqrt + weight + silu gate as separate kernels, ~85 launches/step at b=8) with one row-parallel kernel computing (w * rms(x)) * silu(z); no-grad decode path only. Signed-off-by: Guokai Ma <guokai.ma@intel.com>
The batched quad/dual GEMV kernels (row-split and weight-stationary variants both) lost to the cuBLAS composite on end-to-end b=8 decode: the MLP weights dominate and re-streaming them per feature (or spilling weight-cached registers at in_features=2560) regressed the step by more than the eager cat/activation launches it saved (quad -10%, dual -60% in the slope bench, isolated by bisection). Batched rows keep the cuBLAS GEMM composite; the b=1 fused GEMVs are unaffected. Signed-off-by: Guokai Ma <guokai.ma@intel.com>
_generate_graph dropped cache_position/position_ids at some point after the original commit; HF then derives positions from the cache, which D2H-syncs and aborts stream capture on model/cache combinations like LLaMA + StaticCache (reproduced by the reviewer on PyTorch 2.9.1 / Transformers 4.57.1). Restore the original contract: prefill passes cache_position=arange(prompt_len); warmup and capture read persistent cache_position/position_ids tensors that the Python decode loops advance in place each step. Signed-off-by: Guokai Ma <guokai.ma@intel.com>
log(1+exp(x)) evaluates to exactly 0 in fp32 when x < ~-16 (1+exp(x) rounds to 1), silently zeroing the decay gate; torch's F.softplus uses log1p and keeps full precision. Branch to log1pf(expf(x)) there, keeping the x>20 shortcut. Signed-off-by: Guokai Ma <guokai.ma@intel.com>
…ents The kernel tests validate against our own reference implementations; both sides could share a wrong formula. Add GPU-gated tests that compare injected Qwen3.5 layers (full-attn, GDN, MLP) and the fused add+norm against the original HF math — forward outputs and input gradients with explicit bf16 tolerances, including the rounding-order case where HF materializes the bf16 residual sum before the norm. Signed-off-by: Guokai Ma <guokai.ma@intel.com>
- Log why the native op fell back (builder lookup / not registered / JIT failure), once per process, so users can tell which path is running. - Document throughput semantics (executed steps vs valid tokens; capture overhead excluded), graph-reuse semantics (recapture per generate, no replay across training steps), and the packed-sequence boundary (forwarded as in the native path; packed training untested). Signed-off-by: Guokai Ma <guokai.ma@intel.com>
The kernel indexed a single row only; multi-row launches had every block recompute row 0 and left rows 1+ of the output uninitialized (caught by the new HF-parity test at b=4, invisible to the single-row kernel test). One block now owns one row, matching the gated_norm kernel's layout, and the kernel test covers rows=1 and rows=4. Signed-off-by: Guokai Ma <guokai.ma@intel.com>
Collaborator
Author
|
Hi @nathon-lee , thanks for the review comments. I have followup them with subsequent commits. There is one finding, fused_add_norm didn't act correctly with BS>1, and it is found by comparing with HF forward.
|
Signed-off-by: Guokai Ma <guokai.ma@intel.com>
The torchcuda hook flags .is_cuda; on_accelerator() has different semantics on the CPU accelerator (CPU tensors are 'on accelerator'), which un-guarded the capture check on CPU-only builds. Gate on torch.version.cuda instead. Signed-off-by: Guokai Ma <guokai.ma@intel.com>
Replace the #ignore-cuda markers and the torch.version.cuda guard with a proper is_stream_capturing() method on the abstract accelerator (CUDA: torch.cuda.is_current_stream_capturing; CPU/others: False). Test files use the device_name() availability pattern instead of torch.cuda.is_available(). No CUDA-specific calls remain outside the accelerator layer. Signed-off-by: Guokai Ma <guokai.ma@intel.com>
Contributor
|
Hi @delock, thanks for following up. That was a useful catch: the previous I did notice one issue at |
This branch has not been deployed
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
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.
Summary
Adds a segment-KI framework that accelerates greedy decode in the hybrid-engine rollout by injecting native CUDA kernels into comm-free segments of the model, plus a C++ decode loop that runs the full autoregressive generation with near-zero per-step Python overhead.
Segment detection: structural, not per-model
Detection is attribute-name based: any HF family whose gated MLP uses
gate_proj/up_proj/down_proj(LLaMA, Qwen, Mistral, DeepSeek, Gemma...) or whose GatedDeltaNet block uses thein_proj_{qkv,z,b,a}set is picked up without per-model code. Projections carrying TP collectives are never fused (a defensiveisinstancecheck against the Allreduce layer hierarchy).Co-location by construction
Both GLU (dual-weight GEMV) and GDN (quad-weight GEMV) replacements read the original weight Parameters directly through
autograd.Functions with exact backward, so training and generation share one forward path: zero weight copies, no inject/eject switching between rollout and training. Verified end-to-end by an integration test (below): generate → train one step → generate, output changes, no sync calls.Components
csrc/module_inject/fused_glu.cucsrc/module_inject/decode_loop.custd::functionreplay callable (CUDA graph, eager forward, or future SYCL); argmax and buffer updates run as a private kernel; EOS checked via amortized D2H sync every 16 stepsdeepspeed/module_inject/segment_ki.pyapply_segment_ki)deepspeed/module_inject/kernel_reference.pydeepspeed/ops/module_inject/{fused_glu,decode_loop}.pydeepspeed/utils/static_cache.pydeepspeed/runtime/rollout/hybrid_engine_rollout.pyPerformance
Qwen3.5-4B, RTX 4080 SUPER, greedy, b=1:
Tests
tests/unit/module_inject/test_segment_ki_kernels.py— 11 tests, all passing (executed on RTX 4080 SUPER):kernel_reference.pyoracle on identical inputs (dual_gemv, quad_gemv, gdn_gates, triple_gemv, fused_add_norm, decode_attn)Notes
kernel_reference.pyserves as the executable spec;decode_loop'sstd::functioninterface is backend-agnostic (graph replay or eager forward); SYCL kernels would follow the same op-builder pattern.