Skip to content

[MUSA-60092] Fix DeepSeek-V4 MTP multibatch sparse-prefill OOM - #176

Closed
yeahdongcn wants to merge 7 commits into
v0.24.0-devfrom
fix/musa60092-dsv4-mtp-workspace
Closed

yeahdongcn wants to merge 7 commits into
v0.24.0-devfrom
fix/musa60092-dsv4-mtp-workspace

Conversation

@yeahdongcn

@yeahdongcn yeahdongcn commented Aug 8, 2026 •

Copy link
Copy Markdown
Collaborator

Stack

This draft is stacked on #174 (46550bcc1). Until #174 lands, the GitHub diff
also 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 BF16
result (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

  • Keep the correct H64 MATE 0.2.4 model1 kernel, but write directly into the
    caller-owned DSV4 output.
  • Gate the private adapter shim on the exact DSV4 TP8 MTP contract, MUSA device,
    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.
  • Reserve the exact transient H64 workspace before automatic KV-cache sizing:
    max_num_batched_tokens * 64 * 512 * 2 bytes. For the PR configuration this
    is 537,067,520 bytes (0.5002 GiB). User gpu_memory_utilization is unchanged.
    Explicit kv_cache_memory_bytes remains authoritative and emits a warning.
  • Fix contract drift seen in real workers: accept normalized fp8_ds_mla, and
    let 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:

Metric Control Candidate
Successful requests 4/4 4/4
KV-cache tokens 67,772 66,827
Sampled peak memory 80,932 MiB 80,082 MiB
Output throughput 36.72 tok/s 38.00 tok/s
Mean TTFT 3,288.5 ms 3,231.3 ms
Mean TPOT 23.50 ms 20.97 ms
Mean ITL 76.39 ms 76.08 ms

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:

  • direct vs public H64 output/max/LSE: bitwise equal;
  • direct vs reference output max abs: 0.0078125;
  • direct vs reference LSE max abs: 0.00000143.

Tests:

  • focused contract/direct-output/source tests: 66 passed;
  • tests/test_patches.py: 71 passed, 2 skipped;
  • all 108 build patches apply cleanly to pinned vLLM. The verifier's remaining
    postcondition failure is the pre-existing flash_attn.py tripwire drift.

Draft limitations

  • Runs used --disable-custom-all-reduce because the default custom-AR path
    independently hangs during graph capture on the available 35.88/35.93
    hosts. All causal arms used the same setting.
  • BS2, BS16/64, Qwen regression smokes, and repeated seeded A-R-A performance
    runs are not complete yet.
  • Rebase onto v0.24.0-dev after DSV4 optimization #174 merges.

@yeahdongcn

Copy link
Copy Markdown
Collaborator Author

Closing this separate stacked draft. Per review direction, the MUSA-60092 fix will be appended directly to #174 instead.

@yeahdongcn yeahdongcn closed this Aug 8, 2026
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants