Skip to content

Support MXFP4 expert weights in remote prefetch - #29

Open
lyppg wants to merge 1 commit into
MoonshotAI:masterfrom
bytedance-iaas:jxp/prefetch_support_mxfp4
Open

Support MXFP4 expert weights in remote prefetch#29
lyppg wants to merge 1 commit into
MoonshotAI:masterfrom
bytedance-iaas:jxp/prefetch_support_mxfp4

Conversation

@lyppg

@lyppg lyppg commented Aug 11, 2026

Copy link
Copy Markdown
Contributor

Current PrefetchKernel was hardcoded to bf16, so a quantized expert could not be prefetched. This PR adds support for MXFP4 expert weights in remote prefetch with minimum kernel change.

Testing

Measured on 2x GB200, one K3 layer's expert prefetch, got 3.8x less traffic for 3.3x less time:

# on two GB200 nodes, --node_rank=1 on another node
torchrun --nnodes=2 --node_rank=0 --nproc_per_node=4 --master_addr=$MASTER_ADDR --master_port=$PORT benchmarks/bench_prefetch.py --owner-rank 4
MoonEP Prefetch Benchmark (R=8, warmup=5, iters=20, reader=rank0, owner=rank4, num_sms=32)
Config                   E    B    dt   N       H     Hp   SMs  Slots   Data(MB)  Worst(us)  BW(GB/s)   CommBW
---------------------------------------------------------------------------------------------------------------
slots_1                 64   14  bf16   1    3584   3072    32      1      44.04      35.02   1257.60   628.80
slots_2                 64   14  bf16   1    3584   3072    32      2      88.08      63.15   1394.70   697.35
slots_3                 64   14  bf16   1    3584   3072    32      3     132.12      91.04   1451.31   725.66
slots_5                 64   14  bf16   1    3584   3072    32      5     220.20     147.77   1490.20   745.10
ramp_0_3                64   14  bf16   1    3584   3072    32     12     528.48     343.81   1537.13   768.56
mixed                   64   14  bf16   1    3584   3072    32     12     528.48     343.83   1537.07   768.53
heavy                   64   14  bf16   1    3584   3072    32     17     748.68     485.58   1541.85   770.92
full_3x8                64   14  bf16   1    3584   3072    32     24    1056.96     679.37   1555.79   777.90
dense_B8                64    8  bf16   1    3584   3072    32     15     660.60     428.92   1540.15   770.07
epn4_full               32   14  bf16   1    3584   3072    32     12     528.48     344.47   1534.18   767.09
epn16_mixed            128   16  bf16   1    3584   3072    32     24    1056.96     681.79   1550.29   775.14
thin_7168x128           64   16  bf16   1    7168    128    32     12      44.04      35.52   1239.81   619.91
tall_1024x3072          64   14  bf16   1    1024   3072    32     12     150.99     103.36   1460.84   730.42
tiny_512x512            64   16  bf16   1     512    512    32     17      17.83      18.97    939.78   469.89
mxfp4_gate_up           64   16    u8   1    3584   3072    32     12     132.12      93.07   1419.58   709.79
mxfp4_down              64   16    u8   1    3584   3072    32     12     132.12      93.42   1414.32   707.16
mxfp4_scale             64  128    u8   1    3584   3072    32     12       8.26      24.93    331.26   165.63
k3_layer_bf16           64   14  bf16   3    3584   3072    32     12    1585.45    1030.53   1538.48   769.24
k3_layer_mxfp4          64   16    u8   6    3584   3072    32     12     421.13     315.30   1335.65   667.83

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