perf(linear_attention): halve the prefill chunk and split the scan over dv (-25% op, -1.07 s TTFT, 129 -> 49 MiB scratch)

lib/Runtime/Kernels/include/hip_custom_kernels.h

@@ -2013,6 +2013,11 @@ HIP_KERNEL_API int hip_linear_attention_decode(
 // fall back to the per-token decode loop); 0 on success; <0 on launch error.
 // Only the gated_delta rule with scalar log-decay (decay_per_key_dim==0) is
 // supported; other rules/layouts/oversized smem are declined.
+// scratch / scratch_bytes: caller-owned device scratch for the chunk-parallel
+// path (RuntimeState::la_scratch, grown on demand, freed on session cleanup).
+// Size it with hip_linear_attention_prefill_scratch_bytes() below. When null or
+// under-sized the launcher declines (returns 1) and the caller falls back to
+// the per-token loop.
 HIP_KERNEL_API int hip_linear_attention_prefill_chunked(
     void* stream,
     const void* query,
@@ -2033,7 +2038,15 @@ HIP_KERNEL_API int hip_linear_attention_prefill_chunked(
     int64_t update_rule,
     int64_t decay_per_key_dim,
     int64_t beta_per_head,
-    int64_t type);
+    int64_t type,
+    void* scratch,
+    size_t scratch_bytes);
+
+// Device-scratch bytes the chunk-parallel prefill needs for a given shape.
+// Returns 0 for shapes/params the parallel path will decline. The runtime
+// wrapper uses this to grow RuntimeState::la_scratch before the launch.
+HIP_KERNEL_API size_t hip_linear_attention_prefill_scratch_bytes(
+    int B, int seq_len, int Hkv, int dk, int dv);
 
 // Max memref rank honoured by the strided memref.copy fast path
 // (hip_strided_copy) and the host per-row fallback in memrefCopy. Defined

lib/Runtime/hipdnn_ep_runtime.h

@@ -409,6 +409,15 @@ void *hipdnn_ep_state_get_matmul_dp4a_scratch(RuntimeState *state);
 int hipdnn_ep_state_ensure_matmul_dp4a_scratch(RuntimeState *state,
                                                size_t needed_size);
 
+// Per-session scratch for the linear-attention chunk-parallel gated_delta
+// prefill (hip_linear_attention_prefill_chunked). Lazily grown via
+// hipdnn_ep_state_ensure_la_scratch (same policy as conv_scratch: never
+// shrinks, freed in hipdnn_ep_state_cleanup). Single buffer reused across all
+// linear-attention layers in the session -- safe because the stream is
+// serialised. See runtime_state_internal.h for design rationale.
+void *hipdnn_ep_state_get_la_scratch(RuntimeState *state);
+int hipdnn_ep_state_ensure_la_scratch(RuntimeState *state, size_t needed_size);
+
 // Per-op state slots (see docs/design/op-state-slots-design.md). The generated
 // @hipdnn_ep_op_states_init_fn (built by --generate-op-state-init) calls
 // _alloc once, then per stateful op calls its construct symbol; each construct

lib/Runtime/hipdnn_ep_runtime_state.cpp

@@ -165,6 +165,8 @@ static int initialize_state_handles(RuntimeState **out_state) {
   state->conv_scratch_size = 0;
   state->matmul_dp4a_scratch = nullptr;
   state->matmul_dp4a_scratch_size = 0;
+  state->la_scratch = nullptr;
+  state->la_scratch_size = 0;
   state->zp_unpack_cache = nullptr;
   state->op_profile = hipdnn_ep_perf_enabled() ? op_profile_create() : nullptr;
   state->device_error_flag = nullptr;
@@ -721,6 +723,12 @@ int hipdnn_ep_state_cleanup(RuntimeState *state) {
     HIP_CLEANUP(hipFree(state->matmul_dp4a_scratch));
   }
 
+  // Free the linear-attention chunk-parallel prefill scratch (if allocated).
+  // The stream sync above has drained any in-flight prefill still reading it.
+  if (state->la_scratch) {
+    HIP_CLEANUP(hipFree(state->la_scratch));
+  }
+
   // Tear down per-op state slots. Each entry's deletor destroys its concrete
   // type; slots reference nothing in other slots, so order is irrelevant. The
   // stream sync at the top has drained any in-flight op that may read a slot.
@@ -1415,6 +1423,53 @@ int hipdnn_ep_state_ensure_matmul_dp4a_scratch(RuntimeState *state,
   return 0;
 }
 
+// Linear-attention chunk-parallel prefill scratch pool. Same grow-on-demand /
+// never-shrink policy as conv_scratch / qmoe_scratch, freed in cleanup. Holds
+// the per-(head,chunk) Uloc/W/rlast/alast tiles + chunk-start states for one
+// prefill. Single-buffer reuse across LA layers is safe because the stream is
+// serialised: the three prefill passes fully consume it before the next launch.
+void *hipdnn_ep_state_get_la_scratch(RuntimeState *state) {
+  return state ? state->la_scratch : nullptr;
+}
+
+int hipdnn_ep_state_ensure_la_scratch(RuntimeState *state, size_t needed_size) {
+  if (!state)
+    return -1;
+  if (needed_size == 0)
+    return 0;
+  if (state->la_scratch_size >= needed_size)
+    return 0;
+
+  // 1.5x growth amortisation mirrors conv_scratch / qmoe_scratch.
+  size_t alloc_size = needed_size;
+  if (state->la_scratch_size > 0) {
+    size_t grown = state->la_scratch_size + state->la_scratch_size / 2;
+    if (grown > alloc_size)
+      alloc_size = grown;
+  }
+
+  if (state->la_scratch) {
+    // Drain any in-flight prefill still reading the old buffer before freeing.
+    // Growth is rare (only when a longer sequence is first seen).
+    if (state->stream) {
+      hipStreamSynchronize(state->stream);
+    }
+    HIP_CLEANUP(hipFree(state->la_scratch));
+    state->la_scratch = nullptr;
+    state->la_scratch_size = 0;
+  }
+
+  if (hipMalloc(&state->la_scratch, alloc_size) != hipSuccess) {
+    fprintf(stderr,
+            "hipdnn_ep_state_ensure_la_scratch: hipMalloc failed for %zu "
+            "bytes\n",
+            alloc_size);
+    return -1;
+  }
+  state->la_scratch_size = alloc_size;
+  return 0;
+}
+
 void *hipdnn_ep_state_get_error_flag_device_ptr(RuntimeState *state) {
   return state ? static_cast<void *>(state->device_error_flag) : nullptr;
 }

lib/Runtime/real/linear_attention.cpp

@@ -199,10 +199,23 @@ extern "C" int wrap_linear_attention(
   // returns 1 when it declines an unsupported config; we then fall back to the
   // per-token loop below.
   if (update_rule == kUpdateRuleGatedDelta && seq_len > 1) {
+    // The chunk-parallel path needs a device scratch arena sized to this shape.
+    // It lives in the per-session RuntimeState::la_scratch pool
+    // (grow-on-demand, freed in hipdnn_ep_state_cleanup) -- same policy as
+    // qmoe/conv scratch -- rather than a process-static buffer. If
+    // sizing/growth fails we pass a null scratch and the launcher declines
+    // (rc=1 -> per-token loop below).
+    void *la_scratch = nullptr;
+    size_t la_bytes = hip_linear_attention_prefill_scratch_bytes(
+        (int)B, (int)seq_len, (int)Hkv, (int)dk, (int)dv);
+    if (la_bytes > 0 &&
+        hipdnn_ep_state_ensure_la_scratch(state, la_bytes) == 0) {
+      la_scratch = hipdnn_ep_state_get_la_scratch(state);
+    }
     int rc = hip_linear_attention_prefill_chunked(
         hip_stream, query, key, value, decay, beta, present_state, output, B,
         seq_len, Hq, Hkv, Nk, dk, dv, scale, update_rule, decay_per_key_dim,
-        beta_per_head, type);
+        beta_per_head, type, la_scratch, la_bytes);
     if (rc == 0) {
       RUNTIME_DEBUG_LOG(
           "[linear_attention] prefill via chunked-parallel kernel (%lld "

lib/Runtime/runtime_state_internal.h

@@ -144,6 +144,18 @@ struct RuntimeState {
   void *matmul_dp4a_scratch;
   size_t matmul_dp4a_scratch_size;
 
+  // Per-session scratch for the linear-attention chunk-parallel gated_delta
+  // prefill (hip_linear_attention_prefill_chunked). One contiguous device
+  // buffer holding the per-(head,chunk) Uloc/W/rlast/alast tiles and the
+  // chunk-start states for the whole sequence. Same grow-on-demand /
+  // never-shrink policy as conv_scratch; lazily allocated on first prefill,
+  // freed in hipdnn_ep_state_cleanup. Single-buffer reuse is safe because the
+  // HIP stream is serialised (the three prefill passes consume it before the
+  // next linear-attention layer launches). Replaces a process-static buffer so
+  // the footprint is bounded to the session and released on teardown.
+  void *la_scratch;
+  size_t la_scratch_size;
+
   // NOTE: the GQA GEMM descriptor cache (GqaGemmCache) formerly lived here as
   // gqa_gemm_cache. It is now per-op-instance: each gqa instance owns one in
   // its GqaState op-state slot (see op_states below and