[xijiaat]

📌 Description

For contiguous BF16 inputs to the SM90 MXFP4 pull backend, safe quantization currently produces temporary FP8/scales and then stages them and the routing metadata into symmetric buffers. Write the existing safe quantizer's results directly to those buffers, copy routing in the same kernel, and clear inactive route IDs on every replay.

The change reuses the existing quantization arithmetic, including tiny-value handling and rounding. It keeps the Torch FP32 cast/abs/amax, weight representation, wire format, public API, and expert-parallel protocol. Other dtypes/layouts, prequantized input, and empty input retain the existing paths. The new H/top-k/device cache is warmed by the existing full-forward warmup before graph capture.

This is a fixed-overhead optimization: complete-MoE gains are largest for small batches and diminish as GEMM/communication dominate. It does not establish an end-to-end model speedup.

🔍 Related Issues

Builds on #4843 (Hopper MXFP4 MegaMoE and safe quantizer, by @ruoqianguo and contributors). #6465 tunes GEMM geometry; this change only fuses safe post-amax staging. #4449 concerns a different backend/staging path. Targeted open/merged PR searches found no equivalent change to this SM90 MXFP4 safe staging path.

🚀 Pull Request Checklist

Thank you for contributing to FlashInfer! Before we review your pull request, please make sure the following items are complete.

✅ Pre-commit Checks

• I have installed pre-commit by running pip install pre-commit (or used your preferred method).
• I have installed the hooks with pre-commit install.
• I have run the hooks manually with pre-commit run --all-files and fixed any reported issues.

If you are unsure about how to set up pre-commit, see the pre-commit documentation.

🧪 Tests

• Tests have been added or updated as needed.

• All tests are passing (unittest, etc.). Full suite not run; targeted results below.

• Final PR source: pytest -q tests/moe_ep/test_sm90_mxfp4_fused_staging.py — 4 passed on H20-3e. Tests compare exact payload/scale/routing bytes, guard regions, shrinking and empty batches, changed/restored graph inputs, and the fallback path.

• Additional qualification: 20 shapes / 128 exact staging checks; 30 nonempty staging timing cases, all retained. Hidden sizes 4096/5120/7168, M=1/4/16/128/1024, varying capacity and top-k. Staging latency reduction 4.75–41.11%; this is not the complete-MoE improvement.

• Four-rank complete MoE: all 7 cases on both geometries passed exact finite-output comparison for original/changed/restored inputs, with 3 graph replays per phase. Includes skewed routing and an empty rank.

• Three rounds of 20 AB/BA pairs per case. Each timing graph contains an unmeasured NCCL alignment gate followed by CUDA events around five complete forward calls; all actual MoE communication is timed. Reported samples use the maximum rank latency. Compilation, validation and warmup are untimed. Both arms share weights, workspace and the same fixed upstream K128 qualification tactic; this is not a best-autotuned backend comparison.

• A/A control on H4096/E64/K4: apparent round gains ranged from -0.031% to +0.238%. The +0.122% large-M result below is small and should not be interpreted as a meaningful improvement without further noise controls.

• Separate NCU capture at M16/H5120/top-k6/capacity48: compute kernels 6→4, with four intermediate copy operations also removed. Quantizer/fused kernel both use 22 registers/thread; achieved occupancy 43.62%→43.64%. NCU replay times are not benchmark times. The main MoE kernel was not profiled, so no compute/HBM saturation claim is made.

• Targeted pre-commit hooks passed for all three changed files, including Ruff and mypy. The full repository test suite and pre-commit --all-files were not run.

Environment: 4×H20-3e, SM90; torch 2.13.0+cu130; CuTe DSL 4.7.1; NVSHMEM 3.8.0 / Python 0.4.1. GPU benchmark checkout was 11f41575e13e53f89afd39a5e017577be68ed46e; the two modified files are byte-identical at PR base 8256ac5d070268e03a4dfa4a44ededcf041e291a. Both arms also used the already-existing native cleanup fix from #6517. Final PR source and benchmark prototype have identical ASTs (formatting only).

Both geometries use I=2048, EP4, capacity=1056, canonical packed E2M1 weights and K32 E8M0 scales. All 60 samples per arm/case are retained; the table uses their combined median, and includes every round's result. These are generic synthetic layer inputs, not V4.1 model-quality or model-throughput measurements.

