Remember some constructor args in AllToAll to make life more convenient for users. (#2)

Being able to query some of the setup parameters from the AllToAll class
would make client code a bit simpler/safer, e.g. see
pplx_dispatch_combine.py from
vllm-project/vllm#15956

cc @abcdabcd987 , @tlrmchlsmth

Signed-off-by: Bill Nell <[email protected]>

Author: bnellnm

src/pplx_kernels/all_to_all.py

@@ -21,6 +21,9 @@ def __init__(
         assert world_size % dp_size == 0
         assert world_size // dp_size > 1
 
+        self.world_size = world_size
+        self.dp_size = dp_size
+        self.max_num_tokens = max_num_tokens
         self._has_scales = hidden_dim_scale_bytes > 0
 
         self._ptr = _ops.all_to_all_create(