Flydsl update fix dynmaic layout; use fx.ptr for kernels not using layout

.github/requirements/triton-test.txt

@@ -8,7 +8,7 @@ pybind11==3.0.1
 ninja==1.11.1.4
 psutil
 packaging
-flydsl==0.1.8
+flydsl==0.1.9.dev599
 
 # Test deps.
 pandas==2.2.3

.github/workflows/atom-test.yaml

@@ -184,9 +184,9 @@ jobs:
           RUN pip install --upgrade "pybind11>=3.0.1"
           RUN pip show pybind11
           RUN rm -rf /app/aiter-test
-          RUN git clone ${{ env.GITHUB_REPO_URL }} /app/aiter-test && \\
+          RUN git clone --no-checkout ${{ env.GITHUB_REPO_URL }} /app/aiter-test && \\
               cd /app/aiter-test && \\
-              git checkout ${{ env.GITHUB_COMMIT_SHA }} && \\
+              git checkout --force ${{ env.GITHUB_COMMIT_SHA }} && \\
               git submodule sync && git submodule update --init --recursive && \\
               MAX_JOBS=64 PREBUILD_KERNELS=0 GPU_ARCHS=gfx950 pip install -e . && \\
               ./.github/scripts/install_triton.sh

aiter/aot/flydsl/gemm.py

@@ -197,6 +197,12 @@ def _compile_executable_to_cache(exe, *args) -> None:
         exe(*args)
 
 
