[Zhichenzzz]

Fixes needed for colocated RL on hybrid-state models (Qwen3.8-Flash-Next, GLM-5.3-Flash) on current sglang-miles. Supersedes #40815, which is folded in here as the first commit.

Motivation

1. Hybrid state pools fault with CUDA error 700 at memory release. release_memory_occupation paused (unmapped) the KV cache memory-saver region and only then called flush_cache(). The hybrid pools cleared there (mamba states, qwen4_exp PLE ShortConvPool / NgramPool, KDA state) live inside that region and their clear() writes device memory, so the write hits unmapped memory. With CUDA_LAUNCH_BLOCKING=1 the fault lands on ShortConvPool.clear() via release_memory_occupation -> flush_cache -> HybridReqToTokenPool.clear. The trainer's later /flush_cache while the engine is offloaded faults the same way.
2. bf16 MoE weight update crashes on plain flashinfer_trtllm. The TRT-LLM bf16 postprocess repacks w13/w2 into block layout for every TRT-LLM MoE runner variant, but maybe_restore_flashinfer_trtllm_bf16_weight_shape_for_load only restored the canonical shape for flashinfer_trtllm_routed. On GB300 the auto backend resolves to plain flashinfer_trtllm, so the first weight sync failed with size of tensor a (64) must match tensor b (4096) at dimension 2.

Modifications

1. weight_updater.py: flush before tagging the KV region as offloaded and pausing it (the tag is what makes the later flush a no-op). Scheduler.flush_cache: return success without clearing while the KV region is released (the pools were flushed right before the pause).
2. unquant.py: gate the shape restore on self.use_flashinfer_trtllm_moe, the same flag that gates the repack.
3. qsa/mqa.py: the prefill indexer walks its key tiles with a serial loop. tilelang 0.1.14 only pipelines constant trip counts and rejected T.ceildiv(end_max - start_min, block_n) at codegen (Downcast from tirx.Sub to ir.IntImm failed), which broke every Qwen4Exp prefill on images built with that release (the miles CI image). Logits are bitwise identical on 0.1.9; an 8k x 8k prefill costs about 10% more per QSA layer on GB300.

Validation

Full-model colocated RL e2e, 8 nodes x 4 GB300, latest radixark/miles:dev, --colocate, R3 routing replay, CUDA graph on:

• Qwen3.8-Flash-Next RL training support radixark/miles#2777 (Qwen3.8-Flash-Next): without fix 1, all four engines die with CUDA error 700 at the first memory release. With it: 5 rollouts / 5 train steps, weight-update equality check passes, train_rollout_logprob_abs_diff 0.0083 to 0.0131, raw reward 0.47 to 0.75.
• GLM 5.3 Flash RL Support radixark/miles#2786 (GLM-5.3-Flash bf16): 4 rollouts / 3 train steps, train_rollout_logprob_abs_diff 0.0085 / 0.0098 / 0.0087, raw reward 0.50 to 0.72; routed experts were captured with the triton runner, which the miles launcher now pins for routing replay (the resolver guard was dropped from this PR).

---

CI States

Latest PR Test (Base): ❌ Run #36084620812
Latest PR Test (Extra): ❌ Run #36084620581
Latest PR Test (AMD ROCm 10): ❌ Run #36084620803

[yueming-yuan]

[P1] Flush before marking KV memory as offloaded — scheduler.py:5002
release_memory_occupation() adds kv_cache to offload_tags before invoking self.flush_cache(). This new guard therefore skips the pre-pause flush on every KV release. Resume performs no reset either, leaving radix-cache entries and allocator state referencing discarded KV contents. A CPU harness executing the actual methods confirms one cache reset on the base, zero on this PR. Set the offload tag after flushing.

[yueming-yuan]

random q: do you think it will be better to directly assert moe_backend not in bypass_topk, and use assert message to ask the user to set --backend auto? (in case the user doesn't want to silently change the backend)