[Rocm][fused_moe][fp4] view weight to torch.float4_e2m1fn_x2 when running aiter fused moe for fp4 model #27474
Conversation
zejunchen-zejun requested review from mgoin, pavanimajety, robertgshaw2-redhat, tlrmchlsmth and yewentao256 as code owners 
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
This pull request correctly addresses a dtype mismatch issue for the aiter fused MoE kernel on ROCm when using FP4 models. The change involves viewing the uint8 weight tensor as torch.float4_e2m1fn_x2 before passing it to the kernel. My review identifies a performance improvement opportunity by moving the hasattr check out of the apply method's hot path and into the class initializer. This would prevent redundant checks on every forward pass.
| if hasattr(torch, "float4_e2m1fn_x2"): | ||
| w13_weight = layer.w13_weight.view(torch.float4_e2m1fn_x2) | ||
| w2_weight = layer.w2_weight.view(torch.float4_e2m1fn_x2) | ||
| else: | ||
| w13_weight = layer.w13_weight | ||
| w2_weight = layer.w2_weight |
There was a problem hiding this comment.
The hasattr(torch, "float4_e2m1fn_x2") check is performed on every forward pass within the apply method. This is inefficient as the result of this check is constant for a given PyTorch environment. This check should be performed only once during initialization to avoid repeated overhead in the critical path of inference.
I recommend moving this check to the __init__ method of the QuarkOCP_MX_MoEMethod class and caching the result in an instance variable. For example:
In __init__:
self.fp4_dtype = getattr(torch, "float4_e2m1fn_x2", None)Then, here in apply, you can simplify the logic:
w13_weight = layer.w13_weight w2_weight = layer.w2_weight if self.fp4_dtype: w13_weight = w13_weight.view(self.fp4_dtype) w2_weight = w2_weight.view(self.fp4_dtype)This would be more performant and also makes the code in apply cleaner by removing the if/else block. Since modifying __init__ is outside the current diff, I am not providing a direct code suggestion, but this change is highly recommended for performance.
There was a problem hiding this comment.
Per ai suggestion, this viewing can be moved to process_weights_after_loading to avoid being invoked every forward.
| Hi, @maleksan85 @HaiShaw Could you help review the code changes here? Thank you. |
BowenBao left a comment
BowenBao left a comment There was a problem hiding this comment.
Looks good with suggestions.
| if hasattr(torch, "float4_e2m1fn_x2"): | ||
| w13_weight = layer.w13_weight.view(torch.float4_e2m1fn_x2) | ||
| w2_weight = layer.w2_weight.view(torch.float4_e2m1fn_x2) | ||
| else: | ||
| w13_weight = layer.w13_weight | ||
| w2_weight = layer.w2_weight |
There was a problem hiding this comment.
Per ai suggestion, this viewing can be moved to process_weights_after_loading to avoid being invoked every forward.
| Hi, @BowenBao |
| Hi, @HaiShaw Could you help review this PR? It fixed the FP4 fused MOE functionality issue. Thank you! |
| just hold on merging.. I believe @zejunchen-zejun has yet pushed the changes. |
zejunchen-zejun force-pushed the zejun/fix_fp4_fused_moe_func_issue_for_rocm branch 3 times, most recently from 1c21289 to 5ed92fd Compare 
when running aiter fused moe for fp4 model Signed-off-by: zejunchen-zejun <zejun.chen@amd.com>
zejunchen-zejun force-pushed the zejun/fix_fp4_fused_moe_func_issue_for_rocm branch from 5ed92fd to 15df282 Compare | Thanks @zejunchen-zejun , LGTM |
| Hi, @HaiShaw @SageMoore @BowenBao @gshtras |
| Hi, @HaiShaw @SageMoore @BowenBao @gshtras @LucasWilkinson |
4 participants
doing view for moe weight to torch.float4_e2m1fn_x2 for aiter FP4 fused moe kernel
with this PR, aiter cannot find the suitable kernel for the weight, whose type is uint8
with this PR, the deepseek fp4 can run successfully