int16 Block-Scaled State and Stochastic Rounding for SSU (mamba)

.gitignore

@@ -24,6 +24,10 @@ flashinfer/cute_dsl/benchmark_gated_delta_rule.py
 # vscode
 .vscode/
 
+# zed text editor
+.zed/
+.rules
+
 # Byte-compiled / optimized / DLL files
 __pycache__/
 *.py[cod]

csrc/flashinfer_mamba_binding.cu

@@ -38,10 +38,13 @@ void selective_state_update(
     bool dt_softplus,
     Optional<TensorView> state_batch_indices,  // (batch,)
     int64_t pad_slot_id,
-    TensorView output,  // same as x
+    Optional<TensorView> state_scale,  // float32: (state_cache_size, nheads, dim)
+    TensorView output,                 // same as x
     bool disable_state_update,
     Optional<TensorView> intermediate_states_buffer,  // (batch, cache_steps, nheads, dim, dstate)
     Optional<TensorView> intermediate_state_indices,  // (batch,)
+    Optional<TensorView> intermediate_state_scales,   // float32: (batch, cache_steps, nheads, dim)
+    int64_t rand_seed,                                // used for Philox rounding
     int64_t cache_steps,
     int64_t algorithm);  // SSUAlgorithm: 0=auto, 1=simple, 2=vertical, 3=horizontal