Skip to content

[AMD] [GLM-5.3-Flash Day 0] Enable zero-RoPE TileLang DSA on gfx950 - #37653

Closed
Raiden-Makoto wants to merge 1 commit into
sgl-project:xinyuan/glm-5.3-flash-supportfrom
Raiden-Makoto:RM/glm53-rocm-tilelang-zero-rope
Closed

Raiden-Makoto wants to merge 1 commit into
sgl-project:xinyuan/glm-5.3-flash-supportfrom
Raiden-Makoto:RM/glm53-rocm-tilelang-zero-rope

Conversation

@Raiden-Makoto

@Raiden-Makoto Raiden-Makoto commented Sep 2, 2026

Copy link
Copy Markdown
Contributor

Summary

GLM-5.3-Flash sparse DSA uses a 256-wide NoPE latent with a zero-width RoPE tail. The HIP TileLang BF16 partial kernel still emits zero-extent tail allocation/copy/gather/GEMM work, while the HIP FP8 kernel and scaled cache helpers are specialized for the existing DeepSeek-style 512+64 layout.

This PR enables GLM's 256+0 geometry without changing the established 512+64 behavior.

Contract Before With this PR
HIP BF16 tail zero-width tail buffers and operations are emitted tail allocation, copy, gather, and GEMM are compiled out
HIP FP8 main width fixed to four 128-wide tiles (d_v=512) supports two tiles for d_v=256 and retains four for d_v=512
HIP FP8 RoPE tail fixed 64-wide tail supports zero or 64 elements and compiles out zero-tail work
scaled FP8 cache row fixed 656 bytes (512 FP8 + 16 scale + 128 RoPE) also supports 264 bytes (256 FP8 + 8 scale + 0 RoPE)
paged dequant/gather fixed 512+64 offsets and launch geometry derives dimensions, scale count, offsets, and output width from the packed row

Scope

Four production files:

  • tilelang_kernel.py: compile out HIP BF16 zero-tail work and specialize the HIP FP8 kernel for either 256+0 or 512+64.
  • quant_k_cache.py: infer or accept the NoPE width, derive the number of 128-element scale tiles, and permit an empty RoPE tensor.
  • dequant_k_cache.py: recognize 264-byte and 656-byte packed rows and use their derived dimensions in direct, paged, and gather/dequant/requant paths.
  • memory_pool.py: document the dimension-derived scaled layout and optional RoPE payload at the existing write sites.

One MI35x test file directly covers the TileLang kernels and scaled cache helpers. No model loading, server launch, accuracy benchmark, scheduler/backend policy, dense MHA path, mHC path, MoE path, or CUDA TileLang FP8 behavior is changed.

The FP8 kernel keeps the existing 512+64 allocation, GEMM, accumulation, normalization, and store order. The 256+0 specialization emits only the first two 128-wide tiles and no tail operations.

Test plan

Docker: rocm/sgl-dev:v0.5.18-rocm720-mi35x-20260830.
Hardware: one MI350X (gfx950).
Base: xinyuan/glm-5.3-flash-support @ 545bd6f839.
The same new test file was run in both arms; the four production files are the only variables.

Arm Result
Baseline FAILED (6 failures, 8 errors). All six BF16 256+0 cases hit the generated zero-tail kernel failure; the 256+0 scaled-cache case is rejected by the fixed layout; FP8 256+0 does not match the Torch oracle.
This PR BF16 256+0 matrix: passed; FP8 256+0 matrix: passed; 256+0 and 512+64 cache round-trip/paged-gather checks: passed; BF16/FP8 512+64 controls: passed.

Direct GLM geometry coverage:

  • Q [tokens, 64, 256], KV [slots, 1, 256], output BF16 [tokens, 64, 256];
  • token counts 1, 8, and 17;
  • top-k widths 2048 and 2112;
  • fully valid rows and 2051 logical entries padded with -1 to 2112;
  • explicit Torch gather/mask/score/softmax/value oracle;
  • BF16 and native gfx950 FP8 inputs;
  • finite shape/dtype checks and bit-identical repeated execution;
  • 264-byte (256+0) and 656-byte (512+64) scaled-cache round trips and repeated-page gathers;
  • BF16 and FP8 512+64 sparse-attention regression controls.
Check Detail
Registration register_amd_ci(suite="stage-b-test-1-gpu-small-amd-mi35x")
Formatting File-scoped syntax, isort, Ruff, Black, codespell, and registered-test validation pass. The whole-tree package-location hook separately reports a pre-existing file outside this PR: python/sglang/jit_kernel/tests/test_triton_store_cache_local.py.
Accuracy not measured — this is component-level Day-0 correctness; full-model serving is intentionally outside scope.
Speed not measured — no server or performance benchmark was run.

CI States

Latest PR Test (Base): 🚫 Run #33904283151
Latest PR Test (Extra): ❌ Run #33904282808
Latest PR Test (AMD ROCm 7.2): ❌ Run #33904283075

@Raiden-Makoto Raiden-Makoto changed the title [AMD] Enable GLM-5.3-Flash zero-RoPE TileLang DSA [AMD] [GLM-5.3-Flash Day 0] Enable zero-RoPE TileLang DSA on gfx950 Sep 2, 2026
Compile out zero-width tail work and generalize FP8 TileLang/cache layouts so GLM's 256+0 geometry is covered without changing the existing 512+64 path.
@Raiden-Makoto
Raiden-Makoto force-pushed the RM/glm53-rocm-tilelang-zero-rope branch from 98f81b4 to a437ca4 Compare September 4, 2026 18:08
@Fridge003
Fridge003 deleted the branch sgl-project:xinyuan/glm-5.3-flash-support September 6, 2026 09:28
@Fridge003 Fridge003 closed this Sep 6, 2026
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Projects

None yet

Development

Successfully merging this pull request may close these issues.

3 participants