Skip to content

fix(engine): resolve moe-strategy auto to fused on unified-memory GPUs (GB10) - #445

Open
iamanishx wants to merge 1 commit into
FlashML-org:mainfrom
iamanishx:fix/unified-memory-auto-fused
Open

fix(engine): resolve moe-strategy auto to fused on unified-memory GPUs (GB10)#445
iamanishx wants to merge 1 commit into
FlashML-org:mainfrom
iamanishx:fix/unified-memory-auto-fused

Conversation

@iamanishx

Copy link
Copy Markdown

Summary

  • Detect unified-memory GPUs via cudaDevAttrIntegrated (_is_unified_memory_gpu()),
    with a FREETOKEN_UNIFIED_MEMORY=0/1 override for platforms where the attribute lies.
  • On unified memory, resolve --moe-strategy auto to fused instead of the offload
    family when the model's expert format supports residency (_fused_resident_ok()),
    and skip the benchbw hybrid upgrade since CPU execution adds no bandwidth when CPU
    and GPU share one memory.
  • Warn when --moe-strategy offload is explicitly requested on a unified-memory GPU.

Problem

The auto resolution always picks the offload family because on discrete GPUs a wrong
"it fits" guess is a load-time OOM. On unified-memory parts (GB10/DGX Spark, sm_121)
that rationale inverts: there is no host/device boundary, so the offload path's pinned
staging and slot gather are DRAM-to-DRAM copies, and the slot cache double-allocates
memory the GPU could read directly. Measured on GB10 in #22: offload (the default)
added a fixed 126-141 s stall to every request with ~360 s boot, while fused booted in
~50 s and served normally. Discrete-GPU behavior is unchanged: the new branch only
fires when is_integrated is set, which discrete cards never report.

Validation

Tested on a rented DGX Spark (GB10, sm_121, aarch64, driver 610.57.04, CUDA 13.0):

  • torch.cuda.get_device_properties(0).is_integrated == 1 on GB10.
  • pytest tests/engine tests/moe tests/kernels: 455 passed, 4 skipped.
  • New tests/engine/test_moe_strategy_uma.py: 10 passed.
  • Served Qwen/Qwen3.6-35B-A3B-FP8 with no --moe-strategy flag: log shows
    "Unified-memory GPU detected; auto-selecting 'fused'", 31.4 GB experts loaded
    resident, CUDA graphs captured at bs 1-2, API ready in 54 s.
  • OpenAI-compatible generation verified: warm single request 15 completion tokens
    in 0.357 s; two concurrent requests in 1.86 s.

Refs: #369, #22, roadmap #79 (DGX Spark row).

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.

1 participant