fix(update_weight): restore FlashInfer MoE layout after BF16 hot updates

[LLMShark]

Background

When performing BF16 weight hot-updates with the flashinfer_trtllm MoE backend, such as Qwen3.5 / Qwen3.6 35B-A3B, SGLang restores block-layout weights to canonical layout before copying the updated weights.

However, slime previously called:

post_process_weights(
    restore_weights_before_load=False,
    post_process_quantization=True,
)

only for the compressed-tensors quantization path. The BF16 path skipped this step, leaving MoE weights in canonical layout after the update while the FlashInfer TRT-LLM forward path expected block layout. The next forward pass could therefore fail with:

IndexError: Index 3 out of bounds for tensor with 3 dimensions

at flashinfer::Bf16MoeLauncher::check_moe().

Changes

Updated both weight-update implementations:

• slime/backends/megatron_utils/update_weight/update_weight_from_tensor.py
• slime/backends/megatron_utils/update_weight/update_weight_from_distributed.py

The post-load post_process_weights(...) call now runs after every successful hot-update, rather than only for compressed-tensors.

Additional behavior preserved:

• The existing compressed-tensors pre-load call with restore_weights_before_load=True remains conditional.
• continue_generation is called only after post-processing succeeds.
• Generation remains paused if post-processing raises an exception.
• Backend-specific post-processing is delegated to SGLang; backends without additional processing are expected to no-op.

Relationship to SGLang #28015

This PR and SGLang #28015 are independently mergeable and address
different phases of the weight-update lifecycle:

• SGLang #28015 fixes the pre-load restore step for the non-routed
  flashinfer_trtllm BF16 MoE backend.
• This PR triggers the post-load finalize step after the updated
  canonical weights have been copied.

For the non-routed flashinfer_trtllm backend, both changes (or
equivalent downstream patches) are required for an end-to-end hot
update:

1. restore the runtime block-shaped parameter to canonical shape;
2. copy the updated canonical weights;
3. convert the weights back to the runtime block layout.

For routed backends whose pre-load restore path already works, this PR
can independently fix the missing post-load finalize step.

Tests

Added tests/test_bf16_weight_post_process.py, covering:

• The BF16 tensor update path calls post-processing with:
  • restore_weights_before_load=False
  • post_process_quantization=True
• The BF16 distributed update path calls post-processing.
• The compressed-tensors path preserves the restore-then-post-process order.
• continue_generation is not called when post-processing fails.

The test is included in the CPU CI matrix.

Validation

• git diff --check passed.
• Python syntax checks passed.
• CPU unit tests passed locally: 4 passed.
• GitHub PR Test passed, including
  test_bf16_weight_post_process.py.
• End-to-end control-flow validation was performed on an 8x B200
  colocated setup with the corresponding SGLang patch:
  • Qwen3.5-35B-A3B and Qwen3.6-35B-A3B completed full hot updates
    (132/132 and 136/136 update chunks).
  • Multiple iterations completed through rollout, ref_log_probs,
    and actor_train.
  • The previous shape mismatch and FlashInfer index error were not
    reproduced.

This validation confirms the update/finalize control flow and kernel
execution. It does not by itself establish logit or log-probability
parity between the training and rollout engines. Numerical parity
should be validated separately with a fixed-token, no-op hot-update
test.