Skip to content

[AMD] Support block FP8 MoE on gfx1250 - #3

Draft
yushengsu-thu wants to merge 4 commits into
akao-amd:amd_heliosfrom
yushengsu-thu:agent/mi455-mimo-fp8-moe
Draft

[AMD] Support block FP8 MoE on gfx1250#3
yushengsu-thu wants to merge 4 commits into
akao-amd:amd_heliosfrom
yushengsu-thu:agent/mi455-mimo-fp8-moe

Conversation

@yushengsu-thu

@yushengsu-thu yushengsu-thu commented Aug 6, 2026

Copy link
Copy Markdown

Summary

  • add a gfx1250-only BF16 contraction fallback for 128x128 block-FP8 fused MoE while preserving block scales and FP32 accumulation
  • force num_stages=1 for gfx1250 block-FP8 MoE kernels
  • avoid modulo-wrapped matrix offsets on the gfx1250 block-FP8 path, matching the dense-kernel out-of-bounds workaround
  • add a MiMo-V2.5-Pro-shaped gfx1250 regression (K=6144, TP4 N=512, top-k 8)
  • add a guarded MI455 MiMo-V2.5-Pro smoke/benchmark launcher that requires the Triton path, validates the checkpoint/TP/driver/GPUs, and prepends its own source worktree to PYTHONPATH
  • default the 1.03 TB checkpoint to single-threaded loading to bound per-rank host-memory pressure, with an explicit override available for measured systems

Why

This is stacked on sgl-project#32754. That PR adds the corresponding gfx1250 workaround for dense block-FP8 GEMM, because native FP8 tl.dot faults, multi-stage software pipelining can miscompile, and modulo-wrapped matrix offsets can become out-of-bounds on gfx1250. The generic fused-MoE block-FP8 kernel still used all three patterns, so MiMo-V2.5-Pro's 128x128 FP8 routed experts did not have the same correctness fallback.

The pinned gfx1250 AITER grouped-MoE path does not support this checkpoint's per-128x128 format. The launcher therefore consistently disables AITER/CK and selects Triton for attention and MoE, avoiding AITER weight preshuffling for a Triton consumer.

SGLang's buffered multi-thread loader documents a peak of approximately (workers + 2) * shard size per rank. This checkpoint has 34 roughly 30 GB shards; the default eight workers across four TP ranks can therefore create excessive host-memory pressure during bring-up. The launcher defaults to enable_multithread_load=false until the node's capacity and stability are measured.

Validation

  • python3 -m py_compile python/sglang/kernels/ops/moe/fused_moe_triton_kernels.py test/manual/quant/test_block_fp8.py
  • ruff check --select F python/sglang/kernels/ops/moe/fused_moe_triton_kernels.py
  • shellcheck scripts/ci/amd/mi455_mimo_v25_pro.sh
  • bash -n scripts/ci/amd/mi455_mimo_v25_pro.sh
  • git diff --check
  • launcher metadata/TP/shard success path exercised with a temporary two-shard fixture
  • launcher invoked from outside the checkout to verify source-root resolution and AITER guard ordering

MI455 runtime status

This remains a draft intentionally. The available MI455 host currently fails L2 neighbor resolution from the otherwise reachable AMD bastion (No route to host), so the driver command and GPU numerical test cannot reach the node yet.

Remaining gates:

  • load amdgpu with gpu_recovery=0
  • pass TestW8A8BlockFP8FusedMoE.test_gfx1250_w8a8_block_fp8_fused_moe on gfx1250
  • load MiMo-V2.5-Pro with TP4
  • pass server health and one deterministic generation request
  • run the requested benchmark sweep after correctness

Dependency

Stacked on sgl-project#32754 (akao-amd/sglang:amd_helios). It should be rebased onto sgl-project/sglang:main after that dependency lands.

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