[FlyDSL] Keep FP4 prefill modules alive across async dispatches

[AMD-yanfeiwang]

Summary

• retain compiled paged FP4 prefill launchers for the process lifetime
• prevent LRU eviction from unloading modules that asynchronous GPU dispatches can still reference
• add a cache-policy regression check to the existing direct-script op test

Root cause

The paged FP4 prefill compiler specializes on max_blocks_per_seq and used a 32-entry LRU. A 128K chunked-prefill workload creates 125 guarded page-table widths. Eviction drops the FlyDSL launcher; FlyDSL 0.3.1 then runs GpuJitModule.__del__, which calls mgpuModuleUnload/hipModuleUnload without first synchronizing outstanding launch streams.

In SGLang PR #36581, this produced an eight-GPU memory fault followed by HSA_STATUS_ERROR_ILLEGAL_INSTRUCTION. ROCm Debug Agent captured 2048 stopped waves and an AQL packet with grid=131072, workgroup=256, and a kernel_object absent from the live kernel map. That geometry exactly matches this kernel's 512-CTA x 256-thread prefill launch. Disabling scheduler overlap did not change the signature.

An unbounded cache is intentionally used instead of raising the limit: any finite bound can reintroduce unsafe module unload as compile dimensions vary. For the current 128K/page-size-256 integration, width specialization is practically bounded to about 128 entries. Long term, FlyDSL can make module retirement stream-aware.

Validation

• focused cache-policy test: 1 passed
• complete test_flydsl_pa_mqa_logits_fp4_prefill.py: PASS; 4 prefill and 5 varqlen cases, all exact-reference cosine 1.0
• exact production A/B on 8x MI355X, ROCm 7.2.4, TP8/DP8 attention, EAGLE, radix cache, 128K, global concurrency 16: warmup and measured wave passed; crossed the prior 126720 + 1024 + 256 fault boundary; zero HSA/illegal-instruction fault markers; all scheduler workers remained healthy
• Black 26.3.0, Ruff 0.15.7, py_compile, and git diff --check pass

Related: sgl-project/sglang#36581

[github-actions]

🏷️ CI Guide

Runs automatically on every PR:

• ✅ Pre-checks (submodule verification, code formatting)
• ✅ Aiter op tests (gfx942 + gfx950)
• ✅ Triton tests on MI35X (only when aiter/ops/triton/** or related paths are changed)

Extended tests (opt-in via labels):

Label	Tests
ci:gfx1250-ffm-triton	Run the five-shard gfx1250 FFM Triton test suite
ci:triton-300x	Run an additional Triton test job on MI300X in PRs; main branch always runs both MI35X and MI300X
multigpu	Aiter multi-GPU tests on the 8-GPU runner
ci:sglang	SGLang integration tests: DeepSeek-R1-MXFP4 accuracy, Qwen 3.5 accuracy
ci:atom	ATOM benchmark: DeepSeek-R1-0528, GPT-OSS-120B
ci:atom_full	ATOM accuracy suite for PR and main models from ATOM models_accuracy.json
ci:vllm	vLLM benchmark: GPT-OSS-120B, DeepSeek-R1-0528, Kimi-K2.5
ci:all	All standard extended tests (excludes ci:atom_full)

Only add ci:atom_full for FlyDSL or Triton upgrades.
Add labels via the sidebar or gh pr edit 5126 --add-label <label>

PR title tags:
Component tags ([Triton/Gluon], [HIP], [CK], [ASM], ...) are added to the PR title automatically from the changed files and re-synced on every push — change-type tags like [fix]/[Perf] and op tags like [MLA] are left untouched. Add the no-auto-title label to opt this PR out of title tagging.

[Bernard-Liu]

This prevents unsafe eviction, but permanently retains every compiled launcher and GPU module. Since max_blocks_per_seq varies with the page-table width and is only used as the block-table row stride, could we pass the stride at runtime instead of creating an unbounded number of specializations?

[github-actions]

🏷️ CI Guide

Runs automatically on every PR:

• ✅ Pre-checks (submodule verification, code formatting)
• ✅ Aiter op tests (gfx942 + gfx950)
• ✅ Triton tests on MI35X (only when aiter/ops/triton/** or related paths are changed)

Extended tests (opt-in via labels):

Label	Tests
ci:gfx1250-ffm-triton	Run the five-shard gfx1250 FFM Triton test suite
ci:triton-300x	Run an additional Triton test job on MI300X in PRs; main branch always runs both MI35X and MI300X
multigpu	Aiter multi-GPU tests on the 8-GPU runner
ci:sglang	SGLang integration tests: DeepSeek-R1-MXFP4 accuracy, Qwen 3.5 accuracy
ci:atom	ATOM benchmark: DeepSeek-R1-0528, GPT-OSS-120B
ci:atom_full	ATOM accuracy suite for PR and main models from ATOM models_accuracy.json
ci:vllm	vLLM benchmark: GPT-OSS-120B, DeepSeek-R1-0528, Kimi-K2.5
ci:all	All standard extended tests (excludes ci:atom_full)

Only add ci:atom_full for FlyDSL or Triton upgrades.
Add labels via the sidebar or gh pr edit 5126 --add-label <label>

PR title tags & labels:
Component tags ([Triton/Gluon], [HIP], [CK], [ASM], ...) are added to the PR title and as PR labels automatically from the changed files and re-synced on every push — change-type tags like [fix]/[Perf], op tags like [MLA], and human labels (ci:*) are left untouched. Add the no-auto-title label to opt this PR out.

[github-actions]

🏷️ CI Guide

Runs automatically on every PR:

• ✅ Pre-checks (submodule verification, code formatting)
• ✅ Aiter op tests (gfx942 + gfx950)
• ✅ Triton tests on MI35X (only when aiter/ops/triton/** or related paths are changed)

Extended tests (opt-in via labels):

Label	Tests
ci:gfx1250-ffm-triton	Run the five-shard gfx1250 FFM Triton test suite
ci:triton-300x	Run an additional Triton test job on MI300X in PRs; main branch always runs both MI35X and MI300X
multigpu	Aiter multi-GPU tests on the 8-GPU runner
ci:sglang	SGLang integration tests: DeepSeek-R1-MXFP4 accuracy, Qwen 3.5 accuracy
ci:atom	ATOM benchmark: DeepSeek-R1-0528, GPT-OSS-120B
ci:atom_full	ATOM accuracy suite for PR and main models from ATOM models_accuracy.json
ci:vllm	vLLM benchmark: GPT-OSS-120B, DeepSeek-R1-0528, Kimi-K2.5
ci:all	All standard extended tests (excludes ci:atom_full)

Only add ci:atom_full for FlyDSL or Triton upgrades.
Add labels via the sidebar or gh pr edit 5126 --add-label <label>

PR title tags & labels:
Component tags ([Triton/Gluon], [HIP], [CK], [ASM], ...) are added to the PR title and as PR labels automatically from the changed files and re-synced on every push — change-type tags like [fix]/[Perf], op tags like [MLA], and human labels (ci:*) are left untouched. Add the no-auto-title label to opt this PR out.