Fix FMHA index order and check (#231)

Author: AntonOresten

examples/fmha.jl

@@ -81,7 +81,7 @@ function fmha_kernel(Q::ct.TileArray{T, 4}, K::ct.TileArray{T, 4},
         # QK product
         # K is (D_k, SeqLen_KV, KVH, Batch)
         # Load (TILE_N, TILE_D, 1, 1) with order=(2,1,3,4) to transpose D and N
-        k = reshape(ct.load(K; index=(Int32(1), j + Int32(1), off_kv_h, batch_idx),
+        k = reshape(ct.load(K; index=(j + Int32(1), Int32(1), off_kv_h, batch_idx),
                             shape=(TILE_N, TILE_D, 1, 1), order=(2, 1, 3, 4),
                             latency=2),
                     (TILE_N, TILE_D))  # (TILE_N, TILE_D)
@@ -279,7 +279,7 @@ function verify(data, result)
     expected = ref_fmha(data.Q, data.K, data.V; causal=data.causal)
     actual = Float32.(Array(result.out))
     max_diff = maximum(abs.(actual .- expected))
-    @assert max_diff < 1e-2 "FMHA incorrect! max diff: $max_diff"
+    @assert isapprox(actual, expected, rtol=1e-2, atol=1e-3) "FMHA incorrect! max diff: $max_diff"
 end
 
 

examples/fmha.py

@@ -263,7 +263,7 @@ def verify(data, result):
     expected = ref_fmha(data["Q"], data["K"], data["V"],
                         causal=data["causal"])
     actual = cp.asnumpy(result["out"]).astype(np.float32)
-    assert np.allclose(actual, expected, rtol=1e-2, atol=1e-2), \
+    assert np.allclose(actual, expected, rtol=1e-2, atol=1e-3), \
         f"FMHA incorrect! max diff: {np.max(np.abs(actual - expected))}"
 
 

src/language/operations.jl

@@ -345,6 +345,7 @@ Index is 1-indexed. Shape must be compile-time constant.
 - `order`: Optional tuple specifying the logical-to-physical dimension mapping (1-indexed).
   For example, `order=(2, 1)` indicates dimension 2 is contiguous in memory,
   enabling coalesced loads from transposed/permuted arrays.
+  `index[i]` and `shape[i]` describe tile dim `i`, which maps to source dim `order[i]`.
   Default: `nothing` → identity `(1, 2, ..., N)`.
 
 # Padding Modes
@@ -366,10 +367,10 @@ outside the array, the behavior is undefined regardless of `padding_mode`.
 
 # Example
 ```julia
-tile = ct.load(arr, (bid,), (TILE_N[],); padding_mode=ct.PaddingMode.Zero, latency=3)
+tile = ct.load(arr, (bid,), (TILE_N,); padding_mode=ct.PaddingMode.Zero, latency=3)
 
 # Load from a transposed array with coalesced access
-tile = ct.load(arr, (bidx, bidy), (TM, TN); order=(2, 1))
+tile = ct.load(arr, (bidy, bidx), (TN, TM); order=(2, 1))
 ```
 """
 @inline function load(arr::TileArray, index, shape::NTuple{<:Any, Int};
@@ -442,6 +443,7 @@ behavior is undefined.
 # Dimension Ordering
 - `order`: Optional tuple specifying the logical-to-physical dimension mapping (1-indexed).
   Must match the `order` used in the corresponding `load` for permuted arrays.
+  `index[i]` and `shape[i]` describe tile dim `i`, which maps to destination dim `order[i]`.
   Default: `nothing` → identity `(1, 2, ..., N)`.
 
 # Optimization Hints