result mismatch for mosaic gpu attention on sm100

[qelk123]

Hi,
I tried to run the newly added sm100 mosaic gpu attention kernel on B200 with args:

PallasMosaicGpuFlashAttention(config=None, vjp=PallasMosaicGpuFlashAttentionVjp(config=None, vjp=None, use_base2=True, dbias_intermediate_dtype=None), use_base2=True, use_stable_softmax=<class 'tokamax._src.ops.attention.base.AUTO'>) on NVIDIA B200 with key immutabledict({'q': BatchedShapeDtype(shape=(4, 1024, 32, 128), dtype=bfloat16, vmap_axes=()), 'k': BatchedShapeDtype(shape=(4, 1024, 32, 128), dtype=bfloat16, vmap_axes=()), 'v': BatchedShapeDtype(shape=(4, 1024, 32, 128), dtype=bfloat16, vmap_axes=()), 'bias': None, 'dropout_mask': None, 'dropout_rate': 0.0, 'k_indices': None, 'logits_dtype': dtype('float32'), 'logits_scale': 0.08838834764831843, 'logits_soft_cap': None, 'mask': Mask(bool_mask=None, q_start=None, q_end=None, k_start=None, k_end=None, is_causal=False), 'normalize_output': True, 'paging_info': None, 'precision': (DotAlgorithmPreset.BF16_BF16_F32, DotAlgorithmPreset.BF16_BF16_F32), 'q_indices': None, 'return_residuals': False})

But I got result mismatch with config(which is actually in the current tuning space):

Config(block_q=128, block_kv=128, num_stages=2, fold_q_sequence_heads=False, split_k=1, block_d=128, num_tma_splits=2, collective=False)

I think the potential issue is the code here, when the tma_chunk_size is smaller than block_d, this loop won't be executed at all, and alpha won't be applied to the acc, further lead to a result mismatch error.

[sbodenstein]

Thank you for reporting, we are taking a look.