Repository navigation
Conversation
11231a7 to
d41c914
Compare
delock
changed the title
![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".
| 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 👍 / 👎.
| 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 👍 / 👎.
| 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 👍 / 👎.
| 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 👍 / 👎.
| 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 👍 / 👎.
| 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 👍 / 👎.
| 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 👍 / 👎.
| } | ||
|
|
||
| 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 👍 / 👎.
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.
|
Hi @delock, thanks for submitting this PR. I have left a few comments and questions. |
There was a problem hiding this comment.
Copilot review overview
🟡 Changes recommended
Graph compatibility, packed GDN execution, projection bias handling, and deterministic argmax behavior have unresolved correctness issues.
Review effort: Balanced
Findings: 3
Open (6)
Resolved since last review (18)
| kw.pop("use_cache", None) | ||
| kw.pop("cu_seqlens", None) | ||
| kw.pop("cu_seq_lens_q", None) |
| module._ki_qkv_op = kernel_op if (hasattr(kernel_op, "triple_gemv") | ||
| and os.environ.get("DS_TIER2", "1") == "1") else None |
| attention_mask={ | ||
| "full_attention": None, | ||
| "linear_attention": None | ||
| }, |
| if (tid < stride) { | ||
| if (s_val[tid + stride] > s_val[tid]) { | ||
| s_val[tid] = s_val[tid + stride]; | ||
| s_idx[tid] = s_idx[tid + stride]; | ||
| } | ||
| } |
| m.def("decode_attn", &decode_attn, "b=1 decode attention with GQA, graph-compatible (CUDA)"); | ||
| m.def("decode_step", &decode_step, "fused decode step update (CUDA)"); | ||
| m.def("decode_step_graph", &decode_step_graph, "graph-capturable decode step (CUDA)"); |
| batch2 = rollout.generate(req, greedy) | ||
| assert not torch.equal(batch1.input_ids, batch2.input_ids), \ | ||
| "rollout output unchanged after training — kernels read stale weights" |
Hi @nathon-lee , kindly reminder in case you didn't submit your review comments. Thanks! |
|
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. |
|
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.
|
|
Hi @delock, thanks for following up. That was a useful catch: the previous I did notice one issue at |
Codex Review SummaryThis comment shows the latest Codex review activity on this pull request.
ℹ️ About Codex in GitHubYour team has set up Codex to review pull requests in this repo. Reviews are triggered when you
Codex reacts with 👀 while any review is running, comments if it has suggestions, and reacts with 👍 once all reviews finish with no findings. |
Hi @nathon-lee , sorry for the late reply, let me check with latest status of master. |
Hi @nathon-lee , the latest code should contain this missing implementation. Sorry for inconvience. |
| def total_memory(self, device_index=None): | ||
| ... | ||
|
|
||
| @abc.abstractmethod |
There was a problem hiding this comment.
I think we should not remove this?
| per-segment graphs with collectives issued between replays.""" | ||
| return False | ||
|
|
||
| def is_stream_capturing(self): |
There was a problem hiding this comment.
Yes, we can reuse the interface on master. is_stream_capturing is removed.
| # b>1 attends every slot the mask reveals. | ||
| static_attn[:, :, :, pos] = True | ||
| get_accelerator().replay_graph(graph) | ||
| next_token = static_logits[:, -1, :].argmax(dim=-1, keepdim=True) |
There was a problem hiding this comment.
I think we should have sampling for temp>0 as per the docstring (line 9-10).
We are removing is_greedy check in line 232 so even for sampled req (with b>1) we might reach here to get greedy?
|
The new |
…_greedy Adopt master's routing structure exactly (#8980 + #8939): the local use_graph_capture is the config value narrowed by is_graph_capture_supported(). No is_greedy gate — #8939 added sampling to the graph decode path (token selection runs outside the captured graph, ~0.9% decode cost). segKI's only increment beyond master is the needs_pad_fallback condition. The previous commit referenced the local variable without defining it. (flake8 skipped: sandbox) Signed-off-by: Ma, Guokai <guokai.ma@intel.com>
Master's #8939 refactored _continuous_next_tokens to accept a select_next_tokens parameter (applying repetition penalty via _apply_repetition_penalty, then delegating token selection to the injected selector). Our auto-merged version kept the old 5-arg inline implementation. Updated the method, both call sites, and added select_next_tokens definitions in the enclosing continuous methods. (flake8 skipped: sandbox) Signed-off-by: Ma, Guokai <guokai.ma@intel.com>
The previous commit added @classmethod without removing the existing @staticmethod; both together make the method uncallable as a bound method. (flake8 skipped: sandbox) Signed-off-by: Ma, Guokai <guokai.ma@intel.com>
Master's tests mock the module-level logger and use SimpleNamespace mocks that lack num_attention_heads. Our merge dropped the logger definition and the TP head-count adjustment assumed the attribute exists. (flake8 skipped: sandbox) Signed-off-by: Ma, Guokai <guokai.ma@intel.com>
…iction Same auto-merge pattern: master's #8939 added a sampling parameter to _continuous_prefill_cache and removed the greedy-only restriction in continuous batching. Both call sites updated. (flake8 skipped: sandbox) Signed-off-by: Ma, Guokai <guokai.ma@intel.com>
The previous commit added sampling to _continuous_prefill_cache's signature and its call sites inside _continuous_prefill, but forgot to add sampling to _continuous_prefill itself and its callers — yielding F821 undefined name at the inner call sites. (flake8 skipped: sandbox) Signed-off-by: Ma, Guokai <guokai.ma@intel.com>
The #8939 merge resolution dropped three master behaviours in hybrid_engine_rollout.py: - _generate_graph lost the flash_attention_2 rejection, so graph capture proceeded into StaticCache and crashed inside transformers instead of raising the documented ValueError. - _generate_continuous routed on self.use_graph_capture directly, losing the accelerator capability gate, so backends without the graph-capture opt-in entered the graph path instead of falling back to eager continuous batching. - generate() logged through a local logging.getLogger while DeepSpeed's logger is the one exposing warning_once. Restore master's exact code for all three and import the DeepSpeed logger. Signed-off-by: Ma, Guokai <guokai.ma@intel.com>
|
#9008 needs to be merged before this PR as a better AutoTP solution. |
tohtana
added
the
needs-gpu-test
…utoTP (deepspeedai#9008) ## Summary Size HybridEngine rollout KV caches from each rank's own prefill KV tensors, not from the global model config. Under AutoTP, each rank holds only its share of the KV heads. Today the cache is allocated for the full model's KV-head count, so TP>1 rollout fails on all three decode paths: | Path | TP>1 failure today | Change | | --- | --- | --- | | Plain graph (`_generate_graph`) | Passing batch size, device and dtype makes `DeepSpeedStaticCache` initialize from `module.config` up front, so the lazy per-layer setup never runs. The prefill copy then fails with `4 vs 2` heads, or broadcasts silently when the rank has 1 local head. | Drop those up-front arguments, so the existing lazy setup sizes each layer from the prefill tensors before capture. | | Continuous eager (`_generate_continuous`) | The same up-front setup from config fails on the first prefill copy (`4 vs 2`, `2 vs 8`). | Create the cache lazily. `_continuous_cache_layer` initializes each layer from the prefill tensors, sized to the slot count. The slot count is the decode batch capacity, not the prefill batch. | | Continuous graph | The graph state cache is initialized from config, and capture happens before any prefill. | Create the cache lazily and run one decode forward on the current stream before warmup, so every layer is initialized with local shapes. Without this, those tensors would be allocated on the warmup side stream. | `lazy_initialization` gains an optional `max_batch_size`, which the continuous path uses because its capacity is the slot count. `cache_memory_bytes` is now computed after decode, once the layers exist. No model config is changed and no TP/world group is consulted, so TP×DP layouts work with no extra logic. ## Related work / scope - deepspeedai#9000 / deepspeedai#9004 fix attention-mask construction during capture, not KV shape allocation. This PR is independent of them. For real Hugging Face model validation, deepspeedai#9004's runtime hunk was applied to **both** the baseline and the fixed tree. - deepspeedai#8602 divides cache heads by the global world size, which breaks TP×DP (for example TP2×DP4). This fix reads the tensors directly instead of inferring the TP group. - On Transformers 5.17, continuous eager also needs deepspeedai#8852 (`get_query_offset`), which is still open, so the continuous eager HF runs below apply deepspeedai#8852 to both baseline and fixed trees. The new unit tests use a local model and do not depend on deepspeedai#8852. - **Known limitation, not changed here:** a persistent continuous graph (deepspeedai#8732) that contains NCCL ops makes `torch.distributed.destroy_process_group()` hang in `ProcessGroupNCCL.shutdown()` at exit. Releasing the rollout (`del` + `gc.collect()` + `synchronize()`) before destroying avoids it. Before this PR, TP>1 continuous graph could not run, so the hang was not reachable. A release/close API is left for a follow-up. ## Validation **Hardware/software:** 8× NVIDIA L20, driver 580.173.02, CUDA runtime 13.0, torch 2.14.1+cu130, Transformers 5.17.0, Python 3.10.22. Baseline: `d8241a9ba266405e1a58fadf19e24b2370f6d7fa`. ```bash CUDA_VISIBLE_DEVICES=0 python -m pytest --forked tests/unit/v1/test_hybrid_engine_rollout_v1.py \ -k "static_graph_uses_prefill_kv_shape or continuous_generation_uses_local_kv_shape" # baseline: 6 failed, 3 passed; this PR: 9 passed CUDA_VISIBLE_DEVICES=0 python -m pytest --forked tests/unit/v1/test_hybrid_engine_rollout_v1.py \ tests/unit/runtime/rollout/test_hybrid_engine_rollout.py tests/unit/utils/test_static_cache.py # baseline: 6 failed, 64 passed, 1 skipped; this PR: 70 passed, 1 skipped ``` The new tests write head-dependent cache contents, compare tokens with an independent oracle, and count real graph replays: | Local / config KV heads | Plain graph | Continuous eager | Continuous graph | | --- | --- | --- | --- | | 4 / 4 | pass → pass | pass → pass | pass → pass | | 2 / 4 | fail → pass | fail → pass | fail → pass | | 1 / 4 | fail → pass | fail → pass | fail → pass | **Multi-GPU tiny Llama, continuous path.** 2 layers, 4 KV heads, fp32, 4 requests per rank, 2 slots. Tokens were compared with plain eager generation on every rank. | Layout | Continuous graph (before → after) | Continuous eager + deepspeedai#8852 (before → after) | | --- | --- | --- | | TP1, world 1 | pass → pass | pass → pass | | TP2, world 2 | shape error → pass | `4 vs 2` error → pass | | TP4, world 4 | shape error → pass | `2 vs 8` error → pass | | TP2×DP4, world 8 | shape error → pass | `4 vs 2` error → pass | | TP4×DP2, world 8 | not run → pass | `2 vs 8` error → pass | **Plain graph with real models**, with deepspeedai#9004 applied to both trees: | Model / precision | Layouts | Result | | --- | --- | --- | | Tiny Llama / fp32 | TP1×DP1, TP1×DP4, TP2×DP2, TP2×DP4 | All pass; graph tokens and masks equal native eager generation | | Qwen3-8B / bf16 | TP2×DP4 | 4 prompts × 16 new tokens; graph tokens and masks equal native eager generation | On the plain-graph runs, each rank made three generation calls: 3 captures and 45 replays, and repeated graph outputs matched. The baseline TP2×DP4 Llama run failed at the KV copy (`4 vs 2`). The Qwen3 graph calls took 0.464 / 0.439 s, including prefill and capture, with 8.25 GiB max allocated per GPU. These numbers come from correctness runs; **no speedup is claimed**. **Distributed AutoTP integration test** (`TestAutoTPContinuousRollout`, world size 2): real AutoTP via `deepspeed.initialize` (`autotp_size=2`) on a tiny Llama, 3 requests on 2 continuous slots, tokens compared with `generate()` on the same sharded model. On 2× NVIDIA H20, torch 2.9.0a0+50eac811a6.nv25.09, Transformers 4.57.6, Python 3.12.3: ```bash CUDA_VISIBLE_DEVICES=0,1 python -m pytest tests/unit/v1/test_hybrid_engine_rollout_v1.py -k TestAutoTPContinuousRollout # baseline (rollout + static cache from d8241a9): 2 failed (eager `4 vs 2`, graph `[2, 2, 16]` vs `[2, 4, 16]`); this PR: 2 passed ``` ```bash DS_ACCELERATOR=cpu python -m pytest tests/unit/runtime/rollout/test_hybrid_engine_rollout.py \ tests/unit/utils/test_static_cache.py tests/unit/v1/test_hybrid_engine_rollout_v1.py # 60 passed, 11 skipped; Python 3.12.3 / torch 2.14.0.dev20260804+xpu pre-commit run --files deepspeed/runtime/rollout/hybrid_engine_rollout.py \ deepspeed/utils/static_cache.py tests/unit/v1/test_hybrid_engine_rollout_v1.py # passed ``` **Intel XPU** (1× Arc Pro B70, torch 2.14.1+xpu, Transformers 5.22.1, Intel compute runtime 26.27.39122.11, Level Zero loader 1.28.2). Master has no XPU graph backend, so the new graph tests skip there. On top of deepspeedai#9007 (XPU graph capture): ```bash DS_ACCELERATOR=xpu python -m pytest tests/unit/v1/test_hybrid_engine_rollout_v1.py \ -k "static_graph_uses_prefill_kv_shape or continuous_generation_uses_local_kv_shape" # deepspeedai#9007 alone: 6 failed, 3 passed; deepspeedai#9007 + this PR: 9 passed DS_ACCELERATOR=xpu python -m pytest tests/unit/v1/test_hybrid_engine_rollout_v1.py \ tests/unit/runtime/rollout/test_hybrid_engine_rollout.py # deepspeedai#9007 + this PR: 60 passed, 3 skipped; process exit 0 ``` This clean-exit torch 2.14.1 rerun confirms all nine new tests pass; the same tests against deepspeedai#9007 without this PR give 6 failed / 3 passed (exit 1). The earlier torch 2.15.0+xpu run printed passing tests but aborts during interpreter teardown on this host (exit 134). A minimal XPU tensor program reproduces that abort without DeepSpeed, so that earlier run is not counted as a clean test pass. The earlier master-based tree reported 52 passed / 10 skipped, with graph coverage skipped as expected. The opt-in `test_continuous_graph_matches_pretrained_model_when_enabled` (only runs when `DEEPSPEED_PRETRAINED_ROLLOUT_MODEL` is set) fails the same way on master and on this PR with Transformers 5.17: `'DeepSpeedStaticCache' object has no attribute 'get_query_offset'`, which deepspeedai#8852 addresses. **Validated:** ZeRO-0, no offload, fp32/bf16, AutoTP with data-parallel replicas on CUDA, plain and continuous decode, eager and graph; single-device graph tests on XPU with deepspeedai#9007. **Not validated:** a full training-to-rollout cycle, ZeRO-1/2/3, offload, pipeline/expert/sequence parallelism, older HF preallocated cache APIs, and multi-device XPU graph. No claim is made about training convergence or performance. AI assistance: GitHub Copilot helped with the implementation, the tests and the validation. I reviewed every changed line before signing off. --------- Signed-off-by: Jin, Youzhi <youzhi.jin@intel.com> Signed-off-by: iLeGend <824040212@qq.com> Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Co-authored-by: Ma, Guokai <guokai.ma@gmail.com>
…oups The CPU shm inference_all_reduce_ kernel is initialized with the WORLD topology and takes no group argument, so it can only produce a correct sum for a world-wide reduction. When a proper TP subgroup called inference_all_reduce (autotp_size=2 inside a 4-rank world), the kernel combined the two sibling subgroups' partial sums into a 4-way sum, corrupting outputs and grad norms. This is the cpu-torch-latest-v1-multirank CI failure: TestTpGradNorm [2-1]/[2-2] failed deterministically (base_norm 1.6128 vs tp_norm 2.0052, bit-identical across runs) on avx512bw runners where the shm kernel engages; avx2-only runners silently fell back to gloo and passed, which made the check look flaky run-to-run. Fix: engage the shm kernel only when the requested group spans the whole world; a subgroup with full world size is necessarily the world itself, so a size comparison suffices. Every other group falls through to the group-aware torch.distributed.all_reduce. No other backend or routing change: the accelerator custom-collective path is already bound to its registered group per process and is untouched. Add a regression test that reduces through two 2-rank subgroups of a 4-rank world and asserts subgroup-scoped results (1/5, not the world-sum 6). With the guard removed the test reproduces the CI failure exactly. Validated on CPU (torch 2.12.1+cpu, python 3.11): - unit/v1/autotp/test_autotp_training.py -k TestTpGradNorm: 4 passed, 2 skipped (previously 2 failed, bit-identical to CI) - unit/comm/test_custom_collective.py: 5 passed, 2 skipped - unit/v1/autotp/test_tp_plan_e2e.py: test_tp_plan_basic_training passed - full test_autotp_training.py: no regressions vs the unpatched head Signed-off-by: Ma, Guokai <guokai.ma@gmail.com>
Conflict resolution: - generate path: combine master's ZeRO-3 group-aware stopping criterion (_GlobalEosStoppingCriteria + synced_gpus opt-out) with our sampling kwargs (do_sample/temperature/top_p) in the eager fallback - StaticCache construction: keep our TP-sharded config copy and eager init (graph capture needs pre-allocated buffers); master's simplified lazy-init call site only fits its non-TP default path - GDN slot binding comment + prefill-shape inference comment both kept - test_static_cache: keep both our GDN-protocol test and master's memory_bytes test (pure both-add) Signed-off-by: Guokai Ma <guokai.ma@intel.com>
The TP guard requires the full blocks (the injected forward slices per-rank head ranges); the old column specs tripped it after the merge. Signed-off-by: Guokai Ma <guokai.ma@intel.com>
…kwargs The fallback prefill branch of _continuous_prefill_cache passed do_sample=False explicitly while also unpacking **_hf_sampling_kwargs(sampling), whose return value always contains do_sample -- so every execution of that branch raised "GenerationMixin.generate() got multiple values for keyword argument 'do_sample'". The explicit kwarg is a leftover from the greedy-only era (e53e2dc); dropping it lets _hf_sampling_kwargs own the value, which also keeps first-token selection consistent with the static-cache path's _token_selector(sampling). Surfaced by TestAutoTPContinuousRollout::test_matches_generate_on_sharded_llama [False] from #9008 after the 2026-11-13 master merge: it drives the no-cache-class fallback branch that our own tests (static-cache models) never reach. Also collapse a stray blank line in tests/unit/utils/test_static_cache.py flagged by the Formatting check (yapf). Validated locally: yapf repo-clean, pre-commit green on changed files, unit/comm and test_static_cache unchanged vs baseline; the rollout test itself requires the CI transformers environment. Signed-off-by: Ma, Guokai <guokai.ma@gmail.com>
| # workload at higher batch. | ||
| eos_mask = torch.zeros(batch_size, dtype=torch.bool, device=device) | ||
| for step in range(max_new_tokens - 1): | ||
| if eos_mask.all(): |
There was a problem hiding this comment.
Suggestion (perf, b>1): if eos_mask.all(): is a device→host sync on every decode step. The host has to wait for the replay plus select_next_tokens to finish before it can enqueue the next step's copy_/fill_/mask update/replay. As a result, the ~50 µs/step of Python work mentioned above runs in series with the GPU instead of overlapping it, and the GPU sits idle between replays. It also syncs when eos_token_id is None, where the mask can never become true.
The b=1 path below already avoids this with eos_check_every = 16. The same approach should work here with identical output, because generate() passes the result through _pad_after_eos (L339), which pads everything after each row's first EOS:
eos_check_every = 16
for step in range(max_new_tokens - 1):
... # copy_/fill_/mask/replay/select, unchanged
output_ids.append(next_token)
if eos_token_id is None:
continue
eos_mask |= (next_token.view(batch_size) == eos_token_id)
if step % eos_check_every == eos_check_every - 1 and bool(eos_mask.all()):
remaining = max_new_tokens - 2 - step
if remaining > 0:
output_ids.append(torch.full((batch_size, remaining), pad_token_id, dtype=torch.long, device=device))
breakAt most eos_check_every - 1 extra steps run after the last row finishes, and their tokens get padded anyway. The Python fallback loop at L1780 could use the same pattern, but it matters less there.
|
|
||
| op = self._ki_norm_op | ||
| residual = hidden_states | ||
| hidden_states = self.input_layernorm(hidden_states) |
There was a problem hiding this comment.
input_layernorm (and the final model.norm) still run HF's eager Qwen3_5RMSNorm: x.float(), pow, mean, +eps, rsqrt, two muls, 1 + w.float(), cast back, which is about 10 launches per norm. In a profile of the b=1 and b=64 decode graphs, these 33 norms are one of the larger remaining non-GEMM costs, since only the post-attention norm goes through fused_add_norm.
Two options, cheapest first:
- Replace the eager norm with
F.rms_norm(x.float(), (hidden,), scale_fp32, eps).type_as(x), wherescale_fp32 = 1 + w.float()is a cached buffer (see thepostattn_scalecomment below). That's 3 launches instead of about 10. On H20 / torch 2.13 the output is bitwise-equal to HF at b=1 and within 1 bf16 ulp at b=64, because the reduction order differs. That is the same class of differencefused_add_normalready has. - Fuse across the layer boundary: layer i's trailing
hidden_states + new_residual(L1137) followed by layer i+1'sinput_layernormis exactly thefused_add_normpattern. That removes both the add and the eager norm, but the patched layer would have to hand the un-added pair to the next layer, which is a larger change to the HF loop contract.
Option 1 alone looks like a self-contained win.
| # other batches or no kernel: the composite reference (gradients also | ||
| # correct there); measured faster than the batched GEMV kernels | ||
| from deepspeed.module_inject.kernel_reference import grouped_linear | ||
| glu = grouped_linear(input, [self.gate_proj.weight, self.up_proj.weight], True)[0] |
There was a problem hiding this comment.
For m > 32, this bypasses the native op and calls kernel_reference.grouped_linear, which is the spec/fallback. That path does two separate F.linear calls plus an fp32 decomposed SiLU (exp, add, div, two muls, and the casts), so the epilogue alone is about 6 elementwise launches per layer.
The native grouped_linear already owns the large-M choice (fused_glu.cu L1624–1649). On GPUs with >100 KB opt-in smem (Ampere/Hopper, grouped_wmma_threshold() == 1) it does one GEMM over the cached [gate; up] concat plus silu_mul. On Ada it falls back to the aten composite. So routing every m > 1 through self._ki_dual_op when it is set, and keeping the reference only for the no-kernel case, should reduce both GEMMs and the epilogue at b=64.
This also has a parity benefit. I checked on H20 with random bf16 inputs: the reference epilogue (g / (1 + exp(-g)) * u) in fp32 is not bitwise-equal to HF's F.silu(gate) * up, while the native composite at::silu(g.float()).to(bf16) * u is. So today b ≤ 32 and b > 32 round the MLP differently. If the native op must stay off for m > 32, F.silu(g) * u in bf16 would at least match HF and save most of the epilogue launches.
| else: | ||
| fused = GDNInputProj.apply(hidden_states, self.in_proj_qkv.weight, self.in_proj_z.weight, | ||
| self.in_proj_b.weight, self.in_proj_a.weight, gdn_kernel_op) |
There was a problem hiding this comment.
At TP=1, b>1 decode lands in the native grouped_linear with n=4. For M > grouped_wmma_threshold() that is 4 separate matmuls plus an at::cat (fused_glu.cu L1652–1661). The b/a projections only have num_v_heads (32) rows each, so they are two GEMMs per layer that are almost pure launch cost, plus a cat that copies the full [qkv | z | b | a] output.
The TP>1 branch above already solves this with _ki_tp_cat_w (two groups, refreshed by segki_refresh). Possible variants for TP=1:
- Cheapest: cache only
cat([in_proj_b.weight, in_proj_a.weight])(2 * num_v_heads × hidden, about 0.3 MB per layer). That takes 4 GEMMs down to 3. - Same as TP>1: cache
cat([z, b, a])and callapply6(x2, qkv_w, qkv_w, qkv_w, zba, zba, zba), which gives 2 GEMMs and one code path for TP=1 and TP>1. The cost is about(value_dim + 2 * num_v_heads) × hiddenbf16 per GDN layer, about 0.5 GB total for Qwen3.5-4B.
Either way the cat stays unless the GEMMs write into slices of one preallocated output.
| # directly. During capture the add becomes part of the graph, so replay | ||
| # reads the parameter's memory and reflects post-training values. | ||
| new_residual = hidden_states | ||
| postattn_scale = torch.add(1.0, self.post_attention_layernorm.weight).contiguous() |
There was a problem hiding this comment.
Minor: 1 + w is one extra launch per layer per step (32 per decode step), and the same pattern would apply to the input norm if it moves to F.rms_norm.
I see the comment about picking up post-training values through the captured add. An alternative that keeps that property: compute it once into a persistent buffer, and have segki_refresh update it in place (buf.copy_(1 + w)) instead of rebuilding it. Captured graphs read the same address, so replay still sees the new values.
The same in-place idea would matter for the other caches segki_refresh currently delattrs (_ki_tp_cat_w, _ki_attn_cat_w, g_glu_cat_cache) once the decode graph is reused across generate() calls (#9077 / #9009). Rebuilding them gives new storage, and an already-captured graph would keep reading the stale copy.
| rotary = getattr(getattr(module, "model", module), "rotary_emb", None) | ||
| if rotary is not None: | ||
| with torch.no_grad(): | ||
| positions = torch.arange(max_len, device=device)[None, :] |
There was a problem hiding this comment.
On transformers 5.x, Qwen3.5's rotary_emb is MRoPE (mrope_section) and expects position_ids of shape [3, bs, L]. The 2-D positions here make rotary(...) raise, and the except Exception below swallows it, so install_decode_attention / install_fused_norm / install_lm_head are all silently skipped and decode falls back to the HF SDPA path. For text-only rows the three sections share the same positions, so the resulting [L, head_dim] tables are unchanged:
| positions = torch.arange(max_len, device=device)[None, :] | |
| positions = torch.arange(max_len, device=device)[None, :] | |
| if hasattr(rotary, "mrope_section"): | |
| # MRoPE takes [3, batch, len] positions; text-only rows are all equal. | |
| positions = positions[None].expand(3, 1, -1) |
Also suggest logging a warning (or re-raising) in the except below instead of falling back silently, so a failed install is visible.
|
Two more decode-path suggestions, both specific to the segKI kernels in this PR. I compared the decode step against vLLM on Qwen3.5-4B (H20). At larger batch sizes most of the remaining kernel gap is in these two places. 1. Full attention: use FA3 (or FlashInfer) for decode instead of At b>1,
Like the current kernel, this only applies to unpadded prompts. Padded prompts still fall back to SDPA. 2. GDN decode: use a packed fused-recurrent kernel The current path launches FLA 0.6's
What I tried: allocate the recurrent state as Smaller item: the conv update still uses causal-conv1d CUDA. vLLM's Triton conv update is faster at large batch, but I have not ported it (its state semantics differ from the HF With both changes, greedy outputs are unchanged at b1/b4/b16. At b64 the tokens drift late in the sequence, which is consistent with bf16 reduction-order differences. Reference implementation, prototype quality and opt-in via env flags: jinyouzhi@be8c2d886 ( Related items that are independent of segKI went to separate PRs: graph reuse across calls (#9077) and running the whole decode step inside the graph (#9079). |
tohtana
removed
the
needs-gpu-test
Thanks @jinyouzhi for the investigation. I'll test your suggestions. And these new changes would be added in separate PRs you opened. |
|
Thanks! @delock For the relatively independent or generic parts, I've already opened separate PRs. But for changes that are coupled to this PR itself, it's probably better if you lead them. The comments and prototype are for reference only, so I prefer no need for a separate PR. |
…eedai#8997) ## Summary Adds a latency-optimized one-shot all-reduce for P2P-capable CUDA pairs (TP=2), filling the existing `deepspeed.comm.inference_all_reduce` hook — the same slot the CPU backend fills with shm+AVX. **Split out of deepspeedai#8602**: this is a generic comm-layer optimization with zero dependency on segment-KI; any AutoTP TP=2 workload benefits, including plain `HF generate` whose decode steps all-reduce small hidden-state tensors through the TP-wrapped linears. ### How it works - `csrc/inference/custom_ar.cu`: one-shot P2P kernel over CUDA IPC buffers (2x4MB, sequence-numbered flags) - `deepspeed/comm/custom_ar.py`: topology gating — world==2, P2P capable, bounded-spin handshake probe; **symmetric verdict** via `all_reduce(MIN)` so a unidirectional fabric fault disables the fast path on both ranks (prevents the peer-spin hang) - `op_builder/custom_ar.py` (`DS_BUILD_CUSTOM_AR`): built through the standard op-builder path like every other op (wheel-prebuildable), initialized next to the shm op in `TorchBackend`, mirroring `build_shm_op` - `layers.py`: `RowParallel` routes no-grad small-message decode through `inference_all_reduce` (falls back to NCCL unchanged); inference mode uses reduce-scatter+all-gather - Fallbacks are safe and logged: non-CUDA, world≠2, no P2P, oversized (>4MB), or build failure → NCCL ## Not a duplicate Split from deepspeedai#8602 (segment-KI), which remains the umbrella PR. No other open PR implements a CUDA fast path for the `inference_all_reduce` hook (CPU shm path is `ShareMemCommBuilder`, untouched here). ## Tests `tests/unit/comm/test_custom_ar.py` — routing contract (4 tests, CPU-runnable): - not-ready before initialize; initialize() no-ops without a process group - payload gate is zero when not ready (no message can ever qualify) Run: `pytest tests/unit/comm/test_custom_ar.py tests/unit/comm/ -q` → 38 passed, 28 skipped on CPU (local, Python 3.11, CPU only). **On-device validation**: the kernel + handshake path was exercised on 2x GPU (AutoDL, see deepspeedai#8602 commit `4bcc86b82` verification) through the ZeRO-3→arena→TP2 rollout e2e. A dedicated 2-GPU parity test (custom_ar vs NCCL all-reduce) is planned as follow-up; reviewer may request hardware specifics from deepspeedai#8602's logs. ## Compatibility matrix Validated: CPU (comm unit suite), CUDA TP=2 (via deepspeedai#8602 e2e). Not validated: XPU/HPU/NPU (code routes through `get_accelerator()` and falls back; only `can_device_access_peer` is CUDA-specific, marked `#ignore-cuda`). ZeRO/offload unaffected (inference-path only). ## AI assistance This PR was prepared with AI assistance (code assembly, builder wiring, contract tests, review-hardening fixes — symmetric handshake verdict, `get_accelerator()` routing). The human submitter reviewed every line; the handshake-hang analysis and the split decision originated from human review discussion. --------- Signed-off-by: Ma, Guokai <guokai.ma@intel.com> Signed-off-by: Guokai Ma <guokai.ma@intel.com> Signed-off-by: Ma, Guokai <guokai.ma@gmail.com> Co-authored-by: pengdurice <pengduhit@gmail.com>
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.