Unify expert parallel token dispatcher API

[wwwjn]

Summary

This PR unifies the token-dispatch API across EP backends and separates:

1. The maximum token capacity used to allocate persistent communication storage.
2. The current token capacity used by one dispatch operation.

Persistent communication buffers are initialized before execution and no longer grow inside
dispatch().

API design

Dispatcher configuration carries the lifetime maximum:

num_max_tokens_per_rank: int

Dispatch receives the current logical per-rank capacity:

  dispatch(
      x_TD,
      topk_scores_TK,
      topk_expert_ids_TK,
      num_local_tokens_per_expert_E,
      *,
      num_tokens_per_rank: int,
  )

The intended relationship is:

  x_TD.shape[0]
      <= num_tokens_per_rank
      <= num_max_tokens_per_rank

x_TD.shape[0] is the materialized token count. The dispatch argument includes sequence-dimension
padding and may therefore be larger.

During mesh wiring, backends with persistent communication storage allocate it using
num_max_tokens_per_rank. Dispatch then creates or uses a current logical layout without
reallocating or growing that storage.

This provides stable buffer addresses for CUDA graph capture while allowing different eager calls
or graph buckets to use different current token capacities.

Backend behavior

• DeepEP: preallocates one ElasticBuffer. The current capacity is passed as DeepEP's per-rank
  layout stride.

• HybridEP: preallocates its communication buffer. Non-blocking dispatch uses the current capacity
  to size the fused-permute output.

• MinimalAsyncEP: preallocates symmetric-memory storage. It currently derives dispatch shapes from
  x_TD.shape[0] because TP/SP is unsupported.

• AllToAll/TorchAO: have no dispatcher-owned persistent communication buffer and continue using
  exact collective tensor sizes.

Capacity derivation

Pretraining

The maximum is derived from the configured training shape:

  num_token_shards = (
      context_parallel_degree * tensor_parallel_degree
  )

  num_max_tokens_per_rank = (
      local_batch_size
      * ceil_div(seq_len, num_token_shards)
  )

Ceiling division prevents underallocation when the sequence length is not divisible by the token-
sharding degree.

vLLM inference

The inference bound comes from vLLM's per-step scheduler limit:

  num_max_tokens_per_rank = ceil_div(
      max_num_batched_tokens,
      context_parallel_degree * tensor_parallel_degree,
  )

vLLM filters its finalized CUDA graph capture sizes so they do not exceed max_num_batched_tokens.
Decode-only workers should configure max_num_batched_tokens from their maximum decode workload
rather than retaining a large prefill-oriented default.

If the user explicitly configures num_max_tokens_per_rank, setup validates that it is positive and
large enough for the derived runtime requirement.

Correctness

The change only refactors dispatcher APIs and communication-buffer lifetime. It does not change MoE
routing or expert computation.

With TP enabled, the baseline and this PR were run for 10 training steps using:

  --debug.seed=42
  --debug.deterministic

The runs produced bitwise-identical loss and grad_norm at every compared step.

[tianyu-l]

Let's try to deprecate this, maybe blocked by deprecate non-spmd_types (default / full_dtensor) backend. If so, can wait for @pianpwk

[tianyu-l]

@IvanKobzarev to review code change in this file

[tianyu-l]

@IvanKobzarev if there's no review, we'll proceed with the change

[wwwjn]

Including all fields when constructing a new HybridEPTokenDispatcher.Config, copy all the fields over

[tianyu-l]

@IvanKobzarev if there's no review, we'll proceed with the change

[tianyu-l]

when would it be None?

[wwwjn]

No, both pretraining and inference caller will send training config. This checking is for safty guard purpose, and following the same pattern in maybe_update_minimal_async_ep_config