[KernelGen][MThreads] Add bucketize Moore Threads specialized operator

[Yukun-Cui]

[KernelGen][MThreads] Add bucketize Moore Threads specialized operator

Summary

Add a Moore Threads (MUSA) specialized Triton kernel for bucketize. The kernel keeps the generic vectorized binary-search algorithm (one [lo, hi) window per lane, narrowed over ceil(log2(n_boundaries + 1)) iterations) and adds a MUSA device gate with default_bucketize fallback, Moore-Threads-tuned autotune configs keyed on (n_elements, n_boundaries), and a contiguous flat output allocation so stores land in the returned tensor regardless of the input's memory layout.

Testing

• Reused the existing upstream accuracy tests tests/test_bucketize.py (-m bucketize): 64 passed, 2 skipped
• Validated against the reference on the MUSA device; specialization confirmed active via the GEMS_MTHREADS BUCKETIZE debug log
• Falls back to the generic implementation for unsupported dtype/device/shape (fp64/int64 not supported on Moore Threads hardware; empty boundaries; boundary/input dtype mismatch)
• The 2 skipped cases are the integer-boundary parametrization (test_bucketize_boundary_cases[integer-*]): MUSA's native torch.bucketize raises RuntimeError: Bucketize func unsupported! for integer boundaries, so the test reference itself cannot run on MUSA. This is a hardware/reference limitation, not a kernel bug — gated with pytest.mark.skipif(vendor == "mthreads") (condition-gate only; other backends still exercise the path)

Performance

Compared against the generic FlagGems implementation on Moore Threads (MUSA).

bucketize

dtype	Size	Torch Latency (ms)	Gems Latency (ms)	Speedup
float32	1073741824	42.903439	11.754280	3.650x
float32	64, 64	0.005360	0.007680	0.698x
float32	4096, 4096	0.677480	0.202720	3.342x
float32	64, 512, 512	0.677760	0.202360	3.349x
float32	1024, 1024, 1024	42.893980	11.756540	3.649x
float32	268435456	10.714720	2.952880	3.629x
float32	10000, 1	0.005200	0.007120	0.730x
float32	10000, 256	0.109040	0.037560	2.903x
float32	10000, 65536	26.168280	7.179160	3.645x
float32	100, 1, 100	0.005200	0.007120	0.730x
float32	100, 256, 100	0.109040	0.037480	2.909x
float32	100, 65536, 100	26.175961	7.178960	3.646x

Operator	Arithmetic Mean Speedup
bucketize	2.74x

Only float32 is benchmarked on MUSA: the benchmark baseline is torch.bucketize, and MUSA's native torch.bucketize accepts only float32/int32/int64 — it raises RuntimeError: Bucketize supports dtypes of float32, int32 and int64, but now it is Half/BFloat16 for fp16/bf16, so there is no valid baseline to compare against for those dtypes. The specialized kernel itself does support fp16/bf16 (computing the comparison in the native element type), and the accuracy tests cover fp16/fp32/bf16 because the test reference upcasts the input, sidestepping the native limitation.

Files Changed

• src/flag_gems/runtime/backend/_mthreads/ops/bucketize.py: Moore Threads Triton kernel + fallback
• src/flag_gems/runtime/backend/_mthreads/ops/__init__.py: Register import and __all__
• tests/test_bucketize.py: gate the integer-boundary case for mthreads (native reference unsupported)
• benchmark/test_bucketize.py: restrict dtypes to fp32 on mthreads (native torch.bucketize baseline supports fp32 only)

[Dongxu-H]

/test | bucketize:mthreads

[github-actions]

⚠️ Cannot proceed with operation

• reviewDecision: skip_reviews
• commitStatus: FAILURE

Reviews are not required for this operation but CI checks must be passing in order to continue

[Dongxu-H]

/test | bucketize:mthreads

[github-actions]

⚠️ Cannot proceed with operation

• reviewDecision: skip_reviews
• commitStatus: FAILURE

Reviews are not required for this operation but CI checks must be passing in order to continue

[Dongxu-H]

/test | bucketize:mthreads

[github-actions]

⚠️ Cannot proceed with operation

• reviewDecision: skip_reviews
• commitStatus: FAILURE

Reviews are not required for this operation but CI checks must be passing in order to continue

[Dongxu-H]

/test | bucketize:mthreads

[github-actions]

⚠️ Cannot proceed with operation

• reviewDecision: skip_reviews
• commitStatus: FAILURE

Reviews are not required for this operation but CI checks must be passing in order to continue

[Dongxu-H]

/test | bucketize:mthreads

[github-actions]

⚠️ Cannot proceed with operation

• reviewDecision: skip_reviews
• commitStatus: FAILURE

Reviews are not required for this operation but CI checks must be passing in order to continue

[Dongxu-H]

CI failures in python-op, backend-tests (mthreads-musa520), and code-style checks need to be fixed before merge. Logger format is correct for Mthreads backend. Dtype limitation to fp16/bf16/fp32 is appropriate for Mthreads hardware. Test skipif for integer boundaries is acceptable as hardware limitation.

[Dongxu-H]

LGTM

[Dongxu-H]

/test | bucketize:mthreads

[github-actions]

⚠️ Cannot proceed with operation

• reviewDecision: skip_reviews
• commitStatus: FAILURE

Reviews are not required for this operation but CI checks must be passing in order to continue

[Dongxu-H]

/test | bucketize:mthreads

[github-actions]

⚠️ Cannot proceed with operation

• reviewDecision: skip_reviews
• commitStatus: FAILURE

Reviews are not required for this operation but CI checks must be passing in order to continue

[Dongxu-H]

CI code-style failure needs to be fixed. Run pre-commit or ruff to check and fix formatting issues.