Add V cache write-back with interleaved KV cache layout

Extend the KV cache prefill design to write both K and V caches to DDR
during flash attention computation. Uses a single CacheWB channel with
an interleaved KV cache layout [K_c0, V_c0, K_c1, V_c1, ...] where
both K and V data are staged through kwb_buf before DMA transfer.

Key design choices:
- Single CacheWB channel avoids shim S2MM channel exhaustion (no packet
  switching needed)
- Shared kwb_buf staging buffer prevents DMA race between CacheWB read
  and V2L1 write on the v buffer
- scf.for loop in launch body enables compiler BD folding, preventing
  BD exhaustion at large sequence lengths (tested up to 12h x 4096)

Compiler changes (AIRToAIEPass.cpp):
- Fix packet BD attribute lookup for L1-to-L3 dma_packet channels
  (getExistingPacketFlowOpFromDevice searches both flow maps)
- Place outbound MM2S lock acquire before channel put and release after
  channel put, enabling interleaved lock pattern for multiple puts
  sharing the same staging buffer

Performance: 12 heads x 4096 seq_len achieves 2460 peak GFLOPS with
zero overhead vs K-only writeback.

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>

Author: erwei-xilinx

mlir/lib/Conversion/AIRToAIEPass.cpp

@@ -4685,12 +4685,7 @@ class AIRToAIEPass : public air::impl::AIRToAIEBase<AIRToAIEPass> {
       builder.setInsertionPoint(alloc.getDefiningOp());
     else if (!tileInbound.value() &&
              isa<AIE::BufferOp>(alloc.getDefiningOp())) {
-      auto br = dyn_cast_if_present<cf::BranchOp>(
-          memcpyOpIf->getBlock()->getTerminator());
-      if (br)
-        builder.setInsertionPointToStart(br.getDest());
-      else
-        builder.setInsertionPointToStart(memcpyOpIf->getBlock());
+      builder.setInsertionPoint(memcpyOpIf);
     } else
       builder.setInsertionPoint(memcpyOpIf);
 
@@ -4701,8 +4696,14 @@ class AIRToAIEPass : public air::impl::AIRToAIEBase<AIRToAIEPass> {
                            lockAqValue);
 
     // Try to find the end of lifetime for the data copied by memcpyOpIf, and
-    // put the unlock.
-    if (auto nextWriter = findNextDmaWriteOp(memcpyOpIf, alloc)) {
+    // put the unlock. For outbound puts from AIE::BufferOp, release
+    // immediately after the put to enable interleaved operation when multiple
+    // puts share the same staging buffer.
+    if (!tileInbound.value() && isa<AIE::BufferOp>(alloc.getDefiningOp())) {
+      builder.setInsertionPointAfter(memcpyOpIf);
+      AIE::UseLockOp::create(builder, memcpyOpIf->getLoc(), relLockOp,
+                             AIE::LockAction::Release, lockRelValue);
+    } else if (auto nextWriter = findNextDmaWriteOp(memcpyOpIf, alloc)) {
       // Lifetime ends if dma writes into the same buffer.
       builder.setInsertionPoint(nextWriter);
       AIE::UseLockOp::create(builder, nextWriter->getLoc(), relLockOp,

mlir/test/Conversion/AIRToAIE/air_shimcpy_to_npu.mlir

@@ -1764,3 +1764,48 @@ module {
     return
   }
 }
+
+// -----
+
+// L1-to-L3 packet flow: ChannelPut in herd (L1 source) with ChannelGet in
+// launch (L3 destination). Verifies that the packet attribute is generated on
+// the compute tile MM2S BD despite ChannelPut.getDstMemref() returning null.
+
+// CHECK:      aie.mem(%[[TILE:.*]]) {
+// CHECK:        aie.dma_start(MM2S, 0
+// CHECK:        aie.dma_bd(%{{.*}} : memref<64xbf16, 2>{{.*}}) {{{.*}}packet = #aie.packet_info<pkt_type = 0, pkt_id = 0>
+// CHECK:      aie.packet_flow(0) {
+// CHECK:        aie.packet_source<%[[TILE]], DMA : 0>
+// CHECK:        aie.packet_dest<%{{.*}}, DMA :
+
+// RACECONDFIX: @func21
+
+module {
+  air.channel @L1ToL3Pkt [1, 1] {channel_type = "dma_packet"}
+  func.func @func21(%arg0: memref<64xbf16>) {
+    %c1 = arith.constant 1 : index
+    %c0 = arith.constant 0 : index
+    %0 = air.launch async () in () args(%out=%arg0) : memref<64xbf16> attributes {id = 1 : i32} {
+      %lc0 = arith.constant 0 : index
+      // L3 destination: ChannelGet into L3 memref
+      %1 = air.channel.get async @L1ToL3Pkt[%lc0, %lc0] (%out[] [] []) {id = 1 : i32} : (memref<64xbf16>)
+      %lc1 = arith.constant 1 : index
+      %2 = air.segment @seg async attributes {id = 2 : i32, x_loc = 0 : i64, y_loc = 2 : i64} {
+        %sc1 = arith.constant 1 : index
+        %3 = air.herd @herd async tile (%tx, %ty) in (%htx=%sc1, %hty=%sc1) attributes {id = 3 : i32} {
+          %hc0 = arith.constant 0 : index
+          %async_token, %buf = air.execute -> (memref<64xbf16, 2>) {
+            %alloc = memref.alloc() : memref<64xbf16, 2>
+            air.execute_terminator %alloc : memref<64xbf16, 2>
+          }
+          // L1 source: ChannelPut from L1 buffer
+          %put = air.channel.put async [%async_token] @L1ToL3Pkt[%hc0, %hc0] (%buf[] [] []) {id = 2 : i32} : (memref<64xbf16, 2>)
+          %4 = air.execute [%put] {
+            memref.dealloc %buf : memref<64xbf16, 2>
+          }
+        }
+      }
+    }
+    return
+  }
+}