Rocm changes to tests/e2e/megatron/test_glm5_744b_a40b_4layer.py

[sreerohi]

• Drop DeepEP; use alltoall dispatcher
• Swap flashmla_sparse NSA backend for tilelang
• Set sglang-page-size to 1 (required by tilelang NSA)
  Depends on [glm5] Clamp indexer topk to seq_len for short sequences #1122 and [glm5] Tile sparse-MLA fwd/bwd kernels for gfx950 (MI300/MI355) LDS b… #1123

Supersedes #1172

[gemini-code-assist]

Code Review

This pull request introduces ROCm compatibility to the GLM5 training script by dynamically detecting ROCm environments and adjusting DeepEP settings, NSA backends, KV cache types, page sizes, and environment variables accordingly. The feedback recommends refactoring the construction of sglang_args to avoid mixing implicit string concatenation with explicit + operators, which will improve code readability.

Important

The consumer version of Gemini Code Assist on GitHub is being sunset. Starting June 18, 2026, new organization installations will be blocked, and all code review activity will officially cease on July 17, 2026.
For more details on the timeline and next steps, please review the Help Documentation.

[gemini-code-assist]

Mixing implicit string literal concatenation and explicit + operators inside a parenthesized expression is highly unreadable and error-prone. It is much cleaner and more maintainable to split the string construction or use separate additions.

    sglang_args += (
        f"--sglang-nsa-decode-backend {nsa_backend} "
        f"--sglang-nsa-prefill-backend {nsa_backend} "
    )
    if not IS_ROCM:
        sglang_args += "--sglang-kv-cache-dtype bf16 "
    sglang_args += (
        "--sglang-attention-backend nsa "
        f"--sglang-page-size {1 if IS_ROCM else 64} "
        f"--sglang-cuda-graph-max-bs {sglang_decode_max_bs} "
        # concurrency
        f"--sglang-max-running-requests 512 "
    )