Fix fallback to default tactic (flashinfer autotuner) with trtllm_fp4_block_scale_moe

[danisereb]

Purpose

The flashinfer autotuner expects the first dimension of the MoE tensors to be num_tokens.

The relevant code can be found in flashinfer function get_trtllm_moe_sm100_module:

        # their first dimension is num_tokens which will be tuned
        tuning_config_with_hidden_states_scales = TuningConfig(
            dynamic_tensor_specs=(
                DynamicTensorSpec(
                    (0, 1, 2, 3, 4, 5),
                    (0, 0, 0, 0, 0, 0),
                    get_last_power_of_2_num_tokens_buckets(8192, 1),
                    lambda x: min(last_positive_power_of_2(x), 8192),
                    dynamic_tensor_initializers,
                ),
            )
        )

When running with vLLM main with env var FLASHINFER_LOGGING_LEVEL=DEBUG, we can see the following:

flashinfer.jit: [AutoTunner]: Using fallback tactic for flashinfer::trtllm_fp4_block_scale_moe with input shapes ...

Using reshape(...) instead of flatten() solves the issues (no using fallback tactic prints).

This change is aligned with function Mxfp4MoEMethod.apply_monolithic:

x_scale = x_scale.view(torch.float8_e4m3fn).reshape(*x.shape[:-1], -1)

Test Plan

Verify that there is no drop in accuracy (lm_eval).

Test Result

Used the following MoE NVFP4 model:
https://huggingface.co/RedHatAI/Llama-4-Scout-17B-16E-Instruct-NVFP4

Use env vars:

export VLLM_USE_FLASHINFER_MOE_FP4=1
export VLLM_FLASHINFER_MOE_BACKEND=latency

To use 'FLASHINFER_TRTLLM' NvFp4 MoE backend.

lm_eval gsm8k:

lm_eval \
  --model vllm \
  --model_args pretrained="/my_home/hf_models/RedHatAI/Llama-4-Scout-17B-16E-Instruct-NVFP4",dtype=auto,add_bos_token=True,max_model_len=4096,tensor_parallel_size=2,enable_chunked_prefill=True,enforce_eager=True \
  --tasks gsm8k_llama \
  --apply_chat_template \
  --fewshot_as_multiturn \
  --batch_size auto

### vLLM main commit: 
|   Tasks   |Version|     Filter     |n-shot|  Metric   |   |Value |   |Stderr|
|-----------|------:|----------------|-----:|-----------|---|-----:|---|-----:|
|gsm8k_llama|      3|flexible_extract|     8|exact_match|↑  |0.9371|±  |0.0067|
|           |       |strict_match    |     8|exact_match|↑  |0.9371|±  |0.0067|


### This PR:
|   Tasks   |Version|     Filter     |n-shot|  Metric   |   |Value |   |Stderr|
|-----------|------:|----------------|-----:|-----------|---|-----:|---|-----:|
|gsm8k_llama|      3|flexible_extract|     8|exact_match|↑  |0.9371|±  |0.0067|
|           |       |strict_match    |     8|exact_match|↑  |0.9371|±  |0.0067|

There was no change in vllm bench serve performance (tok/sec).

---

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.
• (Optional) Release notes update. If your change is user facing, please update the release notes draft in the Google Doc.

[gemini-code-assist]

Code Review

This pull request correctly addresses a performance issue with the flashinfer MoE kernels. The original code used .flatten() on the hidden_states_scale tensor, which collapsed its dimensions and prevented the flashinfer autotuner from identifying the optimal tactic, causing a fallback to a less efficient default. The fix replaces .flatten() with .reshape(*hidden_states_fp4.shape[:-1], -1), which correctly preserves the tensor's leading dimensions (including num_tokens) as expected by the kernel. This change is applied consistently in both flashinfer_trtllm_fp4_moe and flashinfer_trtllm_fp4_routed_moe functions, resolving the performance degradation. The change is correct and well-justified.

[mgoin]

Makes sense, LGTM

[mgoin]

Is reshape required or could we use view here? I just ask so we avoid the implicit .contiguous() that reshape has

[danisereb]

Like I mentioned in the PR description, I followed the reshape from Mxfp4MoEMethod.apply_monolithic.

As far as I understand, reshape() uses a view when possible and implicit contiguous if it's required.
And view() may throw an error if the underlying storage/strides does not match the requested shape.

[danisereb]

Rebased to get this change:
#34874

seems to be related to the CI failures I see.

There's still a bug in certain cases (num tokens not power of 2), but the fix is in flashinfer - see this PR:
flashinfer-ai/flashinfer#2617

[pavanimajety]

Seems like the test failures are unrelated (model architectures in the failing tests likely won't invoke trtllm FP4 moe kernels)

There's still a bug in certain cases (num tokens not power of 2), but the fix is in flashinfer - see this PR:
flashinfer-ai/flashinfer#2617

With smaller batch sizes that would be enabled through #34936 in --performance-mode latency, we may see an issue. However, since this is a perf improvement, we should be okay to merge.