Chronos PoC: PTX Transform Selection via RLVR
Proof-of-concept RL agent that selects PTX assembly transforms to optimize GPU kernel performance on NVIDIA L4 (sm_89).
What This Is
An MLP policy trained with DA-GRPO (Demonstration-Anchored Group Relative Policy Optimization) to select sequences of PTX-level transforms that reduce GPU kernel execution cycles. Trained on 64 gemm_tile kernel variants, validated on 33 diverse Triton kernels.
Results
Metric Value Mean cycle reduction (gemm_tile) -29.2% Best single kernel -53.8% (gemm_tile 4,6,8: 1839 -> 849 cycles)Generalization to Triton kernels 17/33 kernels improve Best Triton improvement -47.0% (attention_d64_kv64 with maxnreg_255)Training time ~6 hours on single NVIDIA L4 Model parameters ~20K
Model Architecture
Input: 25 kernel features + 21 action mask + 21 action history = 67 dims
Hidden: Linear(67, 128) -> ReLU -> Dropout(0.1) -> Linear(128, 128) -> ReLU -> Dropout(0.1)
Output: Linear(128, 21) -> mask -> softmax25 features : instruction counts (loads, stores, FMA, branches), vectorization ratios, cache hint coverage, register state, instruction mix ratios.
21 actions : 20 PTX transforms + stop. Transforms include vectorize loads/stores, cache hints (cs/cg/ca/cv), store cache hints (cs/wt/wb), register budget limits (32/64/128/255), instruction reorder (critical_path/interleave/loads_first/stores_last), prefetch, and split vector loads.
Files
checkpoint_best.pt # Best checkpoint (epoch 250, -29.2% mean)
checkpoint_latest.pt # Final checkpoint (epoch 500)
inference.py # Self-contained inference script
training_result.json # Per-kernel results (64 gemm_tile kernels)
bc_stats.json # Behavior cloning warm-start statistics
measure_triton_results.json # Triton kernel measurement results (33 kernels)Inference
Requirements
No GPU required for inference (model runs on CPU). GPU only needed for actually applying transforms and measuring cycles.
Quick Start
1 import torch
2 from inference import load_model , predict_transforms
3
4 # Load model
5 model = load_model ( "checkpoint_best.pt" )
6
7 # Predict transforms for a PTX kernel
8 with open ( "your_kernel.ptx" ) as f :
9 ptx = f . read ( )
10 actions = predict_transforms ( model , ptx )
11 # -> ['maxnreg_128', 'vec_ld', 'vec_st']
Command Line
1 # Demo with synthetic features
2 python inference.py --checkpoint checkpoint_best.pt
3
4 # Run on a PTX file
5 python inference.py --checkpoint checkpoint_best.pt --ptx path/to/kernel.ptx
Programmatic Usage
1 import torch
2 from inference import TransformPolicy , extract_features_from_ptx , get_action_mask , get_action_history , ACTION_NAMES
3
4 # Load
5 model = TransformPolicy ( hidden = 128 )
6 ckpt = torch . load ( "checkpoint_best.pt" , map_location = "cpu" , weights_only = False )
7 model . load_state_dict ( ckpt [ "policy" ] )
8 model . eval ( )
9
10 # Extract features from PTX
11 ptx_source = open ( "kernel.ptx" ) . read ( )
12 features = extract_features_from_ptx ( ptx_source )
13
14 # Predict step by step
15 applied = set ( )
16 for step in range ( 6 ) :
17 feat_t = torch . tensor ( features , dtype = torch . float32 )
18 mask_t = torch . tensor ( get_action_mask ( applied ) , dtype = torch . float32 )
19 hist_t = torch . tensor ( get_action_history ( applied ) , dtype = torch . float32 )
20
21 action_id = model . get_greedy_action ( feat_t , mask_t , hist_t )
22 action = ACTION_NAMES [ action_id ]
23 if action == "stop" :
24 break
25 print ( f"Step { step + 1 } : apply { action } " )
26 applied . add ( action )
Training Details
Algorithm: DA-GRPO
BC warm-start (50 epochs): Clone greedy search trajectories. Best accuracy: 64.5%.GRPO training (450 epochs): Hardware-in-the-loop RL with SM clock() cycle measurement.
Group size: 8 rollouts per kernel (1 anchor from reference policy + 7 with forced diverse first actions)
Advantage: MC-GRPO (median baseline per kernel, global z-normalization)
Reward: log(cycles_before / cycles_after) — outcome-only, terminal
KL penalty: beta=0.01 against BC reference policy
Clipped surrogate: epsilon=0.2
Action Space
20 PTX transforms organized into 5 conflict groups (only one per group):
Cache hints (load): cs, cg, ca, cvStore cache hints : cs, wt, wbRegister budget : maxnreg 32, 64, 128, 255Instruction reorder : critical_path, interleave, loads_first, stores_lastPrefetch : L1, L2Vectorize : loads, stores (independent)Split : vector loads (independent)
Hardware
NVIDIA L4 GPU (sm_89, Ada Lovelace)
SM clock() cycle counter (1-cycle std dev, 200 samples per measurement)
pip-installed CUDA 12.9 ptxas
Limitations
Trained on gemm_tile kernels only (64 variants, m,n,k in {2,4,6,8})
Mode collapse: 52/64 kernels get the same sequence (vec_st + vec_ld + maxnreg_128)
MLP can't read PTX code — relies on 25 scalar features
3% measurement error rate from cudaErrorMisalignedAddress on gemm_tile(4,6,4)
Reorder transforms deadlock on kernels with bar.sync barriers
References
CuAsmRL (CGO 2025) : PPO on SASS schedulingDr. Kernel (2026) : REINFORCE for Triton kernelsDeepSeek-R1 (2025) : GRPO algorithmMC-GRPO (2025) : Median-centered baseline
License
Research prototype. Contact for usage terms.
