[DSV4] Add BF16 and MXFP8 A2A support for flashinfer a2a one sided

[zyongye]

Purpose

Originally Flashinfer one sided a2a only supports nvfp4 dispatch. Add BF16 and MXFP8 dispatch.

Test Plan

gsm8k on V4-Flash

Test Result

|Tasks|Version|     Filter     |n-shot|  Metric   |   |Value |   |Stderr|
|-----|------:|----------------|-----:|-----------|---|-----:|---|-----:|
|gsm8k|      3|flexible-extract|     5|exact_match|↑  |0.9583|±  |0.0055|
|     |       |strict-match    |     5|exact_match|↑  |0.9591|±  |0.0055|

---

Essential Elements of an Effective PR Description Checklist

• The purpose of the PR, such as "Fix some issue (link existing issues this PR will resolve)".
• The test plan, such as providing test command.
• The test results, such as pasting the results comparison before and after, or e2e results
• (Optional) The necessary documentation update, such as updating supported_models.md and examples for a new model.

[claude]

Claude Code Review

This pull request is from a fork — automated review is disabled. A repository maintainer can comment @claude review to run a one-time review.

[mergify]

Documentation preview: https://vllm--40960.org.readthedocs.build/en/40960/

[gemini-code-assist]

Code Review

This pull request introduces support for bf16 dispatch and deferred input quantization within the flashinfer_nvlink_one_sided MoE kernel, specifically to accommodate expert types like trtllm_mxfp4. The implementation updates the MoeAlltoAll workspace to support dynamic payload sizes and modifies the preparation logic to handle cases where quantization is performed post-dispatch. Feedback was provided regarding a potential AttributeError when quant_config is None and the necessity of updating existing validation checks to ensure the new logic is reachable.

[gemini-code-assist]

This logic is currently unreachable for non-nvfp4 models (such as mxfp4) because of the ValueError check at lines 232-239 (outside this diff hunk). To support BF16 and deferred quantization for other formats, that validation block should be removed or updated. Additionally, the code should safely handle cases where quant_config is None (e.g., for unquantized models) to avoid an AttributeError when accessing quant_config.quant_dtype.

        quant_dtype = quant_config.quant_dtype if quant_config is not None else None
        if defer_input_quant or quant_dtype is None:
            # Experts (e.g. trtllm_mxfp4 with mxfp8 activations) quantize
            # post-dispatch; ship bf16 tokens with no per-token scale payload.
            dispatch_dtype_bytes_per_elem, dispatch_has_fp8_scale = 2, False
        elif quant_dtype == "nvfp4":
            dispatch_dtype_bytes_per_elem, dispatch_has_fp8_scale = 0, True
        else:
            raise NotImplementedError(
                "flashinfer_nvlink_one_sided dispatch only supports nvfp4, "
                "bf16, and defer_input_quant paths today; got "
                f"quant_dtype={quant_dtype!r}"
            )

[liuzijing2014]

I think we need to change this condition as well otherwise, the branch would early exit at here for bf16.

[hjjq]

LGTM, thanks!!

[bobboli]

LGTM overall. I have confirmed that the gen step time is ~40ms for AG/RS and ~35ms for NVLinkOneSided A2A under the following config:

• Model: DeepSeek-V4-Flash
• ISL / OSL: 8192 / 1024
• Local batch size: 512 per DP rank
• DP8EP8, 8x B200
• kv_cache_dtype: fp8

[bobboli]

A small missing piece is payload_in_workspace for combine. It means that the MoE module could directly output to the workspace of A2A, so that there is no need for extra copying of the tokens from the local Tensor onto workspace.
Related code in TRTLLM is here.
It requires that the MoE OP could be provided an pre-allocated tensor as output, which is intrusive.

[zyongye]

A small missing piece is payload_in_workspace for combine. It means that the MoE module could directly output to the workspace of A2A, so that there is no need for extra copying of the tokens from the local Tensor onto workspace. Related code in TRTLLM is here. It requires that the MoE OP could be provided an pre-allocated tensor as output, which is intrusive.

Thanks. We do have pre-allocated workspace buffer but now we are just manually copy it. I will take a look at trtllm code!

[bobboli]

Thanks. We do have pre-allocated workspace buffer but now we are just manually copy it. I will take a look at trtllm code!

Yes, the point is that a portion of workspace is viewed as a Tensor, and passed into the MoE OP as output Tensor, so that the MoE OP directly outputs onto the workspace.