[Bugfix][Mamba] - Fix Conv State Kernel FP32 Support

[Josephasafg]

Purpose

There was a bug in the FP32 support recently added to mamba conv state where if you pass this arg --mamba-cache-dtype "float32" you'd get this error:

E       triton.compiler.errors.CompilationError: at 102:8:
E           w_base = w_ptr + (idx_feats * stride_w_dim)  # [BLOCK_N,]
E
E           # Does 2 things:
E           # 1. READ prior-block init-state data - [done by every Triton programs]
E           # 2. update conv_state with new data [only by the Triton program handles chunk_offset=0]
E           if chunk_offset == 0:
E               # read from conv_states
E               load_init_state = False
E               if HAS_INITIAL_STATES:  # the new HAS_INITIAL_STATES
E                   load_init_state = tl.load(has_initial_states_ptr + idx_seq).to(
E                       tl.int1)
E               if load_init_state:
E               ^
E       AssertionError("Mismatched type for col0 between then block (<['256'], fp32>) and else block (<['256'], bf16>)")

The kernel needed to cast the input type to be the same as conv state.

Test Plan

Added two tests to test_hybrid.py as they were missing. Tests pass

Test Result

---

Essential Elements of an Effective PR Description Checklist

• The purpose of the PR, such as "Fix some issue (link existing issues this PR will resolve)".
• The test plan, such as providing test command.
• The test results, such as pasting the results comparison before and after, or e2e results
• (Optional) The necessary documentation update, such as updating supported_models.md and examples for a new model.
• (Optional) Release notes update. If your change is user facing, please update the release notes draft in the Google Doc.

[gemini-code-assist]

Code Review

This pull request correctly addresses a Triton compilation error for Mamba's FP32 convolution state during prefill by casting the input tensor to match the state's data type. The fix in causal_conv1d_fn is appropriate. However, the same bug likely exists in the causal_conv1d_update function, which handles the decode path, and this has not been addressed. This could lead to failures during token generation after the prefill phase. I've added a critical comment highlighting this omission. The test enhancements are good, improving coverage for different cache dtype parameters.

[mergify]

This pull request has merge conflicts that must be resolved before it can be
merged. Please rebase the PR, @Josephasafg.

https://docs.github.com/en/pull-requests/collaborating-with-pull-requests/working-with-forks/syncing-a-fork

[Josephasafg]

@heheda12345 @tdoublep Can I please get a review?

[tdoublep]

Do we definitely want to cast x to the conv_state dtype, rather than casting conv_state to the x_dtype?

[Josephasafg]

Its a good question - Since the user picks fp32 for the cache type, Im afraid that by downcasting it to fp16 and then back to fp32, we could lose accuracy doing it. I had also figured that by choosing fp32, we want the computations to be done in that type dont we?

[tdoublep]

How does it work for the SSM state? I guess we want it to be consistent.

[Josephasafg]

we cast to float in the kernel

[tdoublep]

LGTM