Repository navigation
[MUSA-60092] Fix DeepSeek-V4 MTP multibatch sparse-prefill OOM - #176
Closed
yeahdongcn wants to merge 7 commits into
Closed
yeahdongcn wants to merge 7 commits into
yeahdongcn wants to merge 7 commits into
Conversation
Collaborator
Author
|
Closing this separate stacked draft. Per review direction, the MUSA-60092 fix will be appended directly to #174 instead. |
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.
Stack
This draft is stacked on #174 (
46550bcc1). Until #174 lands, the GitHub diffalso includes its commits. The MUSA-60092 change itself is the single commit
7a0093da0.Problem
DeepSeek-V4 TP8 + MTP4 can OOM during multi-batch sparse prefill at
--gpu-memory-utilization 0.95. The failing path allocates a padded H64 BF16result (about 512 MiB at 8192 tokens) before copying it into an output buffer
that DSV4 already owns.
Controls reproduced the OOM before #166 and with the page-size experiment, so
this is not a #166 regression. A runtime H8 attempt was rejected because it was
numerically wrong.
Change
caller-owned DSV4 output.
non-capture execution, tensor shape/dtype/layout, disjoint storage, dry-run
semantics, and the validated TVM-FFI ABI. Every pre-launch mismatch uses the
public path; launch failures are not retried.
max_num_batched_tokens * 64 * 512 * 2bytes. For the PR configuration thisis 537,067,520 bytes (0.5002 GiB). User
gpu_memory_utilizationis unchanged.Explicit
kv_cache_memory_bytesremains authoritative and emits a warning.fp8_ds_mla, andlet the FlashMLA owner fill a missing backend hint without overriding an
explicit conflicting backend.
Evidence
Environment: 8x MTT S5000 80 GiB, driver 5.2.0-server, Python 3.10.12,
torch/torch_musa 2.9.1.post1+musa5.2.0s5000, torchada 0.1.76, MATE 0.2.4,
flash_mla 0.2.4+musa, flash_attn_3 0.2.4+musa, image
v0.24.0-ph1-5.2.0-torch2.9.1.post1-20260725.Install was skipped for this Python/build-patch-only delta; the source overlay
used the exact #174 native extension artifact
sha256:94bf688d4b720111cef4e4df7184f59dd9271d3d41303f8255565948852d8cdc.BS4, TP8, MTP4, FlashMLA, FP8 KV, 4K input / 1K output, five warmups followed
by 4 requests:
The latency row is not a standalone performance claim: speculative acceptance
differed between the two single runs (55.9% vs 65.3%). It shows no observed
regression. The KV/memory deltas and request success are direct observations.
Numerics on S5000:
Tests:
66 passed;tests/test_patches.py:71 passed, 2 skipped;postcondition failure is the pre-existing
flash_attn.pytripwire drift.Draft limitations
--disable-custom-all-reducebecause the default custom-AR pathindependently hangs during graph capture on the available 35.88/35.93
hosts. All causal arms used the same setting.
runs are not complete yet.
v0.24.0-devafter DSV4 optimization #174 merges.