I found the exact solution. The key API is torch.gather:
import torch
def kmax_pooling(x, dim, k):
    index = x.topk(k, dim = dim)[1].sort(dim = dim)[0]
    return x.gather(dim, index)
x = torch.rand(4, 5, 6, 10)
y = kmax_pooling(x, 3, 5)
print(x[0, 0])
print(y[0, 0])

Output:
 0.2762  0.3788  0.5708  0.3251  0.0568  0.2483  0.3930  0.1249  0.1874  0.1113
 0.9230  0.7428  0.0957  0.2301  0.6187  0.8898  0.3007  0.2653  0.5313  0.1032
 0.6376  0.9639  0.6584  0.1502  0.0250  0.5792  0.9283  0.1783  0.9545  0.1681
 0.8456  0.6135  0.2860  0.9366  0.5178  0.0113  0.4864  0.9308  0.3005  0.5403
 0.3280  0.8755  0.2290  0.0899  0.9093  0.6971  0.1557  0.2412  0.7991  0.9169
 0.5389  0.4603  0.7291  0.4070  0.0113  0.3571  0.3860  0.3354  0.4081  0.0209
[torch.FloatTensor of size 6x10]


 0.2762  0.3788  0.5708  0.3251  0.3930
 0.9230  0.7428  0.6187  0.8898  0.5313
 0.6376  0.9639  0.6584  0.9283  0.9545
 0.8456  0.6135  0.9366  0.9308  0.5403
 0.8755  0.9093  0.6971  0.7991  0.9169
 0.5389  0.4603  0.7291  0.4070  0.4081
[torch.FloatTensor of size 6x5]

Can also be used on autograd.Variable.