Align FROST LA with FLA/FI conventions for state layout and fix context and IMA bug

python/cudnn/_pygraph.py

@@ -2455,13 +2455,12 @@ def _moe_bwd_dweight_dims(node):
 
 
 def _linear_attention_final_state_dims(node):
-    # [N, HO, K, V]: N sequences (cu_seqlens carries N+1 boundaries), HO =
-    # max(q, v) heads — the recurrent state lives at the gate heads
+    # [N, HO, V, K]
     q, v = node.inputs["q"].dim, node.inputs["v"].dim
     cu = node.inputs.get("cu_seqlens")
     if cu is None or not cu.dim:
         return None
-    return [cu.dim[0] - 1, max(q[1], v[1]), q[2], v[2]]
+    return [cu.dim[0] - 1, max(q[1], v[1]), v[2], q[2]]
 
 
 def _linear_attention_state_checkpoints_dims(node):
@@ -2470,7 +2469,7 @@ def _linear_attention_state_checkpoints_dims(node):
     cu = node.inputs["cu_seqlens"].dim if node.inputs.get("cu_seqlens") is not None else None
     if not n or not q or not v or not cu:
         return None
-    return [max(v[0] // n, 1), max(q[1], v[1]), q[2], v[2]]
+    return [max(v[0] // n + (cu[0] - 1), 1), max(q[1], v[1]), v[2], q[2]]
 
 
 def _linear_attention_o_dims(node):

python/cudnn/fla/gated_delta_rule.py

@@ -172,8 +172,8 @@ def fallback(reason):
             return fallback("variant")
         # state_v_first only changes the recurrent-state layout, so it is a no-op
         # for a stateless (training) call; decline only when a state is exchanged.
-        if state_v_first and (initial_state is not None or output_final_state):
-            return fallback("state_v_first")
+        if not state_v_first and (initial_state is not None or output_final_state):
+            return fallback("state_v_first=False")
         if not (q.is_cuda and torch.cuda.get_device_capability(q.device)[0] >= 10):
             return fallback("pre-Blackwell")
         try:

python/cudnn/fla/kda.py

@@ -183,8 +183,8 @@ def fallback(reason):
 
         if allow_neg_eigval or cp_context is not None or return_intermediate_states:
             return fallback("variant")
-        if state_v_first and (initial_state is not None or output_final_state):
-            return fallback("state_v_first")
+        if not state_v_first and (initial_state is not None or output_final_state):
+            return fallback("state_v_first=False")
         if not (q.is_cuda and torch.cuda.get_device_capability(q.device)[0] >= 10):
             return fallback("pre-Blackwell")
         try:

python/cudnn/linear_attention/cutile/gdn_engine.py

@@ -57,7 +57,7 @@ def __init__(self, graph):
             ("A", layout.add(total * HV * BT * isz), io, (total, HV, BT)),
             ("w", layout.add(total * HV * K * isz), io, (total, HV, K)),
             ("u", layout.add(total * HV * V * isz), io, (total, HV, V)),
-            ("state_checkpoints", layout.add(NT_bound * HV * K * V * isz), io, (NT_bound, HV, K, V)),
+            ("state_checkpoints", layout.add(NT_bound * HV * K * V * isz), io, (NT_bound, HV, V, K)),
             ("v_new", layout.add(total * HV * V * isz), io, (total, HV, V)),
         ]
         if l2norm:
@@ -71,7 +71,7 @@ def __init__(self, graph):
             NK = common.cdiv(K, min(max(common.next_power_of_2(K), 16), 64))
             regions += [
                 ("dv", layout.add(total * HV * V * isz), io, (total, HV, V)),
-                ("dstate", layout.add(NT_bound * HV * K * V * isz), io, (NT_bound, HV, K, V)),
+                ("dstate", layout.add(NT_bound * HV * K * V * isz), io, (NT_bound, HV, V, K)),
                 ("dv2", layout.add(total * HV * V * isz), io, (total, HV, V)),
                 ("dg_nk", layout.add(NK * total * HV * 4), f32, (NK, total, HV)),
                 ("dw", layout.add(total * HV * K * isz), io, (total, HV, K)),
@@ -132,7 +132,6 @@ def execute(self, graph, variant_pack, ctx) -> None:
         check_layouts_compact(self.plan_name, self.expect, self.names, views)
         nb = dict(zip(self.names, views))
         stream = ctx.stream if ctx.stream is not None else 0
-        self.common.ensure_cuda_context(stream)
         ws = Workspace.over(variant_pack, self.ws_bytes, self.plan_name)
         region = dict(zip(self.carve_names, ws.carve(self.carve)))
         self.common.build_chunk_table(
@@ -167,6 +166,7 @@ def execute_fwd(self, nb, region, stream) -> None:
             cu_seqlens=nb["cu_seqlens"],
             chunk_indices=region["chunk_table"],
             bufs=region,
+            state_v_first=True,
             stream=stream,
             **gate,
         )
@@ -186,6 +186,7 @@ def execute_bwd(self, nb, region, stream) -> None:
             cu_seqlens=nb["cu_seqlens"],
             chunk_indices=region["chunk_table"],
             bufs=region,
+            state_v_first=True,
             stream=stream,
         )
 

python/cudnn/linear_attention/cutile/kda_engine.py

@@ -64,7 +64,7 @@ def __init__(self, graph):
             ("u", layout.add(total * HV * V * isz), io, (total, HV, V)),
             ("qg", layout.add(total * HV * K * isz), io, (total, HV, K)),
             ("kg", layout.add(total * HV * K * isz), io, (total, HV, K)),
-            ("state_checkpoints", layout.add(NT_bound * HV * K * V * isz), io, (NT_bound, HV, K, V)),
+            ("state_checkpoints", layout.add(NT_bound * HV * K * V * isz), io, (NT_bound, HV, V, K)),
             ("v_new", layout.add(total * HV * V * isz), io, (total, HV, V)),
         ]
         if node.params.get("use_beta_sigmoid", False):
@@ -82,7 +82,7 @@ def __init__(self, graph):
             regions += [
                 ("dAqk", layout.add(total * HV * BT * 4), f32, (total, HV, BT)),
                 ("dv_dAv", layout.add(total * HV * V * isz), io, (total, HV, V)),
-                ("dstate", layout.add(NT_bound * HV * K * V * isz), io, (NT_bound, HV, K, V)),
+                ("dstate", layout.add(NT_bound * HV * K * V * isz), io, (NT_bound, HV, V, K)),
                 ("dv_dstate_u", layout.add(total * HV * V * isz), io, (total, HV, V)),
                 ("dq", layout.add(total * HV * K * 4), f32, (total, HV, K)),
                 ("dk", layout.add(total * HV * K * 4), f32, (total, HV, K)),
@@ -158,7 +158,6 @@ def execute(self, graph, variant_pack, ctx) -> None:
         check_layouts_compact(self.plan_name, self.expect, self.names, views)
         nb = dict(zip(self.names, views))
         stream = ctx.stream if ctx.stream is not None else 0
-        self.common.ensure_cuda_context(stream)
         ws = Workspace.over(variant_pack, self.ws_bytes, self.plan_name)
         region = dict(zip(self.carve_names, ws.carve(self.carve)))
         self.common.build_chunk_table(
@@ -197,6 +196,7 @@ def execute_fwd(self, nb, region, stream) -> None:
             cu_seqlens=nb["cu_seqlens"],
             chunk_indices=region["chunk_table"],
             bufs=region,
+            state_v_first=True,
             stream=stream,
             **gate,
         )
@@ -216,6 +216,7 @@ def execute_bwd(self, nb, region, stream) -> None:
             cu_seqlens=nb["cu_seqlens"],
             chunk_indices=region["chunk_table"],
             bufs=region,
+            state_v_first=True,
             stream=stream,
         )
 

python/cudnn/linear_attention/cutile/kernels/common.py

@@ -63,35 +63,6 @@ def opt(t, bufs, dtype_name: str = "float32"):
     return t
 
 
-def ensure_cuda_context(stream=0) -> None:
-    """Bind a driver context to the calling thread when none is bound.
-
-    ``ct.launch`` and the autotuner read the calling thread's context stack,
-    and an autograd backward runs on a worker thread where ``cudaSetDevice``
-    has only moved the runtime's thread-local slot. Prefer the launch stream's
-    context, else retain the device's primary one. Best-effort: a context this
-    cannot establish fails at the launch, with the launch's own diagnostics."""
-    try:
-        from cuda.bindings import driver as drv
-
-        err, cur = drv.cuCtxGetCurrent()
-        if err == drv.CUresult.CUDA_SUCCESS and int(cur) != 0:
-            return
-        if stream:
-            err, sctx = drv.cuStreamGetCtx(stream)
-            if err == drv.CUresult.CUDA_SUCCESS:
-                drv.cuCtxSetCurrent(sctx)
-                return
-        device = current_device_id()
-        if device is None:
-            return
-        err, pctx = drv.cuDevicePrimaryCtxRetain(device)
-        if err == drv.CUresult.CUDA_SUCCESS:
-            drv.cuCtxSetCurrent(pctx)
-    except Exception:  # noqa: BLE001
-        pass
-
-
 # --- Launch tuning --------------------------------------------------------------------------------
 
 

python/cudnn/linear_attention/frost/common/thd.py

@@ -12,9 +12,6 @@
 * :func:`emit_checkpoint_seq_descs` — its per-chunk-checkpoint sibling; derives the
   per-sequence checkpoint offsets from the token ``cu_seqlens`` in place of a
   caller-computed prefix array.
-* :func:`emit_copy_desc` — verbatim single-slot copy of a fully static
-  descriptor (dense ``[N, HO, K, V]`` state; batch and head are both load
-  coordinates).
 """
 
 import cutlass
@@ -83,7 +80,7 @@ def emit_checkpoint_seq_descs(
     """Per-BATCH descriptor array for the per-chunk checkpoint tensor with the head
     axis as a descriptor dimension (``(dv, dk, chunk, head)``).  Derives the
     per-sequence checkpoint offsets from the TOKEN ``cu_seqlens`` on the fly
-    (``count_b = (seqlen_b - 1) // every_n``, running-prefix-summed) — an
+    (``count_b = (seqlen_b - 1) // every_n + 1``, running-prefix-summed) — an
     address fold no coordinate transform can express — and caps
     GLOBAL_DIM[``seq_ord``] to ``count_b``.  The head index is a load
     coordinate.  Runs on one electing thread; the calling warp elects and
@@ -95,7 +92,7 @@ def emit_checkpoint_seq_descs(
     run = cutlass.Int32(0)
     for b in cutlass.range(0, n_batch, 1, unroll=1):
         s_tok = cutlass.Int32(cu[b + cutlass.Int32(1)]) - cutlass.Int32(cu[b])
-        cnt = (s_tok - cutlass.Int32(1)) // every_n
+        cnt = (s_tok - cutlass.Int32(1)) // every_n + cutlass.Int32(1)
         cnt = cnt if s_tok > 0 else cutlass.Int32(0)
         checkpoint_base = run
         run = run + cnt
@@ -114,15 +111,3 @@ def emit_checkpoint_seq_descs(
             new_value=cnt,
             ord=seq_ord,
         )
-
-
-@cute.jit
-def emit_copy_desc(base_desc, desc_words) -> None:
-    """Verbatim single-slot copy of a fully static descriptor (e.g. the
-    dense ``[N, HO, K, V]`` initial state, whose batch and head are both
-    load coordinates).  Runs on one electing thread; the calling warp
-    elects and fences."""
-    desc_base = desc_words.iterator.raw_ptr()
-    src_words = Pointer(base_desc.get_ptr(), dtype=cutlass.Int64)
-    for i in cutlass.range_constexpr(TENSOR_MAP_QWORDS):
-        (desc_base + i).store((src_words + i).load())