Description:
Context:
Standard autoregressive sampling like greedy or temperature sampling in vLLM breaks the computational graph, requiring REINFORCE-style gradient estimators like PPO. For future alignment algorithms that may rely on reparameterization gradients through discrete tokens, we need a high-performance, differentiable Gumbel-Softmax sampling operator.
Tasks:
- Implement a Triton kernel for the Gumbel-Softmax trick in rl_engine/kernels/ops/triton/sampling.py.
- The kernel should support hard=True returns one-hot vectors on forward, continuous gradients on backward and temperature annealing.
- Ensure it efficiently handles the [Batch, Sequence, Vocab_size] tensor shape without materializing excessive intermediate random noise arrays in global memory.
Description:
Context:
Standard autoregressive sampling like greedy or temperature sampling in vLLM breaks the computational graph, requiring REINFORCE-style gradient estimators like PPO. For future alignment algorithms that may rely on reparameterization gradients through discrete tokens, we need a high-performance, differentiable Gumbel-Softmax sampling operator.
Tasks: