fix: mark fused rope token tensors dynamic (#458)

Author: coderfeli

kernels/fused_rope_cache_kernel.py

@@ -431,8 +431,15 @@ def pack_fp8(vals):
                                 v_elem = vector.extract(v_e, static_position=[vi], dynamic_position=[])
                                 buffer_ops.buffer_store(v_elem, vc_rsrc, vc_nf_off)
 
+    def _mark_token_layout_dynamic(tensor):
+        if hasattr(tensor, "mark_layout_dynamic"):
+            shape = getattr(tensor, "_orig_shape", None)
+            leading_dim = len(shape) - 1 if shape is not None else -1
+            return tensor.mark_layout_dynamic(leading_dim=leading_dim)
+        return flyc.from_dlpack(tensor).mark_layout_dynamic(leading_dim=tensor.ndim - 1)
+
     @flyc.jit
-    def launch_fused_rope_cache(
+    def _jit_launch_fused_rope_cache(
         Q: fx.Tensor,
         K: fx.Tensor,
         V: fx.Tensor,
@@ -459,4 +466,39 @@ def launch_fused_rope_cache(
             stream=stream,
         )
 
+    def launch_fused_rope_cache(
+        Q,
+        K,
+        V,
+        Positions,
+        CosCache,
+        SinCache,
+        SlotMapping,
+        KeyCache,
+        ValueCache,
+        Q_out,
+        K_out,
+        num_tokens,
+        KScale,
+        VScale,
+        stream=fx.Stream(None),
+    ):
+        return _jit_launch_fused_rope_cache(
+            _mark_token_layout_dynamic(Q),
+            _mark_token_layout_dynamic(K),
+            _mark_token_layout_dynamic(V),
+            _mark_token_layout_dynamic(Positions),
+            CosCache,
+            SinCache,
+            _mark_token_layout_dynamic(SlotMapping),
+            KeyCache,
+            ValueCache,
+            _mark_token_layout_dynamic(Q_out),
+            _mark_token_layout_dynamic(K_out),
+            num_tokens,
+            KScale,
+            VScale,
+            stream=stream,
+        )
+
     return launch_fused_rope_cache

kernels/topk_gating_softmax_kernel.py

@@ -122,8 +122,9 @@ def topk_gating_softmax_kernel(
         bid = fx.block_idx.x
         tid = fx.thread_idx.x
 
-        elem_type = dtype_to_elem_type(dtype_str)
+        elem_type = dtype_to_elem_type(dtype_str).ir_type
         compute_type = T.f32
+        register_addr_space = int(fx.AddressSpace.Register)
 
         fm_fast = arith.FastMathFlags.fast
 
@@ -210,7 +211,7 @@ def group_reduce_argmax(val, idx):
         atom_reg_ty_in = fx.MemRefType.get(
             elem_type,
             fx.LayoutType.get(ELEMS_PER_ATOM, 1),
-            fx.AddressSpace.Register,
+            register_addr_space,
         )
         atom_reg_lay_in = fx.make_layout(ELEMS_PER_ATOM, 1)
 
@@ -219,7 +220,7 @@ def group_reduce_argmax(val, idx):
         # near `_store_scalar_i32` below).
         copy_atom_f32 = fx.make_copy_atom(fx.rocdl.BufferCopy32b(), 32)
         scalar_reg_ty_f32 = fx.MemRefType.get(
-            T.f32, fx.LayoutType.get(1, 1), fx.AddressSpace.Register
+            T.f32, fx.LayoutType.get(1, 1), register_addr_space
         )
         scalar_reg_lay = fx.make_layout(1, 1)