moe layer with nvfp4 grouped_mm

[jjsjann123]

Stacked PRs

#5230 moe layer with nvfp4 grouped_mm <-- this PR
#5345 exposing layout op at direct python binding
#5198 refactor number of groups in layout op
#5174 allow layout op in automatic scheduler

This PR

Fixes that I want to be reviewed:

• GatherScatter IterType shouldn't propagate in fusoin.
• packed fp4 dtype is not properly handled in direct binding.

Changes on test_moe.py is just a display-of-concept:

• Enable nvfuser nvfp4 grouped_mm in MoE benchmark.

Note: This has the indexing issue in #5200. This branch also requires some thunder inflight PRs.

[github-actions]

Review updated until commit 13db45d

Description

• Added NVFP4 quantized grouped matrix multiplication support

• Fixed GatherScatter IterType propagation in fusion

• Enhanced MoE layer with FP4 weight handling

• Updated group size calculation with total tokens

---

Changes walkthrough 📝

Relevant files

Enhancement

test_moe.py Added NVFP4 grouped MM with quantization support                  tests/python/test_moe.py Added _group_sizes_from_offsets with total_tokens parameter for accurate group sizing Introduced nvfuser_f16a_nvfp4weight_scaled_grouped_mm custom op with fake implementation Implemented gmm_nvfuser translator for NVFuser with FP4 quantization logic Enhanced GroupedLinear to support FP4 weights, scaling factors, and problem sizes Updated GroupedSwiGLU and Llama4MoE to handle new FP4 parameters +205/-9

---

PR Reviewer Guide 🔍

Here are some key observations to aid the review process:

🧪 PR contains tests

⚡ Recommended focus areas for review Possible Issue The function _group_sizes_from_offsets now takes an additional total_tokens parameter, but its logic assumes offsets[1:] is valid. If offsets is empty or has only one element, this may lead to incorrect group sizes or runtime errors. def _group_sizes_from_offsets(offsets: torch.Tensor, total_tokens: int) -> list[int]: group_sizes = [] prev = 0 for offset in offsets[1:]: group_sizes.append(offset - prev) prev = offset group_sizes.append(total_tokens - prev) return group_sizes Performance Concern The custom op nvfuser_f16a_nvfp4weight_scaled_grouped_mm currently returns the result of grouped_mm using dropme instead of utilizing the quantized fp4_weight and scaling logic, which may defeat the purpose of the nvfp4 optimization. @torch.library.custom_op( "nvf_cutlass::f16a_nvfp4weight_scaled_grouped_mm", mutates_args=() ) def nvfuser_f16a_nvfp4weight_scaled_grouped_mm( activation: torch.Tensor, fp4_weight: torch.Tensor, weight_scaling_factor: torch.Tensor, global_scale: torch.Tensor, offsets: torch.Tensor, blockscale_offsets: torch.Tensor, problem_sizes: torch.Tensor, dropme: torch.Tensor, ) -> torch.Tensor: return grouped_mm(activation, dropme, offsets) Correctness Risk In gmm_nvfuser, the problem_sizes tensor is constructed using tokens_per_expert.unsqueeze(-1) and concatenated with self.n and self.k, but it's unclear if the shape and order match the expected input for cutlass_nvfp4_grouped_mm, risking incorrect kernel behavior. m_size = activation.shape[0] k_size = activation.shape[1] k_tile_size = k_size // 16

[crcrpar]

nit-picking, this function just calls one function twice --

    from thunder.executors.nvfuserex_impl import register_supported
    from thunder.executors.torchex import _always_executable

    register_supported(symbol, translator_for_nvfuser, checker or _always_executable)
    register_supported(symbol.id, translator_for_nvfuser, checker or _always_executable)

I'm taking a bit of time cleaning up the Lightning-AI/lightning-thunder#2481 tests. To reduce the number of cherry-picks, just calling register_supported here might make things easier

[jjsjann123]

!test

[jjsjann123]

!test

[jjsjann123]

I'm undecided how I would want to proceed with tests. But the other smaller pieces should be good for review.

[jjsjann123]

I'll change target and refactor this to only handle nvfp4 via direct binding. since @Priya2698 is cherry-picking the propagation in #5365

[jjsjann123]

All code changes have been merged in separate PRs, it's only the test_moe.py that's being updated in this PR. I'll clean it up and request a sanity check afterwards.