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

[LLMShark]

Background

During BF16 hot updates with the flashinfer_trtllm MoE backend, SGLang restores runtime block-layout weights to canonical layout before copying updated weights. After the copy, the weights must be finalized back into the backend-specific runtime layout.

slime previously ran the post-load post_process_weights(...) step only for compressed-tensors. Unquantized/BF16 updates skipped it, leaving FlashInfer TRT-LLM MoE weights in canonical layout for the next forward pass.

Reproduction and before/after

Environment:

• 8x NVIDIA B200 (SM100), colocated training
• SGLang v0.5.10.post1
• BF16 Qwen3.5-35B-A3B and Qwen3.6-35B-A3B
• flashinfer_trtllm MoE backend

With the corresponding SGLang pre-load canonical-shape restore fix applied, the BF16 hot-update copy completed successfully. Without this PR, the next generation forward pass failed with:

IndexError: Index 3 out of bounds for tensor with 3 dimensions
at flashinfer::Bf16MoeLauncher::check_moe()

The updated MoE weights remained in canonical layout, while the FlashInfer TRT-LLM kernel expected its runtime block layout.

After applying this PR together with the SGLang restore fix:

• Qwen3.5 completed all 132/132 update chunks.
• Qwen3.6 completed all 136/136 update chunks.
• Subsequent generation no longer reproduced the FlashInfer index error.
• Multiple rollout, ref_log_probs, and actor_train iterations completed.

This comparison validates the restore-copy-finalize control flow and crash fix. It does not by itself establish logit or log-probability parity between the training and rollout engines.

Changes

Updated both hot-update implementations:

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

Added a shared requires_post_process_after_update(...) predicate with an explicit scope:

quantization_config	Post-load processing	Reason
None	Yes	Finalize unquantized/BF16 backend-specific runtime layouts
compressed-tensors	Yes	Preserve the existing restore-then-repack lifecycle
FP8, empty, or unknown configs	No	Preserve previous behavior for unrelated quantization backends

The predicate deliberately checks quantization_config is None instead of treating every falsy value as an unquantized model. It also avoids assuming that every quantization backend's process_weights_after_loading hook is a safe or idempotent no-op.

Additional behavior preserved:

• The compressed-tensors pre-load restore call remains conditional.
• continue_generation runs only after required post-processing succeeds.
• Generation remains paused if required post-processing raises an exception.

Relationship to SGLang #28015

This PR and SGLang #28015 address different phases of the update lifecycle:

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

For the non-routed backend, both changes, or equivalent downstream patches, are required:

1. Restore runtime block-shaped parameters to canonical shape.
2. Copy the updated canonical weights.
3. Convert the updated weights back to the runtime block layout.

Tests

tests/test_bf16_weight_post_process.py covers:

• The post-processing scope for None, compressed-tensors, FP8, empty, and unknown configurations.
• BF16 tensor and distributed updates run post-processing.
• compressed-tensors preserves restore-then-post-process ordering.
• Unrelated quantized tensor and distributed updates skip post-processing and still resume generation.
• Generation does not resume when required post-processing fails.

The test is included in the CPU CI matrix. The existing empty colocated bucket test also provides the shared helper in its isolated fake dependency module.

Validation

• git diff --check passed.
• Python 3.10 syntax checks passed.
• tests/test_empty_colocated_weight_bucket.py: 2 passed locally.
• GitHub pre-commit passed on fb35d5e.
• GitHub PR Test passed on fb35d5e.
• End-to-end control-flow validation completed on an 8x B200 colocated setup with the corresponding SGLang patch.

[LLMShark]

Hi maintainers, just a gentle ping on this PR. It fixes the missing post-load finalization step for BF16 weight hot updates, and the CI checks are passing. I would appreciate a review when you have time. Thanks!

[LLMShark]

@huang3eng Please help review~