Skip to content

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

Merged
jason-fxz merged 2 commits into
FlashML-org:mainfrom
iamanishx:fix/unified-memory-auto-fused
Sep 15, 2026
Merged

jason-fxz merged 2 commits into
FlashML-org:mainfrom
iamanishx:fix/unified-memory-auto-fused

Conversation

@iamanishx

Copy link
Copy Markdown
Contributor

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).

@iamanishx

Copy link
Copy Markdown
Contributor Author

@jason-fxz please have a look on this

@jason-fxz

jason-fxz commented Sep 14, 2026

Copy link
Copy Markdown
Collaborator

@iamanishx

Thanks for the PR.

Two requests before merging:

  1. Please add a FIXME to _fused_resident_ok that names the formats the gate admits today, so it gets dropped once the other quant formats support fused, e.g.:
def _fused_resident_ok(model_config) -> bool:
    """Whether the resident ('fused') MoE path can hold this model's experts.

    FIXME: auto resolves to fused only for bf16 and fp8_block experts; drop this gate once the other quant formats support fused.
    """
  1. Please drop tests/engine/test_moe_strategy_uma.py; the monkeypatch in tests/moe/test_offload.py is right and enough, keep that one.

Signed-off-by: iamanishx <manishbiswal754@gmail.com>
@iamanishx

Copy link
Copy Markdown
Contributor Author

@jason-fxz i have made the changes

@jason-fxz

Copy link
Copy Markdown
Collaborator

Reviewed the update, LGTM.

@jason-fxz
jason-fxz merged commit bea8d06 into FlashML-org:main Sep 15, 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