Fall back to the for-loop grouped_mm on CPU (#46743)
* Fall back to the for-loop grouped_mm on CPU
torch._grouped_mm has no CPU kernel on supported torch versions (e.g. 2.8 raises
NotImplementedError), but _can_use_grouped_mm returned True for CPU tensors as
long as they were 16-byte aligned, so the grouped-mm experts path called
torch._grouped_mm on CPU and every MoE model using it crashed. Route CPU through
the for-loop fallback for torch <= 2.10, which the comment already described as
the intended CPU path.
Signed-off-by: Ting Sun <suntcrick@gmail.com>
* Only fall back on CPU below torch 2.9; keep alignment gate for 2.9/2.10
torch 2.8 has no CPU kernel for _grouped_mm (raises NotImplementedError), while
2.9/2.10 have one that just needs 16 byte alignment. Restore the original
misalignment gate and only force the for-loop unconditionally below 2.9, so the
CPU kernel is still used where it works. Drop the dedicated test: the crash only
happens on torch <= 2.8 CPU, which CI (torch 2.11) never exercises.
Signed-off-by: Ting Sun <suntcrick@gmail.com>
* List the grouped_mm CPU fall-back cases one per line
Per review: flatten the condition so each fall-back case is its own or-clause
(dynamo non-bf16; CPU <= 2.10 misaligned; CPU <= 2.8 with no kernel at all), and
list them one by one in the comment.
Signed-off-by: Ting Sun <suntcrick@gmail.com>
---------
Signed-off-by: Ting Sun <suntcrick@gmail.com>
