[BUG] Blackwell fmha_bwd overflows int32 spec

Which component has the problem?

CuTe DSL

Bug Report

Describe the bug

Hello!

I think the workspace tensor created here goes above 2**31 (int32 range) which causes some odd behaviours: https://github.com/NVIDIA/cutlass/blob/8cd5bef43a2b0d3f9846b026c271593c6e4a8e8a/examples/python/CuTeDSL/blackwell/fmha_bwd.py#L3004-L3008

Steps/Code to reproduce bug

Reproduction command in examples/python/CuTeDSL/blackwell/:

python -m fmha_bwd     --acc_dtype Float32     --element_dtype Float16     --mma_tiler_mn 128,128     --s_q '1185, 417, 607, 1559, 107, 221, 674, 117, 188, 152, 266, 631, 284, 1365, 419, 342, 715, 311, 22, 214, 317, 1214, 961, 3422, 674'     --s_k '1185, 417, 607, 1559, 107, 221, 674, 117, 188, 152, 266, 631, 284, 1365, 419, 342, 715, 311, 22, 214, 317, 1214, 961, 3422, 674'     --d 128     --h_q 128     --h_k 128     --b 25

Error:

Running Blackwell SM100 FMHA bwd test with:
  s_q: (1185, 417, 607, 1559, 107, 221, 674, 117, 188, 152, 266, 631, 284, 1365, 419, 342, 715, 311, 22, 214, 317, 1214, 961, 3422, 674)
  s_k: (1185, 417, 607, 1559, 107, 221, 674, 117, 188, 152, 266, 631, 284, 1365, 419, 342, 715, 311, 22, 214, 317, 1214, 961, 3422, 674)
  h_q: 128
  h_k: 128
  d: 128
  b: 25
  is_causal: False
  bottom_right_align: False
  element_dtype: Float16
  acc_dtype: Float32
  mma_tiler_mn: (128, 128)
  scale_softmax: 0.0
  window_size: (-1, -1)
  warmup_iterations: 0
  iterations: 1
  skip_ref_check: False
Compiling kernel with cute.compile ...
Traceback (most recent call last):
  File "<frozen runpy>", line 198, in _run_module_as_main
  File "<frozen runpy>", line 88, in _run_code
  File "/mnt/vast/home/thomas/code/mistral/mistral/mistral/model/attention/cutlass/fmha_bwd.py", line 3547, in <module>
    run(
  File "/mnt/vast/home/thomas/code/mistral/mistral/mistral/model/attention/cutlass/fmha_bwd.py", line 3011, in run
    compiled_fmha_bwd = cute.compile(
                        ^^^^^^^^^^^^^
  File "/usr/local/lib/python3.12/dist-packages/nvidia_cutlass_dsl/python_packages/cutlass/cute/runtime.py", line 376, in __c_pointers__
    self._memref_desc = self._dltensor_wrapper.build_memref_desc(
                        ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
OverflowError: Value overflow: 5697536000 exceeds range of l

Expected behavior

I'm hoping that:

• maybe workspace size if overshot by a lot?
• migrate from int32 -> int64 would perhaps solve this?

Environment details (please complete the following information):

Additional context

Start in examples/python/CuTeDSL/blackwell/fmha_bwd.py around lines 3004-3008 and run the provided fmha_bwd reproduction command. Trace the workspace tensor size and integer type through cute.compile, then verify the reproduction compiles without the reported overflow while preserving correct workspace sizing.