Update the megablox kernel utility function to support float16

[wang2yn84]

Description

The current function only switches between bf16 and f32. It should also take into consideration of f16.

Tests

No functional change. Reply on presubmit tests.

Checklist

Before submitting this PR, please make sure (put X in square brackets):

• [x] I have performed a self-review of my code.
• [x] I have necessary comments in my code, particularly in hard-to-understand areas.
• [x] I have run end-to-end tests tests and provided workload links above if applicable.
• [x] I have made or will make corresponding changes to the doc if needed.

[wang2yn84]

LGTM. Re: No functional change, rely on presubmit test.

Is there a presubmit test that will test what you need here? Or do we need to try running with these dtypes separately?

That's a good question. Which script shall I run to validate the change?

[bvandermoon]

LGTM. Re: No functional change, rely on presubmit test. Is there a presubmit test that will test what you need here? Or do we need to try running with these dtypes separately?

That's a good question. Which script shall I run to validate the change?

@yangyuwei do you have a GPU-specific test that @wang2yn84 could use here?