H / experts / top-k	M per rank / routing	Baseline median (us)	Fused median (us)	Throughput gain	Three round gains
4096 / 64 / 4	1 / uniform	163.757	156.064	+4.929%	+4.847%, +5.020%, +4.874%
4096 / 64 / 4	4 / uniform	215.747	208.682	+3.386%	+3.397%, +3.410%, +3.373%
4096 / 64 / 4	16 / uniform	286.970	279.459	+2.687%	+2.689%, +2.600%, +2.711%
4096 / 64 / 4	128 / uniform	417.149	408.506	+2.116%	+2.147%, +2.086%, +2.123%
4096 / 64 / 4	1024 / uniform	1962.531	1952.963	+0.490%	+0.456%, +0.486%, +0.488%
4096 / 64 / 4	16 / skew	171.158	163.421	+4.735%	+4.763%, +4.819%, +4.666%
4096 / 64 / 4	4 / empty-rank	215.846	208.237	+3.654%	+3.755%, +3.600%, +3.710%
7168 / 384 / 8	1 / uniform	316.346	309.133	+2.333%	+2.318%, +2.268%, +2.396%
7168 / 384 / 8	4 / uniform	744.829	737.565	+0.985%	+0.996%, +0.982%, +0.981%
7168 / 384 / 8	16 / uniform	1769.280	1761.664	+0.432%	+0.421%, +0.443%, +0.432%
7168 / 384 / 8	128 / uniform	2244.717	2236.358	+0.374%	+0.393%, +0.382%, +0.392%
7168 / 384 / 8	1024 / uniform	7191.859	7183.126	+0.122%	+0.148%, +0.133%, +0.071%
7168 / 384 / 8	16 / skew	462.365	454.502	+1.730%	+1.798%, +1.727%, +1.762%
7168 / 384 / 8	4 / empty-rank	655.616	648.154	+1.151%	+1.131%, +1.166%, +1.142%

All-sample ranges

H / experts	M / routing	Baseline min–max (us)	Fused min–max (us)
4096 / 64	1 / uniform	162.752–164.467	155.136–157.702
4096 / 64	4 / uniform	215.072–217.331	208.077–210.675
4096 / 64	16 / uniform	286.406–287.885	278.554–280.211
4096 / 64	128 / uniform	416.192–418.022	407.725–409.280
4096 / 64	1024 / uniform	1961.421–1963.738	1949.843–1955.021
4096 / 64	16 / skew	170.496–172.282	162.854–164.070
4096 / 64	4 / empty-rank	215.251–216.512	207.450–208.838
7168 / 384	1 / uniform	315.648–317.293	308.646–309.888
7168 / 384	4 / uniform	744.326–745.440	736.640–738.112
7168 / 384	16 / uniform	1768.128–1773.722	1760.666–1766.118
7168 / 384	128 / uniform	2228.352–2253.120	2219.565–2244.314
7168 / 384	1024 / uniform	7179.353–7206.937	7168.109–7197.939
7168 / 384	16 / skew	461.312–464.224	453.510–456.019
7168 / 384	4 / empty-rank	654.694–657.792	647.526–650.163

🔬 Experimental Track

• This PR is experimental: it adds or changes code under flashinfer/experimental/ and/or an @flashinfer_experimental_api. Tracking issue: #
  • The tracking issue names an owner, the reason for the experimental path, and a graduation plan with a target release.
  • Core changes are limited to a thin entry point (signature, shared validation, feature-gate check, backend selection, handoff).
  • Tests live in tests/experimental/ and were validated on the intended hardware; a runnable example is included.
  • Nothing is registered in flashinfer/aot.py, and no experimental backend is reachable from backend="auto" without FLASHINFER_ALLOW_EXPERIMENTAL_AUTO_BACKENDS=1. (Calling an @flashinfer_experimental_api or naming a backend explicitly is itself the opt-in and needs no environment variable.)
  • Test scope declared below. The experimental CI lane runs exactly these targets, so keep them as narrow as the change allows.

