Upload ai/research/cuda_kernels.py with huggingface_hub

ai/research/cuda_kernels.py

@@ -0,0 +1,1368 @@
+"""
+CUDA Kernels for GPU-Accelerated VectorEnv.
+
+This module contains CUDA kernel implementations for:
+- Environment reset
+- Game step (integrated with opponent, phases, scoring)
+- Observation encoding
+- Action mask computation
+
+All kernels are designed for the VectorGameStateGPU class.
+"""
+
+import numpy as np
+
+try:
+    from numba import cuda
+    from numba.cuda.random import xoroshiro128p_normal_float32, xoroshiro128p_uniform_float32
+
+    HAS_CUDA = True
+except ImportError:
+    HAS_CUDA = False
+
+    # Mock for type checking
+    class MockCuda:
+        def jit(self, *args, **kwargs):
+            def decorator(f):
+                return f
+
+            return decorator
+
+        def grid(self, x):
+            return 0
+
+    cuda = MockCuda()
+
+    def xoroshiro128p_uniform_float32(rng, i):
+        return 0.5
+
+
+# ============================================================================
+# CONSTANTS (Must match fast_logic.py)
+# ============================================================================
+SC = 0
+OS = 1
+TR = 2
+HD = 3
+DI = 4
+EN = 5
+DK = 6
+OT = 7
+PH = 8
+OD = 9
+
+# Opcodes
+O_DRAW = 10
+O_BLADES = 11
+O_HEARTS = 12
+O_RECOV_L = 13
+O_BOOST = 14
+O_RECOV_M = 15
+O_BUFF = 16
+O_CHARGE = 17
+O_TAP_O = 18
+O_CHOOSE = 19
+O_ADD_H = 20
+O_RETURN = 999
+O_JUMP = 100
+O_JUMP_F = 101
+
+# Conditions
+C_TR1 = 200
+C_CLR = 202
+C_STG = 203
+C_HND = 204
+C_CTR = 206
+C_LLD = 207
+C_GRP = 208
+C_OPH = 210
+C_ENR = 213
+C_CMP = 220
+
+# Unique ID (UID) System
+BASE_ID_MASK = 0xFFFFF
+
+
+@cuda.jit(device=True)
+def get_base_id_device(uid: int) -> int:
+    """Extract the base card definition ID (0-1999) from a UID."""
+    return uid & BASE_ID_MASK
+
+
+# ============================================================================
+# DEVICE FUNCTIONS (Callable from kernels)
+# ============================================================================
+
+
+@cuda.jit(device=True)
+def check_deck_refresh_device(p_deck, p_trash, p_global_ctx, DK_idx, TR_idx):
+    """Shuffle trash back into deck if deck is empty."""
+    if p_global_ctx[DK_idx] <= 0:
+        # Count trash
+        tr_count = 0
+        for t in range(60):
+            if p_trash[t] > 0:
+                tr_count += 1
+
+        if tr_count > 0:
+            # Move trash to deck
+            d_ptr = 0
+            for t in range(60):
+                if p_trash[t] > 0:
+                    p_deck[d_ptr] = p_trash[t]
+                    p_trash[t] = 0
+                    d_ptr += 1
+
+            p_global_ctx[DK_idx] = d_ptr
+            p_global_ctx[TR_idx] = 0
+
+
+@cuda.jit(device=True)
+def move_to_trash_device(card_id, p_trash, p_global_ctx, TR_idx):
+    """Move a card to trash zone."""
+    for t in range(60):
+        if p_trash[t] == 0:
+            p_trash[t] = card_id
+            p_global_ctx[TR_idx] += 1
+            break
+
+
+@cuda.jit(device=True)
+def draw_cards_device(count, p_hand, p_deck, p_trash, p_global_ctx):
+    """Draw cards from deck to hand."""
+    for _ in range(count):
+        check_deck_refresh_device(p_deck, p_trash, p_global_ctx, DK, TR)
+
+        if p_global_ctx[DK] <= 0:
+            break
+
+        # Find top card in deck
+        top_card = 0
+        d_idx_found = -1
+        for d in range(60):
+            if p_deck[d] > 0:
+                top_card = p_deck[d]
+                d_idx_found = d
+                break
+
+        if top_card > 0:
+            # Find empty hand slot
+            for h in range(60):
+                if p_hand[h] == 0:
+                    p_hand[h] = top_card
+                    p_deck[d_idx_found] = 0
+                    p_global_ctx[DK] -= 1
+                    p_global_ctx[HD] += 1
+                    break
+
+
+@cuda.jit(device=True)
+def resolve_bytecode_device(
+    bytecode,
+    flat_ctx,
+    global_ctx,
+    player_id,
+    p_hand,
+    p_deck,
+    p_stage,
+    p_energy_vec,
+    p_energy_count,
+    p_cont_vec,
+    p_cont_ptr,
+    p_tapped,
+    p_live,
+    opp_tapped,
+    p_trash,
+    bytecode_map,
+    bytecode_index,
+):
+    """
+    GPU Device function for resolving bytecode.
+    Returns (new_cont_ptr, status, bonus).
+    """
+    ip = 0
+    cptr = p_cont_ptr
+    bonus = 0
+    cond = True
+    blen = bytecode.shape[0]
+    safety_counter = 0
+
+    while ip < blen and safety_counter < 500:
+        safety_counter += 1
+        op = bytecode[ip, 0]
+        v = bytecode[ip, 1]
+        a = bytecode[ip, 2]
+        s = bytecode[ip, 3]
+
+        if op == 0:
+            ip += 1
+            continue
+        if op == O_RETURN:
+            break
+
+        # Jumps
+        if op == O_JUMP:
+            new_ip = ip + v
+            if 0 <= new_ip < blen:
+                ip = new_ip
+            else:
+                break
+            continue
+
+        if op == O_JUMP_F:
+            if not cond:
+                new_ip = ip + v
+                if 0 <= new_ip < blen:
+                    ip = new_ip
+                else:
+                    break
+                continue
+            ip += 1
+            continue
+
+        # Conditions (op >= 200)
+        if op >= 200:
+            if op == C_TR1:
+                cond = global_ctx[TR] == 1
+            elif op == C_STG:
+                ct = 0
+                for j in range(3):
+                    if p_stage[j] != -1:
+                        ct += 1
+                cond = ct >= v
+            elif op == C_HND:
+                cond = global_ctx[HD] >= v
+            elif op == C_LLD:
+                cond = global_ctx[SC] > global_ctx[OS]
+            elif op == C_ENR:
+                cond = global_ctx[EN] >= v
+            elif op == C_CMP:
+                if v > 0:
+                    cond = global_ctx[SC] >= v
+                else:
+                    cond = global_ctx[SC] > global_ctx[OS]
+            elif op == C_OPH:
+                cond = global_ctx[OT] >= v if v > 0 else global_ctx[OT] > 0
+            else:
+                cond = True
+            ip += 1
+        else:
+            # Effects
+            if cond:
+                if op == O_DRAW:
+                    draw_cards_device(v, p_hand, p_deck, p_trash, global_ctx)
+                elif op == O_CHARGE:
+                    # Move cards from deck to energy (simplified)
+                    amt = min(v, global_ctx[DK])
+                    for _ in range(amt):
+                        for d in range(60):
+                            if p_deck[d] > 0:
+                                p_deck[d] = 0
+                                global_ctx[DK] -= 1
+                                global_ctx[EN] += 1
+                                break
+                elif op == O_HEARTS:
+                    # Add hearts (points)
+                    bonus += v
+                    # Register continuous effect
+                    if cptr < 32:
+                        p_cont_vec[cptr, 0] = 2
+                        p_cont_vec[cptr, 1] = v
+                        p_cont_vec[cptr, 5] = a
+                        p_cont_vec[cptr, 9] = 1
+                        cptr += 1
+                elif op == O_BLADES:
+                    if cptr < 32:
+                        p_cont_vec[cptr, 0] = 1
+                        p_cont_vec[cptr, 1] = v
+                        p_cont_vec[cptr, 2] = 4
+                        p_cont_vec[cptr, 3] = s
+                        p_cont_vec[cptr, 9] = 1
+                        cptr += 1
+                elif op == O_RECOV_M:
+                    if 0 <= s < 3:
+                        p_tapped[s] = 0
+                elif op == O_RECOV_L:
+                    if 0 <= s < p_live.shape[0]:
+                        p_live[s] = 0
+                elif op == O_TAP_O:
+                    if 0 <= s < 3:
+                        opp_tapped[s] = 1
+                elif op == O_BUFF:
+                    if cptr < 32:
+                        p_cont_vec[cptr, 0] = 8
+                        p_cont_vec[cptr, 1] = v
+                        p_cont_vec[cptr, 2] = s
+                        p_cont_vec[cptr, 9] = 1
+                        cptr += 1
+                elif op == O_BOOST:
+                    bonus += v
+            ip += 1
+
+    return cptr, 0, bonus
+
+
+@cuda.jit(device=True)
+def step_player_device(
+    act_id,
+    player_id,
+    rng_state,
+    i,
+    p_hand,
+    p_deck,
+    p_stage,
+    p_energy_vec,
+    p_energy_count,
+    p_tapped,
+    p_live,
+    p_scores,
+    p_global_ctx,
+    p_trash,
+    p_continuous_vec,
+    p_continuous_ptr,
+    opp_tapped,
+    card_stats,
+    bytecode_map,
+    bytecode_index,
+):
+    """
+    Device function for single player step.
+    Returns bonus score from this action.
+    """
+    bonus = 0
+
+    if act_id == 0:
+        # Pass -> Next Phase
+        ph = p_global_ctx[PH]
+        if ph == -1:
+            p_global_ctx[PH] = 0
+        elif ph == 0:
+            p_global_ctx[PH] = 4  # Skip to Main
+        elif ph == 4:
+            p_global_ctx[PH] = 8  # Performance
+        return 0
+
+    # Member Play (1-180)
+    if 1 <= act_id <= 180:
+        adj = act_id - 1
+        hand_idx = adj // 3
+        slot = adj % 3
+
+        if hand_idx < 60:
+            card_id = p_hand[hand_idx]
+            if card_id >= 0:
+                bid = get_base_id_device(card_id)
+                if bid < card_stats.shape[0]:
+                    # Cost calculation
+                    cost = card_stats[bid, 0]
+                    effective_cost = cost
+                    prev_cid = p_stage[slot]
+                    if prev_cid >= 0:
+                        prev_bid = get_base_id_device(prev_cid)
+                        if prev_bid < card_stats.shape[0]:
+                            prev_cost = card_stats[prev_bid, 0]
+                            effective_cost = max(0, cost - prev_cost)
+
+                # Pay cost by tapping energy
+                ec = min(p_global_ctx[EN], 12)
+                paid = 0
+                if effective_cost > 0:
+                    for e_idx in range(ec):
+                        if 3 + e_idx < 16:
+                            if p_tapped[3 + e_idx] == 0:
+                                p_tapped[3 + e_idx] = 1
+                                paid += 1
+                                if paid >= effective_cost:
+                                    break
+
+                # Move to stage
+                p_stage[slot] = card_id
+                p_hand[hand_idx] = 0
+                p_global_ctx[HD] -= 1
+                p_global_ctx[51 + slot] = 1  # Mark played
+
+                # Resolve auto-ability
+                bid = get_base_id_device(card_id)
+                if bid < bytecode_index.shape[0]:
+                    map_idx = bytecode_index[bid, 0]
+                    if map_idx >= 0:
+                        flat_ctx = cuda.local.array(64, dtype=np.int32)
+                        for j in range(64):
+                            flat_ctx[j] = 0
+
+                        new_ptr, _, ab_bonus = resolve_bytecode_device(
+                            bytecode_map[map_idx],
+                            flat_ctx,
+                            p_global_ctx,
+                            player_id,
+                            p_hand,
+                            p_deck,
+                            p_stage,
+                            p_energy_vec,
+                            p_energy_count,
+                            p_continuous_vec,
+                            p_continuous_ptr[0],
+                            p_tapped,
+                            p_live,
+                            opp_tapped,
+                            p_trash,
+                            bytecode_map,
+                            bytecode_index,
+                        )
+                        p_continuous_ptr[0] = new_ptr
+                        bonus += ab_bonus
+
+    # Activate Ability (200-202)
+    elif 200 <= act_id <= 202:
+        slot = act_id - 200
+        card_id = p_stage[slot]
+        if card_id >= 0 and p_tapped[slot] == 0:
+            bid = get_base_id_device(card_id)
+            if bid < bytecode_index.shape[0]:
+                map_idx = bytecode_index[bid, 0]
+                if map_idx >= 0:
+                    flat_ctx = cuda.local.array(64, dtype=np.int32)
+                    for j in range(64):
+                        flat_ctx[j] = 0
+
+                    new_ptr, _, ab_bonus = resolve_bytecode_device(
+                        bytecode_map[map_idx],
+                        flat_ctx,
+                        p_global_ctx,
+                        player_id,
+                        p_hand,
+                        p_deck,
+                        p_stage,
+                        p_energy_vec,
+                        p_energy_count,
+                        p_continuous_vec,
+                        p_continuous_ptr[0],
+                        p_tapped,
+                        p_live,
+                        opp_tapped,
+                        p_trash,
+                        bytecode_map,
+                        bytecode_index,
+                    )
+                    p_continuous_ptr[0] = new_ptr
+                    bonus += ab_bonus
+                    p_tapped[slot] = 1
+
+    # Set Live Card (400-459)
+    elif 400 <= act_id <= 459:
+        hand_idx = act_id - 400
+        if hand_idx < 60:
+            card_id = p_hand[hand_idx]
+            if card_id > 0:
+                # Find empty live zone slot
+                for j in range(50):
+                    if p_live[j] == 0:
+                        p_live[j] = card_id
+                        p_hand[hand_idx] = 0
+                        p_global_ctx[HD] -= 1
+                        break
+
+    return bonus
+
+
+@cuda.jit(device=True)
+def resolve_live_device(
+    live_id, p_stage, p_live, p_scores, p_global_ctx, p_deck, p_hand, p_trash, card_stats, p_cont_vec, p_cont_ptr
+):
+    """
+    Device function to resolve a live card.
+    Returns the score value if successful, 0 otherwise.
+    """
+    bid = get_base_id_device(live_id)
+    if live_id < 0 or bid >= card_stats.shape[0]:
+        return 0
+
+    # Get required hearts from card_stats (indices 12-18)
+    required = cuda.local.array(7, dtype=np.int32)
+    for c in range(7):
+        required[c] = card_stats[bid, 12 + c]
+
+    total_required = 0
+    for c in range(7):
+        total_required += required[c]
+
+    if total_required <= 0:
+        # No requirements - auto-succeed
+        return card_stats[bid, 38]  # Score value
+
+    # Calculate provided hearts from stage members
+    provided = cuda.local.array(7, dtype=np.int32)
+    for c in range(7):
+        provided[c] = 0
+
+    for slot in range(3):
+        cid = p_stage[slot]
+        if cid > 0:
+            s_bid = get_base_id_device(cid)
+            if s_bid < card_stats.shape[0]:
+                for c in range(7):
+                    provided[c] += card_stats[s_bid, 12 + c]
+
+    # Check if requirements met
+    for c in range(6):  # Colors (not All)
+        if required[c] > provided[c]:
+            return 0  # Failed
+
+    # All requirements met
+    return card_stats[bid, 38]
+
+
+@cuda.jit(device=True)
+def run_opponent_turn_device(
+    rng_state,
+    i,
+    opp_hand,
+    opp_deck,
+    opp_stage,
+    opp_energy_vec,
+    opp_energy_count,
+    opp_tapped,
+    opp_live,
+    opp_scores,
+    opp_global_ctx,
+    opp_trash,
+    p_tapped,
+    opp_history,
+    card_stats,
+    bytecode_map,
+    bytecode_index,
+):
+    """
+    Simple heuristic opponent turn.
+    Plays members if possible, activates abilities, sets lives.
+    """
+    # Play up to 3 members in empty slots
+    for slot in range(3):
+        if opp_stage[slot] == -1:
+            # Find playable member in hand
+            for h in range(60):
+                cid = opp_hand[h]
+                if cid >= 0:
+                    bid = get_base_id_device(cid)
+                    if bid < card_stats.shape[0]:
+                        ctype = card_stats[bid, 10]
+                        if ctype == 1:  # Member
+                            cost = card_stats[bid, 0]
+                            if cost <= opp_global_ctx[EN]:
+                                # Play it
+                                opp_stage[slot] = cid
+                                opp_hand[h] = 0
+                                opp_global_ctx[HD] -= 1
+                                # Update History
+                                for k in range(5, 0, -1):
+                                    opp_history[i, k] = opp_history[i, k - 1]
+                                opp_history[i, 0] = cid
+                                break
+
+    # Set a live card if possible
+    for h in range(60):
+        cid = opp_hand[h]
+        if cid >= 0:
+            bid = get_base_id_device(cid)
+            if bid < card_stats.shape[0]:
+                ctype = card_stats[bid, 10]
+                if ctype == 2:  # Live
+                    for lz in range(50):
+                        if opp_live[lz] == 0:
+                            opp_live[lz] = cid
+                            opp_hand[h] = 0
+                            opp_global_ctx[HD] -= 1
+                            # Update History
+                            for k in range(5, 0, -1):
+                                opp_history[i, k] = opp_history[i, k - 1]
+                            opp_history[i, 0] = cid
+                            break
+                    break
+
+
+# ============================================================================
+# MAIN KERNELS
+# ============================================================================
+
+
+@cuda.jit
+def reset_kernel(
+    indices,
+    batch_stage,
+    batch_energy_vec,
+    batch_energy_count,
+    batch_continuous_vec,
+    batch_continuous_ptr,
+    batch_tapped,
+    batch_live,
+    batch_scores,
+    batch_flat_ctx,
+    batch_global_ctx,
+    batch_hand,
+    batch_deck,
+    batch_trash,
+    batch_opp_history,
+    opp_stage,
+    opp_energy_vec,
+    opp_energy_count,
+    opp_tapped,
+    opp_live,
+    opp_scores,
+    opp_global_ctx,
+    opp_hand,
+    opp_deck,
+    opp_trash,
+    ability_member_ids,
+    ability_live_ids,
+    rng_states,
+    force_start_order,
+    obs_buffer,
+    card_stats,
+):
+    """
+    CUDA Kernel to reset environments.
+    """
+    tid = cuda.grid(1)
+    if tid >= indices.shape[0]:
+        return
+
+    i = indices[tid]
+
+    # Clear agent state
+    for j in range(3):
+        batch_stage[i, j] = -1
+    for j in range(3):
+        for k in range(32):
+            batch_energy_vec[i, j, k] = 0
+        batch_energy_count[i, j] = 0
+    for j in range(32):
+        for k in range(10):
+            batch_continuous_vec[i, j, k] = 0
+    batch_continuous_ptr[i] = 0
+    for j in range(16):
+        batch_tapped[i, j] = 0
+    for j in range(50):
+        batch_live[i, j] = 0
+    batch_scores[i] = 0
+    for j in range(64):
+        batch_flat_ctx[i, j] = 0
+    for j in range(128):
+        batch_global_ctx[i, j] = 0
+    for j in range(60):
+        batch_trash[i, j] = 0
+    for j in range(6):
+        batch_opp_history[i, j] = 0
+
+    # Clear opponent state
+    for j in range(3):
+        opp_stage[i, j] = -1
+    for j in range(3):
+        for k in range(32):
+            opp_energy_vec[i, j, k] = 0
+        opp_energy_count[i, j] = 0
+    for j in range(16):
+        opp_tapped[i, j] = 0
+    for j in range(50):
+        opp_live[i, j] = 0
+    opp_scores[i] = 0
+    for j in range(128):
+        opp_global_ctx[i, j] = 0
+    for j in range(60):
+        opp_trash[i, j] = 0
+
+    # Generate deck
+    n_members = ability_member_ids.shape[0]
+    n_lives = ability_live_ids.shape[0]
+
+    # Members (0-47)
+    for k in range(48):
+        if n_members == 48:
+            batch_deck[i, k] = ability_member_ids[k]
+            opp_deck[i, k] = ability_member_ids[k]
+        else:
+            # Random pick using RNG
+            r = xoroshiro128p_uniform_float32(rng_states, i)
+            idx = int(r * n_members) % n_members
+            batch_deck[i, k] = ability_member_ids[idx]
+            r = xoroshiro128p_uniform_float32(rng_states, i)
+            idx = int(r * n_members) % n_members
+            opp_deck[i, k] = ability_member_ids[idx]
+
+    # Lives (48-59)
+    for k in range(12):
+        if n_lives == 12:
+            batch_deck[i, 48 + k] = ability_live_ids[k]
+            opp_deck[i, 48 + k] = ability_live_ids[k]
+        else:
+            r = xoroshiro128p_uniform_float32(rng_states, i)
+            idx = int(r * n_lives) % n_lives
+            batch_deck[i, 48 + k] = ability_live_ids[idx]
+            r = xoroshiro128p_uniform_float32(rng_states, i)
+            idx = int(r * n_lives) % n_lives
+            opp_deck[i, 48 + k] = ability_live_ids[idx]
+
+    # Shuffle decks (Fisher-Yates)
+    for k in range(59, 0, -1):
+        r = xoroshiro128p_uniform_float32(rng_states, i)
+        j = int(r * (k + 1)) % (k + 1)
+        tmp = batch_deck[i, k]
+        batch_deck[i, k] = batch_deck[i, j]
+        batch_deck[i, j] = tmp
+
+        r = xoroshiro128p_uniform_float32(rng_states, i)
+        j = int(r * (k + 1)) % (k + 1)
+        tmp = opp_deck[i, k]
+        opp_deck[i, k] = opp_deck[i, j]
+        opp_deck[i, j] = tmp
+
+    # Place 2 cards in Live Zone
+    batch_live[i, 0] = batch_deck[i, 0]
+    batch_live[i, 1] = batch_deck[i, 1]
+    batch_deck[i, 0] = 0
+    batch_deck[i, 1] = 0
+
+    opp_live[i, 0] = opp_deck[i, 0]
+    opp_live[i, 1] = opp_deck[i, 1]
+    opp_deck[i, 0] = 0
+    opp_deck[i, 1] = 0
+
+    # Draw hand (6 cards)
+    for j in range(60):
+        batch_hand[i, j] = 0
+        opp_hand[i, j] = 0
+
+    drawn = 0
+    for k in range(2, 60):
+        if batch_deck[i, k] > 0 and drawn < 6:
+            batch_hand[i, drawn] = batch_deck[i, k]
+            batch_deck[i, k] = 0
+            drawn += 1
+
+    drawn_o = 0
+    for k in range(2, 60):
+        if opp_deck[i, k] > 0 and drawn_o < 6:
+            opp_hand[i, drawn_o] = opp_deck[i, k]
+            opp_deck[i, k] = 0
+            drawn_o += 1
+
+    # Set initial global context
+    batch_global_ctx[i, HD] = 6
+    batch_global_ctx[i, DK] = 52
+    batch_global_ctx[i, EN] = 3
+    batch_global_ctx[i, PH] = 4  # Start in Main phase (simplified)
+    batch_global_ctx[i, 54] = 1  # Turn 1
+
+    opp_global_ctx[i, HD] = 6
+    opp_global_ctx[i, DK] = 52
+    opp_global_ctx[i, EN] = 3
+    opp_global_ctx[i, PH] = 4
+    opp_global_ctx[i, 54] = 1
+
+    # Start order
+    if force_start_order == -1:
+        r = xoroshiro128p_uniform_float32(rng_states, i)
+        is_second = 1 if r > 0.5 else 0
+    else:
+        is_second = force_start_order
+    batch_global_ctx[i, 10] = is_second
+
+
+@cuda.jit
+def step_kernel(
+    num_envs,
+    actions,
+    batch_hand,
+    batch_deck,
+    batch_stage,
+    batch_energy_vec,
+    batch_energy_count,
+    batch_continuous_vec,
+    batch_continuous_ptr,
+    batch_tapped,
+    batch_live,
+    batch_scores,
+    batch_flat_ctx,
+    batch_global_ctx,
+    opp_hand,
+    opp_deck,
+    opp_stage,
+    opp_energy_vec,
+    opp_energy_count,
+    opp_tapped,
+    opp_live,
+    opp_scores,
+    opp_global_ctx,
+    card_stats,
+    bytecode_map,
+    bytecode_index,
+    obs_buffer,
+    rewards,
+    dones,
+    prev_scores,
+    prev_opp_scores,
+    prev_phases,
+    terminal_obs_buffer,
+    batch_trash,
+    opp_trash,
+    batch_opp_history,
+    term_scores_agent,
+    term_scores_opp,
+    ability_member_ids,
+    ability_live_ids,
+    rng_states,
+    game_config,
+    opp_mode,
+    force_start_order,
+):
+    """
+    Main integrated step kernel.
+    Processes one environment per thread.
+    """
+    i = cuda.grid(1)
+    if i >= num_envs:
+        return
+
+    # Config
+    CFG_TURN_LIMIT = int(game_config[0])
+    CFG_STEP_LIMIT = int(game_config[1])
+    CFG_REWARD_WIN = game_config[2]
+    CFG_REWARD_LOSE = game_config[3]
+    CFG_REWARD_SCALE = game_config[4]
+    CFG_REWARD_TURN_PENALTY = game_config[5]
+
+    act_id = actions[i]
+    ph = int(batch_global_ctx[i, PH])
+
+    # Sync score to context
+    batch_global_ctx[i, SC] = batch_scores[i]
+
+    # Increment step counter
+    batch_global_ctx[i, 58] += 1
+
+    # Get continuous pointer slice
+    cont_ptr_arr = batch_continuous_ptr[i : i + 1]
+    score_arr = batch_scores[i : i + 1]
+
+    # Execute action
+    bonus = step_player_device(
+        act_id,
+        0,
+        rng_states,
+        i,
+        batch_hand[i],
+        batch_deck[i],
+        batch_stage[i],
+        batch_energy_vec[i],
+        batch_energy_count[i],
+        batch_tapped[i],
+        batch_live[i],
+        score_arr,
+        batch_global_ctx[i],
+        batch_trash[i],
+        batch_continuous_vec[i],
+        cont_ptr_arr,
+        opp_tapped[i],
+        card_stats,
+        bytecode_map,
+        bytecode_index,
+    )
+    batch_scores[i] += bonus
+
+    # Handle turn end (Pass in Main Phase)
+    if act_id == 0 and ph == 4:
+        # Run opponent turn
+        run_opponent_turn_device(
+            rng_states,
+            i,
+            opp_hand[i],
+            opp_deck[i],
+            opp_stage[i],
+            opp_energy_vec[i],
+            opp_energy_count[i],
+            opp_tapped[i],
+            opp_live[i],
+            opp_scores[i : i + 1],
+            opp_global_ctx[i],
+            opp_trash[i],
+            batch_tapped[i],
+            batch_opp_history,
+            card_stats,
+            bytecode_map,
+            bytecode_index,
+        )
+
+        # Resolve lives for both players
+        agent_live_score = 0
+        opp_live_score = 0
+
+        for z in range(10):
+            lid = batch_live[i, z]
+            if lid > 0:
+                s = resolve_live_device(
+                    lid,
+                    batch_stage[i],
+                    batch_live[i],
+                    batch_scores[i : i + 1],
+                    batch_global_ctx[i],
+                    batch_deck[i],
+                    batch_hand[i],
+                    batch_trash[i],
+                    card_stats,
+                    batch_continuous_vec[i],
+                    cont_ptr_arr,
+                )
+                agent_live_score += s
+                # Clear used live
+                if s > 0:
+                    move_to_trash_device(lid, batch_trash[i], batch_global_ctx[i], TR)
+                    batch_live[i, z] = 0
+
+        for z in range(10):
+            lid = opp_live[i, z]
+            if lid > 0:
+                s = resolve_live_device(
+                    lid,
+                    opp_stage[i],
+                    opp_live[i],
+                    opp_scores[i : i + 1],
+                    opp_global_ctx[i],
+                    opp_deck[i],
+                    opp_hand[i],
+                    opp_trash[i],
+                    card_stats,
+                    batch_continuous_vec[i],
+                    cont_ptr_arr,
+                )
+                opp_live_score += s
+                if s > 0:
+                    move_to_trash_device(lid, opp_trash[i], opp_global_ctx[i], TR)
+                    opp_live[i, z] = 0
+
+        # Scoring comparison
+        if agent_live_score > 0 and opp_live_score == 0:
+            batch_scores[i] += 1
+        elif agent_live_score == 0 and opp_live_score > 0:
+            opp_scores[i] += 1
+        elif agent_live_score > 0 and opp_live_score > 0:
+            if agent_live_score > opp_live_score:
+                batch_scores[i] += 1
+            elif opp_live_score > agent_live_score:
+                opp_scores[i] += 1
+            else:
+                # Tie - both score
+                batch_scores[i] += 1
+                opp_scores[i] += 1
+
+        # Next turn setup
+        batch_global_ctx[i, 54] += 1
+        opp_global_ctx[i, 54] += 1
+
+        # Untap and energy
+        for j in range(16):
+            batch_tapped[i, j] = 0
+            if j < opp_tapped.shape[1]:
+                opp_tapped[i, j] = 0
+
+        batch_global_ctx[i, EN] = min(batch_global_ctx[i, EN] + 1, 12)
+        opp_global_ctx[i, EN] = min(opp_global_ctx[i, EN] + 1, 12)
+
+        # Draw card
+        draw_cards_device(1, batch_hand[i], batch_deck[i], batch_trash[i], batch_global_ctx[i])
+        draw_cards_device(1, opp_hand[i], opp_deck[i], opp_trash[i], opp_global_ctx[i])
+
+    # Calculate rewards
+    current_score = batch_scores[i]
+    score_diff = float(current_score) - float(prev_scores[i])
+    opp_score_diff = float(opp_scores[i]) - float(prev_opp_scores[i])
+
+    r = (score_diff * CFG_REWARD_SCALE) - (opp_score_diff * CFG_REWARD_SCALE)
+    r += CFG_REWARD_TURN_PENALTY
+
+    win = current_score >= 3
+    lose = opp_scores[i] >= 3
+
+    if win:
+        r += CFG_REWARD_WIN
+    if lose:
+        r += CFG_REWARD_LOSE
+
+    rewards[i] = r
+
+    # Sync Opp Stats to Agent Context (for Attention features)
+    batch_global_ctx[i, 4] = opp_global_ctx[i, 3]  # HD
+    batch_global_ctx[i, 9] = opp_global_ctx[i, 6]  # DK
+    batch_global_ctx[i, 7] = opp_global_ctx[i, 2]  # TR
+
+    # Check done
+    is_done = win or lose or batch_global_ctx[i, 54] >= CFG_TURN_LIMIT or batch_global_ctx[i, 58] >= CFG_STEP_LIMIT
+    dones[i] = is_done
+
+    if is_done:
+        term_scores_agent[i] = batch_scores[i]
+        term_scores_opp[i] = opp_scores[i]
+        # Note: Auto-reset should be called separately
+
+    # Update prev scores
+    prev_scores[i] = batch_scores[i]
+    prev_opp_scores[i] = opp_scores[i]
+
+
+@cuda.jit
+def compute_action_masks_kernel(
+    num_envs,
+    batch_hand,
+    batch_stage,
+    batch_tapped,
+    batch_global_ctx,
+    batch_live,
+    card_stats,
+    masks,  # Output: (N, 2000)
+):
+    """
+    Compute legal action masks on GPU.
+    """
+    i = cuda.grid(1)
+    if i >= num_envs:
+        return
+
+    # Reset all to False
+    for a in range(2000):
+        masks[i, a] = False
+
+    ph = batch_global_ctx[i, PH]
+
+    # Action 0: Pass is always legal in Main Phase
+    if ph == 4:
+        masks[i, 0] = True
+
+    # Member Play (1-180): HandIdx * 3 + Slot + 1
+    for h_idx in range(60):
+        cid = batch_hand[i, h_idx]
+        if cid > 0:
+            bid = get_base_id_device(cid)
+            if bid < card_stats.shape[0]:
+                ctype = card_stats[bid, 10]
+                cost = card_stats[bid, 0]
+
+                if ctype == 1:  # Member
+                    for slot in range(3):
+                        # Check if slot empty or can upgrade
+                        old_cid = batch_stage[i, slot]
+                        effective_cost = cost
+                        if old_cid >= 0:
+                            old_bid = get_base_id_device(old_cid)
+                            if old_bid < card_stats.shape[0]:
+                                effective_cost = max(0, cost - card_stats[old_bid, 0])
+
+                        # Check energy
+                        available_energy = 0
+                        for e in range(12):
+                            if batch_tapped[i, 3 + e] == 0:
+                                available_energy += 1
+
+                        if available_energy >= effective_cost:
+                            action_id = h_idx * 3 + slot + 1
+                            if action_id < 181:
+                                masks[i, action_id] = True
+
+    # Activate Ability (200-202)
+    for slot in range(3):
+        cid = batch_stage[i, slot]
+        if cid > 0 and batch_tapped[i, slot] == 0:
+            masks[i, 200 + slot] = True
+
+    # Set Live (400-459)
+    for h_idx in range(60):
+        cid = batch_hand[i, h_idx]
+        if cid > 0:
+            bid = get_base_id_device(cid)
+            if bid < card_stats.shape[0]:
+                ctype = card_stats[bid, 10]
+                if ctype == 2:  # Live
+                    # Check if there's an empty live zone slot
+                    for lz_idx in range(50):
+                        if batch_live[i, lz_idx] == 0:
+                            if h_idx < 60:  # This check is redundant due to outer loop
+                                masks[i, 400 + h_idx] = True
+                                break  # Only need one empty slot to make it legal
+
+
+@cuda.jit
+def encode_observations_kernel(
+    num_envs,
+    batch_hand,
+    batch_stage,
+    batch_energy_count,
+    batch_tapped,
+    batch_scores,
+    opp_scores,
+    opp_stage,
+    opp_tapped,
+    card_stats,
+    batch_global_ctx,
+    batch_live,
+    turn_number,
+    obs_buffer,
+):
+    """
+    Encode observations on GPU (STANDARD mode).
+    """
+    i = cuda.grid(1)
+    if i >= num_envs:
+        return
+
+    obs_dim = obs_buffer.shape[1]
+
+    # Clear observation
+    for j in range(obs_dim):
+        obs_buffer[i, j] = 0.0
+
+    # Metadata
+    obs_buffer[i, 0] = float(batch_scores[i]) / 3.0
+    obs_buffer[i, 1] = float(opp_scores[i]) / 3.0
+    obs_buffer[i, 2] = float(batch_global_ctx[i, EN]) / 12.0
+    obs_buffer[i, 3] = float(batch_global_ctx[i, HD]) / 60.0
+    obs_buffer[i, 4] = float(batch_global_ctx[i, DK]) / 60.0
+    obs_buffer[i, 5] = float(batch_global_ctx[i, 54]) / 100.0  # Turn
+
+    offset = 10
+
+    # Stage (3 slots x 20 features)
+    for slot in range(3):
+        cid = batch_stage[i, slot]
+        base = offset + slot * 20
+        if cid > 0:
+            bid = get_base_id_device(cid)
+            if bid < card_stats.shape[0]:
+                obs_buffer[i, base] = 1.0  # Presence
+                obs_buffer[i, base + 1] = float(cid) / 2000.0
+                obs_buffer[i, base + 2] = float(card_stats[bid, 0]) / 10.0  # Cost
+                obs_buffer[i, base + 3] = float(card_stats[bid, 1]) / 5.0  # Blades
+                obs_buffer[i, base + 4] = float(card_stats[bid, 2]) / 10.0  # Hearts
+                obs_buffer[i, base + 5] = 1.0 if batch_tapped[i, slot] > 0 else 0.0
+
+    offset += 60
+
+    # Opponent Stage
+    for slot in range(3):
+        cid = opp_stage[i, slot]
+        base = offset + slot * 20
+        if cid > 0:
+            bid = get_base_id_device(cid)
+            if bid < card_stats.shape[0]:
+                obs_buffer[i, base] = 1.0
+                obs_buffer[i, base + 1] = float(cid) / 2000.0
+                obs_buffer[i, base + 2] = float(card_stats[bid, 0]) / 10.0
+                obs_buffer[i, base + 3] = float(card_stats[bid, 1]) / 5.0
+                obs_buffer[i, base + 4] = float(card_stats[bid, 2]) / 10.0
+
+    offset += 60
+
+    # Hand (up to 20 cards shown)
+    h_count = 0
+    for h_idx in range(60):
+        cid = batch_hand[i, h_idx]
+        if cid > 0 and h_count < 20:
+            base = offset + h_count * 20
+            if base + 10 < obs_dim:
+                obs_buffer[i, base] = 1.0
+                obs_buffer[i, base + 1] = float(cid) / 2000.0
+                bid = get_base_id_device(cid)
+                if bid < card_stats.shape[0]:
+                    obs_buffer[i, base + 2] = float(card_stats[bid, 0]) / 10.0
+                    obs_buffer[i, base + 3] = float(card_stats[bid, 10])  # Type
+            h_count += 1
+
+    offset += 400
+
+    # Live zone (up to 10 cards)
+    l_count = 0
+    for l_idx in range(50):
+        cid = batch_live[i, l_idx]
+        if cid > 0 and l_count < 10:
+            base = offset + l_count * 10
+            if base + 5 < obs_dim:
+                obs_buffer[i, base] = 1.0
+                obs_buffer[i, base + 1] = float(cid) / 2000.0
+            l_count += 1
+
+
+@cuda.jit
+def encode_observations_attention_kernel(
+    num_envs,
+    batch_hand,
+    batch_stage,
+    batch_energy_count,
+    batch_tapped,
+    batch_scores,
+    opp_scores,
+    opp_stage,
+    opp_tapped,
+    card_stats,
+    batch_global_ctx,
+    batch_live,
+    batch_opp_history,
+    opp_global_ctx,  # Added
+    turn_number,
+    obs_buffer,
+):
+    """
+    Encode observations for Attention Architecture (2240-dim).
+    """
+    i = cuda.grid(1)
+    if i >= num_envs:
+        return
+
+    # Constants
+    FEAT = 64
+    MAX_HAND = 15  # +1 overflow
+
+    # Offsets
+    HAND_START = 0
+    HAND_OVER_START = HAND_START + (MAX_HAND * FEAT)  # 960
+    STAGE_START = HAND_OVER_START + FEAT  # 1024
+    LIVE_START = STAGE_START + (3 * FEAT)  # 1216
+    LIVE_SUCC_START = LIVE_START + (3 * FEAT)  # 1408
+    OPP_STAGE_START = LIVE_SUCC_START + (3 * FEAT)  # 1600
+    OPP_HIST_START = OPP_STAGE_START + (3 * FEAT)  # 1792
+    GLOBAL_START = OPP_HIST_START + (6 * FEAT)  # 2176
+
+    # Clear buffer
+    for k in range(2240):
+        obs_buffer[i, k] = 0.0
+
+    # --- A. HAND (16 slots) ---
+    hand_count = 0
+    for j in range(60):
+        cid = batch_hand[i, j]
+        if cid > 0:
+            bid = get_base_id_device(cid)
+            if bid < card_stats.shape[0]:
+                if hand_count < 16:
+                    base = HAND_START + hand_count * FEAT
+
+                    obs_buffer[i, base + 0] = 1.0  # Presence
+                    obs_buffer[i, base + 1] = float(card_stats[bid, 10]) / 2.0  # Type
+                    obs_buffer[i, base + 2] = float(card_stats[bid, 0]) / 10.0  # Cost
+                    obs_buffer[i, base + 3] = float(card_stats[bid, 1]) / 5.0  # Blades
+                    obs_buffer[i, base + 5] = float(cid) / 2000.0  # Card ID (New)
+                    obs_buffer[i, base + 6] = 0.2  # Location: Hand
+
+                    # Hearts (8-14)
+                    for k in range(7):
+                        if 12 + k < card_stats.shape[1]:
+                            obs_buffer[i, base + 8 + k] = float(card_stats[bid, 12 + k]) / 5.0
+
+                    # Group (22-28)
+                    raw_group = card_stats[bid, 11]
+                    obs_buffer[i, base + 22 + (raw_group % 7)] = 1.0
+
+                    # Context
+                    obs_buffer[i, base + 58] = float(hand_count) / 10.0
+                    obs_buffer[i, base + 59] = 1.0  # Mine
+
+                    hand_count += 1
+
+    # --- B. MY STAGE (3 slots) ---
+    for slot in range(3):
+        cid = batch_stage[i, slot]
+        if cid > 0:
+            bid = get_base_id_device(cid)
+            if bid < card_stats.shape[0]:
+                base = STAGE_START + slot * FEAT
+
+                obs_buffer[i, base + 0] = 1.0
+                obs_buffer[i, base + 1] = float(card_stats[bid, 10]) / 2.0
+                obs_buffer[i, base + 2] = float(card_stats[bid, 0]) / 10.0
+                obs_buffer[i, base + 3] = float(card_stats[bid, 1]) / 5.0
+                obs_buffer[i, base + 4] = 1.0 if batch_tapped[i, slot] > 0 else 0.0
+                obs_buffer[i, base + 5] = float(cid) / 2000.0  # Card ID (New)
+                obs_buffer[i, base + 6] = 0.4  # Location: Stage
+
+                for k in range(7):
+                    if 12 + k < card_stats.shape[1]:
+                        obs_buffer[i, base + 8 + k] = float(card_stats[bid, 12 + k]) / 5.0
+
+                raw_group = card_stats[bid, 11]
+                obs_buffer[i, base + 22 + (raw_group % 7)] = 1.0
+
+                obs_buffer[i, base + 58] = float(slot) / 10.0
+                obs_buffer[i, base + 59] = 1.0
+
+    # --- C. LIVE ZONE (6 slots) ---
+    live_count = 0
+    for j in range(50):
+        cid = batch_live[i, j]
+        if cid > 0:
+            bid = get_base_id_device(cid)
+            if bid < card_stats.shape[0] and live_count < 6:
+                base = LIVE_START + live_count * FEAT
+
+                obs_buffer[i, base + 0] = 1.0
+                obs_buffer[i, base + 1] = float(card_stats[bid, 10]) / 2.0
+                obs_buffer[i, base + 2] = float(card_stats[bid, 0]) / 10.0
+                obs_buffer[i, base + 5] = float(cid) / 2000.0  # Card ID (New)
+                obs_buffer[i, base + 6] = 0.6  # Location: Live
+
+                for k in range(7):
+                    if 12 + k < card_stats.shape[1]:
+                        obs_buffer[i, base + 8 + k] = float(card_stats[bid, 12 + k]) / 5.0
+
+                obs_buffer[i, base + 58] = float(live_count) / 10.0
+                obs_buffer[i, base + 59] = 1.0
+                live_count += 1
+
+    # --- D. OPP STAGE (3 slots) ---
+    for slot in range(3):
+        cid = opp_stage[i, slot]
+        if cid > 0:
+            bid = get_base_id_device(cid)
+            if bid < card_stats.shape[0]:
+                base = OPP_STAGE_START + slot * FEAT
+
+                obs_buffer[i, base + 0] = 1.0
+                obs_buffer[i, base + 1] = float(card_stats[bid, 10]) / 2.0
+                obs_buffer[i, base + 2] = float(card_stats[bid, 0]) / 10.0
+                obs_buffer[i, base + 3] = float(card_stats[bid, 1]) / 5.0
+                obs_buffer[i, base + 4] = 1.0 if opp_tapped[i, slot] > 0 else 0.0
+                obs_buffer[i, base + 5] = float(cid) / 2000.0  # Card ID (New)
+                obs_buffer[i, base + 6] = 0.8  # Location: Opp Stage
+
+                for k in range(7):
+                    if 12 + k < card_stats.shape[1]:
+                        obs_buffer[i, base + 8 + k] = float(card_stats[bid, 12 + k]) / 5.0
+
+                obs_buffer[i, base + 58] = float(slot) / 10.0
+                obs_buffer[i, base + 59] = -1.0
+
+    # --- E. OPP HISTORY (6 slots) ---
+    for h in range(6):
+        cid = batch_opp_history[i, h]
+        if cid > 0:
+            bid = get_base_id_device(cid)
+            if bid < card_stats.shape[0]:
+                base = OPP_HIST_START + h * FEAT
+
+                obs_buffer[i, base + 0] = 1.0
+                obs_buffer[i, base + 1] = float(card_stats[bid, 10]) / 2.0
+                obs_buffer[i, base + 2] = float(card_stats[bid, 0]) / 10.0
+                obs_buffer[i, base + 5] = float(cid) / 2000.0  # Card ID (New)
+                obs_buffer[i, base + 6] = 1.0  # Location: History
+
+                obs_buffer[i, base + 58] = float(h) / 10.0
+                obs_buffer[i, base + 59] = -1.0
+
+    # --- F. GLOBAL SCALARS ---
+    obs_buffer[i, GLOBAL_START + 0] = float(batch_scores[i]) / 10.0
+    obs_buffer[i, GLOBAL_START + 1] = float(opp_scores[i]) / 10.0
+    obs_buffer[i, GLOBAL_START + 2] = float(batch_global_ctx[i, 54]) / 20.0  # Turn from Context
+    obs_buffer[i, GLOBAL_START + 3] = float(batch_global_ctx[i, 8]) / 10.0
+    obs_buffer[i, GLOBAL_START + 4] = float(batch_global_ctx[i, 5]) / 10.0
+    obs_buffer[i, GLOBAL_START + 5] = float(batch_global_ctx[i, 6]) / 40.0
+    obs_buffer[i, GLOBAL_START + 6] = float(hand_count) / 15.0
+
+    # Opponent Resources (New)
+    obs_buffer[i, GLOBAL_START + 7] = float(opp_global_ctx[i, 5]) / 10.0  # Opp Energy
+    obs_buffer[i, GLOBAL_START + 8] = float(batch_global_ctx[i, 4]) / 10.0  # Opp Hand (from ctx[4])
+    obs_buffer[i, GLOBAL_START + 9] = float(batch_global_ctx[i, 9]) / 40.0  # Opp Deck (from ctx[9])
+    obs_buffer[i, GLOBAL_START + 10] = float(batch_global_ctx[i, 7]) / 10.0  # Opp Trash (from ctx[7])