Skip to content

Unified-memory platforms (DGX Spark, sm_121): moe-backend auto resolves to offload, adding ~130s/request of pure overhead #369

Description

@iamanishx

Splitting this out of #22 based on the field data @jlacroix82 posted there.

Problem

On GB10 (sm_121, 121 GB unified memory) --moe-backend auto resolves to the
offload family. But unified memory has no host/device split: the offload path
streams experts from pinned "host" banks into a GPU slot cache, which on this
platform is a DRAM-to-DRAM copy of memory the GPU could have read directly.

Measured in #22 (same box, same checkpoint, idle scheduler):

  • offload (the auto default): fixed 126-141 s stall added to every request,
    ~360 s boot
  • fused (explicit): 0.22-0.77 s probes, 37.8 tok/s decode, ~612 tok/s
    prefill, ~50 s boot

So the default experience on DGX Spark is a multi-minute stall per request,
and the fix is currently tribal knowledge (--moe-backend fused).

Root cause

engine/engine.py (_adjust_config, the moe_backend == "auto" branch):
auto always resolves to offload, with the comment noting it can never pick
fused because nothing knows whether the experts fit in HBM, and a wrong guess
is a load-time OOM. That rationale is sound for discrete GPUs but inverts on
unified memory, where fused residency is the safe default and offload is the
regression.

Proposal

  1. Detect unified-memory GPUs (cudaDevAttrIntegrated, or equivalent) and
    resolve auto to fused there, with a log line stating why. offload
    stays available as an explicit override.
  2. Audit the sm_121 backend probes while in the area. Two known hazards:
    flashinfer's mm_fp8 is sm100a-only and crashes at module load on SM121,
    and [Bug] b12x NVFP4 backend on sm_120 (RTX 5090) serves batch 1 only: not CUDA-graph capturable, and the worker dies with a tile-config ValueError at batch >= 2 (Qwen3.8-Flash-Next-NVFP4) #335 shows the b12x NVFP4 backend is not CUDA-graph capturable and dies
    at batch >= 2 even on sm_120. kernel/backend.py and
    select_nvfp4_backend should not select backends that cannot actually run
    on the detected arch.
  3. Document the GB10 path in docs/install.md: the source-build recipe from
    Feature request: ARM64/aarch64 support for NVIDIA DGX Spark #22 (TVM_FFI_CUDA_ARCH_LIST="12.1"), and why fused is the default there.

Scope

Intentionally narrow: config resolution, probe guards, docs. The deeper
unified-memory execution mode (zero-copy expert reads, rethinking the slot
cache entirely) is a separate design discussion I will open once this lands
and there is benchmark data to anchor it.

Activity

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Metadata

Metadata

Assignees

No one assigned

    Labels

    No labels
    No labels

    Projects

    No projects

      Milestone

      No milestone

      Relationships

      None yet

      Development

      No branches or pull requests

      Issue actions