# One target per line: a directory or a file. (A pytest ::selector is not
# supported -- the sharding runner cannot consume one.) Must be under
# tests/experimental/ and must exist. Delete these comment lines and add yours, e.g.
#
#   tests/experimental/test_my_backend.py
#   tests/experimental/my_backend/
#
# Declaring the whole tree (tests/experimental/) is allowed but means every
# experimental PR pays for every other feature's tests, in every matrix cell.

Reviewer Notes

The H4096 A/B run completed all measurements and correctness checks, then hung during graph/workspace teardown and was stopped by its bounded controller. Its cleanup verified all four GPUs idle. The standalone reproducer explicitly resets captured graphs before destroying the layer; the A/A control and H7168 A/B used this cleanup correction and exited normally. Computation and timing are unchanged.

No extra weight representation is allocated. Quantitative whole-layer peak-memory savings were not isolated, so none are claimed. Timing is warm-cache and uses the same tactic in both arms; it does not prove an advantage over Humming, another backend, or a fully tuned model. At H7168/E384/M1024 the gain is only 0.122%, so this PR should be assessed primarily as a small-batch overhead reduction.

AI assistance (Codex) was used for implementation, validation, and this report.

# Run from the baseline checkout with its dependencies installed; candidate files
# are taken from this PR. Four idle, assigned Hopper GPUs are required.
KERNEL=flashinfer/moe_ep/backends/mega/kernel/sm90/fp8_mxfp4_bf16_pull_cutedsl
mkdir -p candidate
git show 1a666e0a3c803a5d2e2d91f634d31bc17e290e39:$KERNEL/safe_quantize.py > candidate/safe_quantize.py
git show 1a666e0a3c803a5d2e2d91f634d31bc17e290e39:$KERNEL/staging.py > candidate/staging.py
# Save the script below as reproduce_fused_staging.py.
torchrun --nnodes=1 --master-addr=127.0.0.1 --master-port=29741 --nproc-per-node=4 reproduce_fused_staging.py --candidate candidate --out results-small --geometry small --experts 64
torchrun --nnodes=1 --master-addr=127.0.0.1 --master-port=29741 --nproc-per-node=4 reproduce_fused_staging.py --candidate candidate --out results-large --geometry large --experts 384
torchrun --nnodes=1 --master-addr=127.0.0.1 --master-port=29741 --nproc-per-node=4 reproduce_fused_staging.py --candidate candidate --out results-aa --geometry small --experts 64 --control

Standalone complete-MoE reproducer (writes every rank/sample)

"""Same public FlashInfer MoEEpLayer, weights, recipe and tactic; staging A/B only.

Run with torchrun --nnodes=1 --master-addr=127.0.0.1 --master-port=29741 --nproc-per-node=4. Compilation and validation
precede CUDA-event timing. No model, weight conversion or backend comparison.
"""
import argparse
import dataclasses
from datetime import timedelta
import gc
import importlib
import json
import os
from pathlib import Path
import statistics
import time

import hashlib
import importlib.util
import sys
import types

PKG = 'flashinfer.moe_ep.backends.mega.kernel.sm90.fp8_mxfp4_bf16_pull_cutedsl'


def sha(p):
    return hashlib.sha256(Path(p).read_bytes()).hexdigest()


def load_stagers(candidate):
    baseline = importlib.import_module(PKG + '.staging')
    name = PKG + '._fusion_candidate_safe'
    spec = importlib.util.spec_from_file_location(name, candidate / 'safe_quantize.py')
    safe = importlib.util.module_from_spec(spec)
    sys.modules[name] = safe
    spec.loader.exec_module(safe)
    module = types.ModuleType(PKG + '._fusion_candidate_stage')
    module.__package__ = PKG
    module.__file__ = str(candidate / 'staging.py')
    source = Path(module.__file__).read_text()
    source = source.replace('from .safe_quantize import', 'from ._fusion_candidate_safe import')
    exec(compile(source, module.__file__, 'exec'), module.__dict__)
    return baseline.stage_mega_moe_inputs, module.stage_mega_moe_inputs


def save(p, obj):
    tmp = p.with_suffix('.tmp')
    tmp.write_text(json.dumps(obj, indent=2, allow_nan=False) + '\n')
    tmp.replace(p)