+def _ptr_view_safe(t):
+    from aiter.ops.flydsl.gemm_kernels import _ptr_view_safe as _wrap
+
+    return _wrap(t)
+
+
 def _compile_hgemm_to_cache(
     *,
     m: int,
@@ -271,7 +277,15 @@ def _compile_hgemm_to_cache(
     # optional bias and split-K sync tensors.
     launch_bias = bias if has_bias else b
     _compile_executable_to_cache(
-        exe, out, a, b, launch_bias, m, semaphore, signal, stream
+        exe,
+        _ptr_view_safe(out),
+        _ptr_view_safe(a),
+        _ptr_view_safe(b),
+        _ptr_view_safe(launch_bias),
+        m,
+        _ptr_view_safe(semaphore),
+        _ptr_view_safe(signal),
+        stream,
     )
 
 
@@ -322,7 +336,18 @@ def _compile_preshuffle_to_cache(
         waves_per_eu=None if waves_per_eu <= 0 else waves_per_eu,
         xcd_swizzle=xcd_swizzle,
     )
-    _compile_executable_to_cache(exe, out, a, b, scale_a, scale_b, bias, m, n, stream)
+    _compile_executable_to_cache(
+        exe,
+        _ptr_view_safe(out),
+        _ptr_view_safe(a),
+        _ptr_view_safe(b),
+        _ptr_view_safe(scale_a),
+        _ptr_view_safe(scale_b),
+        _ptr_view_safe(bias),
+        m,
+        n,
+        stream,
+    )
 
 
 def compile_one_config(

aiter/aot/flydsl/moe.py

@@ -38,6 +38,7 @@
 from aiter.jit.core import AITER_CONFIGS
 from aiter.ops.flydsl.moe_kernels import (
     _get_compiled_silu_fused,
+    _ptr_view_safe,
     _run_compiled,
     _s1_args_fp4,
     _s1_args_std,
@@ -579,13 +580,13 @@ def _make_a_user(a_dtype_user_shape):
                 _run_compiled(
                     silu_fused,
                     (
-                        tmp_out.view(-1, inter_dim * 2),
-                        out.view(-1).view(torch.uint8),
-                        out_scale_sorted_flat,
-                        sorted_token_ids,
-                        num_valid_ids,
-                        sorted_token_ids.view(-1),
-                        torch.empty(0, device=dev, dtype=torch.float32),
+                        _ptr_view_safe(tmp_out.view(-1, inter_dim * 2)),
+                        _ptr_view_safe(out.view(-1).view(torch.uint8)),
+                        _ptr_view_safe(out_scale_sorted_flat),
+                        _ptr_view_safe(sorted_token_ids),
+                        _ptr_view_safe(num_valid_ids),
+                        _ptr_view_safe(sorted_token_ids.view(-1)),
+                        _ptr_view_safe(torch.empty(0, device=dev, dtype=torch.float32)),
                         tokens,
                         sorted_token_ids.shape[0],
                         0,

aiter/ops/flydsl/gemm_kernels.py

@@ -14,6 +14,7 @@
 from torch import Tensor
 
 import flydsl.expr as fx
+import flydsl.compiler as flyc
 from aiter import logger
 from flydsl.runtime.device import get_rocm_arch
 
@@ -67,6 +68,15 @@ def _get_dtypes():
 SPLIT_K_GLOBAL_SEMAPHORE: dict[SplitKStreamKey, torch.Tensor] = {}
 SPLIT_K_GLOBAL_SIGNAL: dict[SplitKStreamKey, torch.Tensor] = {}
 
+
+def _ptr_view_safe(t: torch.Tensor):
+    type_name = type(t).__name__
+    module_name = type(t).__module__
+    if type_name == "FakeTensor" or "fake_tensor" in module_name:
+        return flyc.from_c_void_p(fx.Uint8, 0)
+    return flyc.from_c_void_p(fx.Uint8, t.data_ptr())
+
+
 # Keep the generic auto-generated catalog aligned with the upstream FlyDSL
 # reference tuning space. The wider local one-off search space introduced
 # gfx950-faulting candidates (for example tile_k=160 and tile_n=160/192),
@@ -800,13 +810,13 @@ def launcher(
         semaphore, signal = _get_split_k_tensors(a.device, launch_stream)
         return _run_compiled(
             kernel,
-            out,
-            a,
-            b,
-            launch_bias,
+            _ptr_view_safe(out),
+            _ptr_view_safe(a),
+            _ptr_view_safe(b),
+            _ptr_view_safe(launch_bias),
             runtime_m,
-            semaphore,
-            signal,
+            _ptr_view_safe(semaphore),
+            _ptr_view_safe(signal),
             fx.Stream(launch_stream),
         )
 
@@ -1004,12 +1014,12 @@ def _as_i8(t):
     _dummy_bias = torch.empty(0, dtype=Out.dtype, device=Out.device)
     _run_compiled(
         exe,
-        out_contig.view(-1),
-        _as_i8(XQ.contiguous()).view(-1),
-        _as_i8(WQ.contiguous()).view(-1),
-        x_scale.contiguous().view(-1),
-        w_scale.contiguous().view(-1),
-        _dummy_bias,
+        _ptr_view_safe(out_contig.view(-1)),
+        _ptr_view_safe(_as_i8(XQ.contiguous()).view(-1)),
+        _ptr_view_safe(_as_i8(WQ.contiguous()).view(-1)),
+        _ptr_view_safe(x_scale.contiguous().view(-1)),
+        _ptr_view_safe(w_scale.contiguous().view(-1)),
+        _ptr_view_safe(_dummy_bias),
         m,
         n,
         fx.Stream(torch.cuda.current_stream()),

aiter/ops/flydsl/kernels/flash_attn_func_gfx1201.py

@@ -52,7 +52,6 @@
 from flydsl.utils.smem_allocator import SmemAllocator, SmemPtr
 from flydsl._mlir import ir
 from flydsl._mlir.dialects import (
-    fly as _fly,
     llvm as _llvm,
     memref as _memref,
 )
@@ -72,9 +71,10 @@ def _llvm_ptr_ty():
     return ir.Type.parse("!llvm.ptr")
 
 
-def _extract_aligned_pointer(tensor) -> ir.Value:
-    """Extract the aligned LLVM pointer from a FlyDSL tensor/memref."""
-    return _fly.extract_aligned_pointer_as_index(_llvm_ptr_ty(), _llvm_value(tensor))
+def _pointer_to_llvm_ptr(ptr) -> ir.Value:
+    """Convert a FlyDSL pointer argument to the LLVM pointer used by raw loads."""
+    ptr_i64 = arith.index_cast(T.i64, fx.ptrtoint(ptr))
+    return _llvm.IntToPtrOp(_llvm_ptr_ty(), ptr_i64).result
 
 
 def _pointer_load(result_type: ir.Type, ptr: ir.Value) -> ir.Value:
@@ -202,18 +202,18 @@ def build_flash_attn_func_module_primary(
 
     @flyc.kernel(known_block_size=[BLOCK_SIZE, 1, 1])
     def flash_attn_func_kernel(
-        Q: fx.Tensor,
-        K: fx.Tensor,
-        V: fx.Tensor,
-        O: fx.Tensor,  # noqa: E741
+        Q: fx.Pointer,
+        K: fx.Pointer,
+        V: fx.Pointer,
+        O: fx.Pointer,  # noqa: E741
         seq_len: fx.Int32,
     ):
         elem_type = dtype_to_elem_type(dtype_str)
         elem_dtype = elem_numeric_cls
-        q_ptr = _extract_aligned_pointer(Q)
-        k_ptr = _extract_aligned_pointer(K)
-        v_ptr = _extract_aligned_pointer(V)
-        o_ptr = _extract_aligned_pointer(O)
+        q_ptr = _pointer_to_llvm_ptr(Q)
+        k_ptr = _pointer_to_llvm_ptr(K)
+        v_ptr = _pointer_to_llvm_ptr(V)
+        o_ptr = _pointer_to_llvm_ptr(O)
         fm_fast = arith.FastMathFlags.fast
 
         # Local fast-math arithmetic helpers — preserve fastmath flag while using
@@ -682,10 +682,10 @@ def _load_v_rowmajor(st_kv_base_val, pks_val, dc_val):
 
     @flyc.jit
     def launch_flash_attn_func(
-        Q: fx.Tensor,
-        K: fx.Tensor,
-        V: fx.Tensor,
-        O: fx.Tensor,  # noqa: E741
+        Q: fx.Pointer,
+        K: fx.Pointer,
+        V: fx.Pointer,
+        O: fx.Pointer,  # noqa: E741
         batch_size: fx.Int32,
         seq_len: fx.Int32,
         stream: fx.Stream = fx.Stream(None),
@@ -756,18 +756,40 @@ def launch_flash_attn_func(
         "llvm_options": {"enable-post-misched": False, "lsr-drop-solution": True},
     }
 
+    def _ptr_arg(t):
+        if hasattr(t, "data_ptr"):
+            type_name = type(t).__name__
+            module_name = type(t).__module__
+            ptr = (
+                0
+                if type_name == "FakeTensor" or "fake_tensor" in module_name
+                else t.data_ptr()
+            )
+            return flyc.from_c_void_p(fx.Uint8, ptr)
+        return t
+
+    def _wrap_qkvo(args, kwargs):
+        args = list(args)
+        for idx in range(min(4, len(args))):
+            args[idx] = _ptr_arg(args[idx])
+        for name in ("Q", "K", "V", "O"):
+            if name in kwargs:
+                kwargs[name] = _ptr_arg(kwargs[name])
+        return tuple(args), kwargs
+
     def _launch(*args, **kwargs):
+        args, kwargs = _wrap_qkvo(args, kwargs)
         with CompilationContext.compile_hints(_fmha_compile_hints):
             return launch_flash_attn_func(*args, **kwargs)
 
     def _compile(Q, K, V, O, batch_size, seq_len, stream=None):  # noqa: E741
         with CompilationContext.compile_hints(_fmha_compile_hints):
             return flyc.compile(
                 launch_flash_attn_func,
-                Q,
-                K,
-                V,
-                O,
+                _ptr_arg(Q),
+                _ptr_arg(K),
+                _ptr_arg(V),
+                _ptr_arg(O),
                 batch_size,
                 seq_len,
                 fx.Stream(stream),