[Fix][Kernel] Reuse P as beta tiles in gated deltanet fused prepare

[3402956340]

• Fix MACA launch failure (mcErrorInvalidValue, dyn_smem_bytes=66048) in Gated DeltaNet fused prepare when chunk_size=dim_k=dim_v=64.
• After Neumann, P is dead and has shape [BC,BC]; alias k_beta/v_beta onto P when BC==DK==DV, so peak shared drops from ~64.5KB to ~50KB.
• When shapes differ (e.g. chunk=32, dim=64), keep separate alloc_shared for beta tiles.

[gemini-code-assist]

Code Review

This pull request changes the allocation of beta_shared from shared memory (T.alloc_shared) to a fragment (T.alloc_fragment). The reviewer identified a critical correctness bug introduced by this change, as fragment allocations are thread-local and will lead to cross-thread access violations when threads attempt to read distributed data. The reviewer recommends reverting this change and instead optimizing shared memory usage by reusing the P_shared buffer for non-overlapping lifetimes of other shared variables.

Important

The consumer version of Gemini Code Assist on GitHub is being sunset. Starting June 18, 2026, new organization installations will be blocked, and all code review activity will officially cease on July 17, 2026.
For more details on the timeline and next steps, please review the Help Documentation.

[gemini-code-assist]

🚨 Critical Correctness Bug & Shared Memory Optimization

Changing beta_shared to a fragment (T.alloc_fragment) introduces a critical correctness bug:

1. Thread-Local Registers: T.alloc_fragment allocates registers which are strictly thread-local.
2. Distributed Copy: When T.copy is called to load beta into beta_shared, the elements are distributed across the threads in the thread block. No single thread holds the entire beta_shared array.
3. Cross-Thread Access Violation: In the subsequent parallel loops (e.g., for i, j in T.Parallel(block_C, block_C)), threads access beta_shared[i] or beta_shared[i_s]. Since threads cannot access other threads' registers, this will result in reading incorrect/undefined data or compilation failures.

beta_shared only requires block_C * 4 bytes (256 bytes for block_C=64), so keeping it in shared memory has negligible impact on the shared memory limit.

---

💡 Correct Solution: Shared Memory Buffer Reuse

The shared memory limit is exceeded because P_shared (16 KB), k_beta_shared (16 KB), and v_beta_shared (16 KB) are all allocated concurrently, even though their lifetimes are completely non-overlapping:

• P_shared is only used in lines 90-106.
• k_beta_shared is only used in lines 114-120.
• v_beta_shared is only used in lines 123-129.

We can save 32 KB of shared memory by completely eliminating k_beta_shared and v_beta_shared and reusing P_shared as a temporary buffer.

Step 1: Update Allocations

Modify the allocations to reuse P_shared (allocated with the maximum required dimension):

                beta_shared = T.alloc_shared([block_C], accum_dtype)
                S_shared = T.alloc_shared([block_C, block_C], accum_dtype)
                # Reuse P_shared for k_beta_shared and v_beta_shared
                P_shared = T.alloc_shared([block_C, max(block_C, dim_k, dim_v)], accum_dtype)

Step 2: Update Usages

Replace k_beta_shared and v_beta_shared with P_shared in the rest of the kernel:

                # k_beta = k * beta (reusing P_shared)
                for i_s, i_k in T.Parallel(block_C, dim_k):
                    P_shared[i_s, i_k] = k_shared[i_s, i_k] * beta_shared[i_s]

                # w = A_g^{-1} @ k_beta
                T.clear(w_frag)
                T.gemm(S_shared, P_shared, w_frag)
                T.copy(w_frag, w[bid, hid, by * block_C : (by + 1) * block_C, :], disable_tma=True)

                # v_beta = v * beta (reusing P_shared)
                for i, j in T.Parallel(block_C, dim_v):
                    P_shared[i, j] = v_shared[i, j] * beta_shared[i]

                # u = A_g^{-1} @ v_beta
                T.clear(u_frag)
                T.gemm(S_shared, P_shared, u_frag)
                T.copy(u_frag, u[bid, hid, by * block_C : (by + 1) * block_C, :], disable_tma=True)

Suggested change

	beta_shared = T.alloc_fragment([block_C], accum_dtype)
	beta_shared = T.alloc_shared([block_C], accum_dtype)

[Five-HZ]

LGTM