def main():
    p = argparse.ArgumentParser()
    p.add_argument('--candidate', type=Path, required=True)
    p.add_argument('--out', type=Path, required=True)
    p.add_argument('--geometry', choices=['small', 'large'], required=True)
    p.add_argument('--experts', type=int, default=64)
    p.add_argument('--control', action='store_true')
    a = p.parse_args()
    import torch
    import torch.distributed as dist
    from flashinfer.moe_ep import BootstrapConfig, FleetParams, MegaConfig, MoEEpLayer, MoEEpTensors, Sm90_Fp8_Mxfp4_Bf16_PullCutedsl_MegaMoeConfig
    from flashinfer.moe_ep.weights import PrequantizedMoEWeights
    rank = int(os.environ['LOCAL_RANK'])
    torch.cuda.set_device(rank)
    assert torch.cuda.device_count() == 4 and torch.cuda.get_device_capability() == (9, 0)
    dist.init_process_group('nccl', timeout=timedelta(minutes=15), device_id=torch.device('cuda', rank))
    group = dist.new_group(backend='gloo', timeout=timedelta(minutes=15))
    assert dist.get_world_size() == 4
    baseline, candidate = load_stagers(a.candidate)
    backend = importlib.import_module(PKG + '.backend')
    assert backend.stage_mega_moe_inputs is baseline
    h, intermediate, topk = (4096, 2048, 4) if a.geometry == 'small' else (7168, 2048, 8)
    experts = a.experts
    assert experts >= topk and experts % 4 == 0
    if a.control: candidate = baseline
    cap = 1056
    a.out.mkdir(parents=True, exist_ok=True)
    dest = a.out / ('rank%d.json' % rank)
    assert not dest.exists()
    torch.manual_seed(20261018 + rank)
    # Canonical packed E2M1 + K32 E8M0. Both arms share the same transformed
    # weights and workspace; no extra weight representation is allocated.
    local_e = experts // 4
    def raw(rows, k):
        return (torch.randint(0, 256, (local_e, rows, k//2), device='cuda', dtype=torch.uint8),
                torch.randint(113, 121, (local_e, rows, k//32), device='cuda', dtype=torch.uint8))
    w1,s1 = raw(2*intermediate,h); w2,s2 = raw(h,intermediate)
    config = Sm90_Fp8_Mxfp4_Bf16_PullCutedsl_MegaMoeConfig(
        intermediate_size=intermediate, top_k=topk, gate_up_clamp=4.0, fast_math=True,
        preprocess_expert_chunk_size=4, swap_ab=True, pingpong=False,
        mma_tiler_mnk=(128,32,128), cluster_shape_mnk=(1,1,1),
        load_balance_mode='static', token_back_mode='epi_warps')
    report = {'status':'PREPARING', 'rank':rank, 'geometry':a.geometry, 'hidden':h,
              'intermediate':intermediate,'experts':experts,'topk':topk,'ep':4,'capacity':cap,
              'control':a.control, 'timing_protocol':'Unmeasured NCCL GPU gate, then captured external events and five complete MoE calls per graph; rank maximum', 'config':dataclasses.asdict(config), 'script_sha':sha(__file__),
              'candidate_sha':{f:sha(a.candidate/f) for f in ('staging.py','safe_quantize.py')},
              'scope':'complete fused MoE including safe BF16 input quantization, staging, EP dispatch, FC1, SwiGLU, A2 quantization, FC2 and combine; fixed same upstream K128 qualification tactic, not best-autotuned claim',
              'cases':[], 'started':time.time()}
    save(dest,report)
    layer = None
    try:
        layer=MoEEpLayer(bootstrap=BootstrapConfig(world_size=4,rank=rank,process_group=dist.group.WORLD,device=rank,auto_bootstrap=True),
                        fleet_params=FleetParams(num_experts=experts,max_tokens_per_rank=cap,token_hidden_size=h),
                        weights=PrequantizedMoEWeights(w13=w1,w2=w2,w13_scale=s1,w2_scale=s2),
                        backend=MegaConfig(megakernel=config,quantize_input=True,preprocess_weights=True))
        report['allocated_after_layer']=torch.cuda.memory_allocated()
        # Include uniform and skewed routes plus a rank with no tokens. Large
        # and small graphs share one capacity; stale route tails must not leak.
        specs=[(m,'uniform') for m in (1,4,16,128,1024)] + [(16,'skew'),(4,'empty-rank')]
        for m,routing in specs:
            t=0 if routing=='empty-rank' and rank==0 else m
            x=torch.randn((t,h),device='cuda',dtype=torch.bfloat16)
            ids=torch.stack([torch.randperm(experts,device='cuda')[:topk] for _ in range(t)]) if t else torch.empty((0,topk),device='cuda',dtype=torch.int64)
            if routing=='skew': ids.copy_(torch.arange(topk,device='cuda').expand(t,-1))
            weights=torch.softmax(torch.randn((t,topk),device='cuda'),dim=-1)
            tensors=MoEEpTensors(hidden_states=x,topk_ids=ids,topk_weights=weights)
            row={'tokens_per_rank_nominal':m,'tokens_this_rank':t,'routing':routing,'status':'PREPARING','checks':[]}
            report['cases'].append(row);save(dest,report)
            backend.stage_mega_moe_inputs=baseline
            layer.warmup(tensors)
            def call(): return layer.forward(tensors,return_workspace_view=True)
            for fn in (baseline,candidate):
                backend.stage_mega_moe_inputs=fn
                for _ in range(3): call()
                torch.cuda.synchronize();dist.barrier(group=group)
            graphs=[]; outputs=[]
            for fn in (baseline,candidate):
                backend.stage_mega_moe_inputs=fn
                dist.barrier(group=group)
                g=torch.cuda.CUDAGraph()
                with torch.cuda.graph(g): y=call()
                graphs.append(g);outputs.append(y)
                torch.cuda.synchronize();dist.barrier(group=group)
            original=(x.clone(),ids.clone(),weights.clone())
            for phase in ('original','changed','restored'):
                if phase=='changed':
                    x.normal_();weights.copy_(torch.softmax(torch.randn_like(weights),dim=-1))
                    if t: ids.copy_((ids+7)%experts)
                elif phase=='restored':
                    x.copy_(original[0]);ids.copy_(original[1]);weights.copy_(original[2])
                snapshots=[]
                for index,g in enumerate(graphs):
                    for _ in range(3): g.replay()
                    torch.cuda.synchronize();snapshots.append(outputs[index].clone());dist.barrier(group=group)
                assert torch.isfinite(snapshots[0]).all() and torch.isfinite(snapshots[1]).all()
                assert torch.equal(snapshots[0],snapshots[1]), (rank,m,routing,phase,'output differs')
                row['checks'].append({'phase':phase,'replays':3,'bitwise_output':True})
            torch.cuda.synchronize();dist.barrier(group=group)
            for _ in range(10):
                for g in graphs: g.replay()
            torch.cuda.synchronize();dist.barrier(group=group)
            # Three complete independent timing rounds, alternating pair order.
            # All rank latency samples are retained; cross-rank maximum is the
            # complete operation's critical latency, reduced outside timing.
            timed_graphs=[]; events=[]; gates=[]
            for fn in (baseline,candidate):
                backend.stage_mega_moe_inputs=fn
                gate=torch.zeros(1,device='cuda',dtype=torch.int32)
                dist.all_reduce(gate);torch.cuda.synchronize();dist.barrier(group=group)
                start=torch.cuda.Event(enable_timing=True,external=True)
                end=torch.cuda.Event(enable_timing=True,external=True)
                g=torch.cuda.CUDAGraph()
                with torch.cuda.graph(g):
                    dist.all_reduce(gate)
                    start.record()
                    for _ in range(5):call()
                    end.record()
                timed_graphs.append(g);events.append((start,end));gates.append(gate)
                torch.cuda.synchronize();dist.barrier(group=group)
            for _ in range(10):
                for g in timed_graphs:g.replay()
            torch.cuda.synchronize();dist.barrier(group=group)
            row['rounds']=[]
            for round_id in range(3):
                samples=[[],[]];rank_samples=[[],[]];orders=[]
                for i in range(20):
                    order=(0,1) if (i+round_id)%2==0 else (1,0);orders.append(order)
                    for index in order:
                        dist.barrier(group=group)
                        start,end=events[index]
                        timed_graphs[index].replay()
                        end.synchronize()
                        elapsed=start.elapsed_time(end)*1000/5
                        values=[None]*4;dist.all_gather_object(values,elapsed,group=group)
                        samples[index].append(max(values));rank_samples[index].append(values)
                row['rounds'].append({'samples_us':samples,'all_rank_us':rank_samples,'order':orders,
                                      'median_us':[statistics.median(s) for s in samples]})
                save(dest,report)
            row['status']='PASS'
            report['memory_allocated_after_case']=torch.cuda.memory_allocated()
            save(dest,report)
            if rank==0: print(json.dumps({'m':m,'routing':routing,'round_medians':[r['median_us'] for r in row['rounds']]}),flush=True)
            for captured in graphs + timed_graphs:
                captured.reset()
            del graphs,timed_graphs,events,gates,outputs,snapshots,original,tensors,x,ids,weights
            gc.collect();torch.cuda.synchronize();dist.barrier(group=group)
        report['status']='PASS'
    except BaseException as e:
        report.update(status='FAIL',error=repr(e));raise
    finally:
        backend.stage_mega_moe_inputs=baseline
        report['finished']=time.time();save(dest,report)
        if layer is not None and report['status']=='PASS':
            report['cleanup_phase']='before_layer_destroy';save(dest,report)
            layer.destroy()
            report['cleanup_phase']='after_layer_destroy';save(dest,report)
        if report['status']=='PASS':
            dist.destroy_process_group(group)
            dist.destroy_process_group()
            report['cleanup_phase']='complete';save(dest,report)


if __name__=='__main__':main()

Summary by CodeRabbit

• New Features
  • Added a fused staging path for eligible MoE inputs that quantizes input data while staging routing IDs and weights.
  • Staging supports varying token counts, including empty batches, and preserves existing behavior for unsupported inputs.
  • Routing outputs beyond the active token count are cleared to -1.
  • CUDA Graph replays support changing and repeated batch sizes while keeping staged outputs consistent.
  • Staged routing data and quantized inputs are written to provided output buffers, including strided workspaces.

[coderabbitai]

No actionable comments were generated in the recent review. 🎉

ℹ️ Recent review info

⚙️ Run configuration

• Configuration used: defaults
• Review profile: CHILL
• Plan: Advanced
• Run ID: 84f82d28-4336-4a41-9511-3b6b4636bbc9

📥 Commits

Reviewing files that changed from the base of the PR and between bfd3072 and ebfa1a6.

📒 Files selected for processing: 1

• flashinfer/moe_ep/backends/mega/kernel/sm90/fp8_mxfp4_bf16_pull_cutedsl/safe_quantize.py

Included review availability: This review used your included allowance. Your plan provides up to 8 included reviews per hour; 7 remain after this review.

📝 Walkthrough

Safe quantization now stages quantized activations, scales, and routing data. Eligible SM90 inputs use a cached fused staging kernel. New tests cover graph replay, workspace guards, and strided workspaces.

Changes: Fused Safe-Quantization Staging

Layer / File(s)	Summary
Quantizer routing outputs flashinfer/moe_ep/backends/mega/kernel/sm90/fp8_mxfp4_bf16_pull_cutedsl/safe_quantize.py	The quantizer accepts routing buffers. With routing enabled, it copies live IDs and weights, fills extra ID rows with -1, and writes scales to four columns.
Staging compiler and entry point flashinfer/moe_ep/backends/mega/kernel/sm90/fp8_mxfp4_bf16_pull_cutedsl/safe_quantize.py	A cached compiler specializes the staging kernel by hidden width, top_k, and device. quantize_and_stage_from_amax compiles and invokes the kernel on the input device.
SM90 staging integration and validation flashinfer/moe_ep/backends/mega/kernel/sm90/fp8_mxfp4_bf16_pull_cutedsl/staging.py, tests/moe_ep/test_sm90_mxfp4_fused_staging.py	Eligible inputs use fused safe quantization to stage activations and routing data. Tests cover graph replay outputs, guard bytes, and strided workspaces.

Priority: ⬇️ Low

Estimated code review effort: 3 (Moderate) | ~20 minutes

Change: Refactor

Sequence Diagram(s)

sequenceDiagram
  participant Staging as stage_mega_moe_inputs
  participant Entry as quantize_and_stage_from_amax
  participant Compiler as cached staging compiler
  participant Kernel as _SafeE4M3Quantize
  Staging->>Staging: Compute per-token full-hidden absmax
  Staging->>Entry: Pass activations, absmax, and routing buffers
  Entry->>Compiler: Get specialized staging kernel
  Compiler-->>Entry: Return compiled kernel
  Entry->>Kernel: Launch quantization and routing staging
  Kernel-->>Staging: Write quantized activations, scales, IDs, and weights

Loading

Suggested reviewers: ruoqianguo

Merge Risk: ⚪ Minimal · up to ebfa1

No merge-blocking issue is established; merge after normal checks.

Pre-merge checks: 4 · 1

Failed checks: 1 warning

Check name	Status	Explanation	Resolution
Docstring Coverage		Docstring coverage is 20.00% which is insufficient. The required threshold is 80.00%. Docstring coverage is scoped to functions touched by this diff. Analyzed 15 functions across 3 files.	Write docstrings for the functions missing them to satisfy the coverage threshold.

Passed checks: 4 passed

Check name	Status	Explanation
Title check		The title clearly and concisely describes the primary change: fusing safe FP8 quantization with SM90 MXFP4 staging for performance.
Description check		The description follows the required template, explains the change and scope, references related issues, documents tests and validation results, and identifies that the full suite and full pre-commit …
Linked Issues check		Check skipped because no linked issues were found for this pull request.
Out of Scope Changes check		Check skipped because no linked issues were found for this pull request.

Full details: Description check

Explanation

The description follows the required template, explains the change and scope, references related issues, documents tests and validation results, and identifies that the full suite and full pre-commit run were not completed.

• Fix all pre-merge checks with AI

✨ Finishing Touches: 💡 1

🛠️ Fix failing CI checks: 💡

• Commit to this branch
• Create a new PR

🧪 Generate unit tests (beta)

• Create a new PR

• Autofix · Keep fixing CodeRabbit findings and required CI, and resolving merge conflicts

Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out.

❤️ Share

• X
• Mastodon
• Reddit
• LinkedIn

_{Comment @coderabbitai help to get the list of available commands.}

[xijiaat]

Hi @ruoqianguo, since you introduced the SM90 MXFP4 backend and safe quantizer in #4843, could you take a look at this staging optimization when you have time? It writes the existing safe quantizer outputs directly into symmetric buffers and fuses routing staging, preserving the quantization arithmetic and EP protocol. Correctness checks and four-rank H20 complete-MoE benchmark results are included in the description. I would appreciate your feedback on whether this fits the backend design. Thanks!

[ruoqianguo]

Hi @ruoqianguo, since you introduced the SM90 MXFP4 backend and safe quantizer in #4843, could you take a look at this staging optimization when you have time? It writes the existing safe quantizer outputs directly into symmetric buffers and fuses routing staging, preserving the quantization arithmetic and EP protocol. Correctness checks and four-rank H20 complete-MoE benchmark results are included in the description. I would appreciate your feedback on whether this fits the backend design. Thanks!

Thanks @xijiaat , this is a clean, well-scoped change and the write-up is very thorough. I checked it against the #4843 design and it fits. The four-lane x_sf scale wire is preserved. So is the per-stage topk_idx[n:] = -1 tail mask, which the megakernel relies on because it launches over the full padded workspace. The fused path runs the same quantization kernel body, and warmup() compiles the new stager before graph capture.

Before merge, could you double-check that the backend's 4-GPU suite still passes on the PR head (bash tests/moe_ep/run_tests.sh mega_sm90_mxfp4)? The layer path now always takes the fused stager, but the 1-GPU oracle stages with safe_quantization=False and never reaches it. These multirank tests, especially test_fused_mxfp4_tiny_full_output, are therefore the only end-to-end coverage, and PR CI doesn't run them on 4 GPUs.

A non-blocking suggestion: _SafeE4M3Quantize currently switches between quantizer and stager via top_k=0 and four optional None tensors. I'd rather keep the quantizer's signature unchanged and add a separate staging kernel that shares the per-element math through a @cute.jit helper.

Possible follow-up, not for this PR: the remaining Torch cast/abs/amax are exact for finite inputs. Reading BF16 directly and reducing amax in-kernel would therefore stay bit-identical while removing three more launches and the FP32 [T, H] temporaries. That is likely where the remaining large-M gain is.

[xijiaat]

Thanks @ruoqianguo — confirmed on the current PR head bfd3072c362822c58098ddb104d6e22ebba001d3, using 4 × NVIDIA H20-3e (SM90).

bash tests/moe_ep/run_tests.sh mega_sm90_mxfp4

The official suite exited 0, with 33 passed on each of the four ranks: 33 unique cases / 132 rank-case executions, no failures, errors, or skips. In particular:

• All four test_fused_mxfp4_tiny_full_output parameterizations passed on all four ranks: (1e-15, 0), (1e-18, 0), (1e-19, 0), and (1e-35, 55).
• Cross-rank routing and independent global-oracle checks passed, including workspace reuse/reset and boundary cases.
• Both outer-CUDA-Graph cases (tail-off and tail-on) passed, with 16 replays each.

I also ran the staging regression separately on one GPU:

python -m pytest tests/moe_ep/test_sm90_mxfp4_fused_staging.py -v -ra

Result: 4 passed, exit 0, covering changed-input graph replay, workspace guards/reuse, and the strided-workspace fallback.

Environment: PyTorch 2.13.0+cu130, CUDA 13.0, driver 595.71.05, CuTe DSL 4.7.1, NVSHMEM 3.8.0, nvshmem4py 0.4.1, Python 3.12.3. All 10,266 tracked source files matched the frozen PR head before and after testing. No test, kernel, numerical threshold, or safe-quantization setting was changed. I rechecked the GitHub head before posting; it still matches the tested commit.

Official four-rank log excerpts — all four pytest summaries

================== 33 passed, 3 warnings in 282.75s (0:04:42) ==================
================== 33 passed, 3 warnings in 283.19s (0:04:43) ==================
================== 33 passed, 4 warnings in 283.19s (0:04:43) ==================
================== 33 passed, 4 warnings in 283.47s (0:04:43) ==================
### sm90_fp8_mxfp4_bf16_pull_cutedsl mega multirank (Hopper): PASS
all sections passed.

Supplementary staging regression log excerpts

tests/moe_ep/test_sm90_mxfp4_fused_staging.py::test_safe_staging_reuses_workspace_across_graphs[129-2] PASSED [ 25%]
tests/moe_ep/test_sm90_mxfp4_fused_staging.py::test_safe_staging_reuses_workspace_across_graphs[4096-4] PASSED [ 50%]
tests/moe_ep/test_sm90_mxfp4_fused_staging.py::test_safe_staging_reuses_workspace_across_graphs[7168-8] PASSED [ 75%]
tests/moe_ep/test_sm90_mxfp4_fused_staging.py::test_safe_staging_preserves_strided_workspace_fallback PASSED [100%]
======================== 4 passed, 3 warnings in 1.63s =========================

An earlier cold-JIT attempt encountered a host-memory cgroup OOM and subsequently timed out. For the successful run, I prebuilt/loaded the standard JIT module in a separate process, used MAX_JOBS=4 / FLASHINFER_NVCC_THREADS=1, a 256 GiB container memory limit, and explicit release compilation (FLASHINFER_JIT_DEBUG=0). Failed and successful run logs are retained. The successful logs contain deprecation and Torch disk-cache-directory warnings; the test durations above are correctness-test timings, not performance benchmarks.

Thanks also for the API-cleanup and BF16/amax-fusion suggestions. Those will be evaluated separately; this report validates the existing PR head without adding either change.

[xijiaat]

Hi @ruoqianguo, thank you again for your time and helpful feedback. I've added the four-GPU test results above. Whenever you have a chance, please let me know if there's anything else I could clarify or add. I'd be happy to follow up. Thanks!

[ruoqianguo]

Thanks for running the four-GPU suite and sharing the results. I have no further blocking comments. The API cleanup and BF16/amax fusion can remain follow-ups.
I don’t have approval permissions. @jiahanc, could you help trigger CI and review this PR?

[xijiaat]

Hi @jiahanc, when you have a chance, could you please advise whether any additional tests or validation would help with reviewing this PR? I have posted the four-GPU SM90 suite and targeted staging regression results above. I would be happy to add coverage for any specific shapes, edge cases, or configurations you think are missing. Thank you for your time!