Fix(attention): adapt the FlyDSL paged decode to aiter #5809 (explicit work plans)

atom/model_ops/attention_mha.py

@@ -624,11 +624,11 @@ def paged_attention_triton(
             sinks=self.sinks,
             sliding_window=self.sliding_window,
             ps=True,
-            # The one site whose context lengths are the real per-request ones,
-            # so the only one a work plan can rebalance -- same boundary as the
-            # `dense_decode_splits` call above. Read off the context this call
+            # The capture-owned dense plan, read off the context this call
             # already holds, never the thread-local one: a TBO worker thread
-            # that never installed its own shares the other ubatch's.
+            # that never installed its own shares the other ubatch's. No
+            # max_context_length: the block-table bound is the same in graph
+            # and eager, so both pick the same kernel specialization.
             work_plan=getattr(attn_metadata, "flydsl_work_plan", None),
         )
 
@@ -1678,6 +1678,8 @@ def _sparse_decode(
             v_scale=v_scale,
             sparse_bt=sparse_bt,
             sparse_ctx=sparse_ctx,
+            # Once per step only when the tail block is forced (see sparse_attn).
+            plan_step_owner=decode_md if self.local_blocks >= 1 else None,
         )
         self._index_q = None
         self._index_q_cache_key_info = None

atom/model_ops/attentions/aiter_attention.py

@@ -795,20 +795,19 @@ def get_kv_transfer_tensors(self):
         )
 
     def refresh_flydsl_plan(self, context_lens, *, create=False):
-        """Build or refresh aiter #5546's work plan for this forward.
+        """Build or refresh aiter's FlyDSL work plan for this forward.
 
         Depends only on context_lens, which every layer of one forward shares,
         so it is built here rather than per pa_decode call. Decode replays a
         captured graph, so the plan must be visible to the op at capture time
         (see build_for_cudagraph_capture) and its tensors must never be
         reallocated afterwards -- hence one entry per batch, kept forever, and
         refreshed in place before each replay. The refresh is a GPU kernel with
-        no readback. None when FlyDSL or the planner is off, or when the batch
-        is out of range.
+        no readback. None when FlyDSL is off, or when the batch is out of range.
         """
         # The plan only feeds FlyDSL; building one with FlyDSL off is a
         # refresh kernel per step that nothing reads.
-        if not (envs.ATOM_PA_FLYDSL and envs.ATOM_PA_FLYDSL_PLAN):
+        if not envs.ATOM_PA_FLYDSL:
             return None
         # From base_attention, not duplicated: the op checks the same bound.
         from aiter.ops.flydsl.pa_decode import plan_pa_decode
@@ -834,20 +833,21 @@ def refresh_flydsl_plan(self, context_lens, *, create=False):
         key = (n, self._flydsl_kv_heads, context_lens.device.index)
         plan = self._flydsl_plans.get(key)
         if plan is None and not create:
-            # Plans are only ever minted during cudagraph capture, where the
+            # This builder mints plans only during cudagraph capture, where the
             # batch is a ladder rung and the cost lands at startup. aiter's
             # planner takes batch as a tl.constexpr, so a new value is a kernel
             # specialization -- 65-72 ms cold -- and a plan that is never freed
             # because some captured graph may have baked its pointers in. A
-            # runtime batch with no plan is one no graph will replay, so the
-            # static path is the right answer for it rather than a stall.
+            # runtime batch with no capture-owned plan gets a runtime plan in
+            # run_pa_decode.
             # Only once a capture has happened: before the first one every
             # call lands here (profile run, eager warmup), and logging then
             # burns the one shot on a step that says nothing.
             if self._flydsl_plans and not self._flydsl_plan_unplanned:
                 self._flydsl_plan_unplanned = True
                 logger.info(
-                    "flydsl: batch %d has no plan, running the static path; "
+                    "flydsl: batch %d has no capture-owned plan; an explicit "
+                    "runtime plan will be used; "
                     "captured batches are %s",
                     n,
                     sorted({k[0] for k in self._flydsl_plans}),

atom/model_ops/attentions/gdn_attn.py

@@ -1577,9 +1577,8 @@ def build_for_cudagraph_capture(self, bs: int):
         attn_metadata.gdn_metadata = self._build_gdn_capture_metadata(bs)
 
         # Decode replays this graph, so the op must see a plan HERE: absent at
-        # capture time, the static path is what gets recorded and every later
-        # refresh feeds a graph that never reads it -- with no error, and an
-        # A/B of the planner that measures pure overhead.
+        # capture time, the graph records a per-call runtime-plan refresh
+        # instead of this builder's shared plan.
         attn_metadata.flydsl_work_plan = self.refresh_flydsl_plan(
             attn_metadata.context_lens, create=True
         )