[DSA] Support group-scaled FP8 NoPE KV reads for CUDA TileLang

[HanHan009527]

Motivation

CUDA TileLang DSA cannot currently consume group-scaled FP8 NoPE KV rows: the existing dequantization helpers require the 656-byte RoPE layout, and startup validation rejects FP8 TileLang on CUDA. This draft adds a reader and BF16 consumer path for the 528-byte layout already produced by the upstream writer: 512 FP8 values and four FP32 scales per token.

Related to #36830. This is one consumer capability toward GLM NoPE FP8 support; complete model-serving support is not established by this draft.

Modifications

• Extend the existing full and paged Triton readers to accept 528-byte NoPE rows while retaining 656-byte RoPE behavior and int64 physical addressing.
• Dequantize selected physical rows when there are fewer selected entries than physical pool rows; otherwise use the full-pool reader. Remap indices to the selected workspace, retaining invalid-entry masks. BF16 payload workspace is bounded by min(physical_tokens, selected_entries) * 1024 bytes, with no data-dependent unique operation or host synchronization.
• Connect this reader to the existing BF16 TileLang consumer. Allow CUDA FP8 through the production _check_dsa_backend_constraints entry point only for BF16 latent512/RoPE0 queries, both DSA phases using TileLang, and DCP size 1. Preserve HIP restrictions and reject mixed backend pairs and HiSparse FP8 TileLang.

This reuses upstream storage and writer code. It does not add a new cache-layout option, change backend defaults, or include raw512, native FP8 attention, H16, or HIP kernels.

Accuracy Tests

Validated source 9a576563fd880a550fcfa1fcd90dbafd72b5ae10, based on main ad5af539cd39e9af10ef318d3f8a50f5d6f9c153, in a Kubernetes Pod with one NVIDIA H20 (SM90). The Pod checked out the pushed commit and installed that checkout as an editable package. Python 3.12.3, PyTorch 2.13.0+cu130, Triton 3.7.1, TileLang 0.1.12.

python -m pytest -q \
  test/registered/unit/test_dsa_tilelang_fp8_validation.py \
  test/registered/unit/test_model_overrides.py::TestGoldenModelOverrides::test_dsa_split_backend_resolution_pass \
  test/registered/kernel/attention/test_dsa_nope_dequant.py

24 tests and 15 subtests passed. Coverage includes both layouts, scale words, duplicate indices, empty selection, invalid-mask preservation in the selected reader, startup constraints, and fixed-width valid-index reader→TileLang execution. CUDA Graph replay tests change indices, queries, and scale data. Configured pre-commit checks and git diff --check pass.

This remains a draft for review:

• Downstream attention masking for invalid indices and all-empty rows is a separate unresolved contract; preserving a reader mask does not establish consumer correctness for those inputs.
• Unfused KPool variable-width transforms are tracked by [AMD] [GLM-5.3-Flash Day 0] Support non-2048 top-k widths in the DSA page-table transform #39340, replacing [AMD] [GLM-5.3-Flash Day 0] Accept k-pool top-k widths in the DSA page-table transform #38543/[AMD] [GLM-5.3-Flash Day 0] Accept k-pool top-k widths in the DSA page-table transform #37573. Cached-prefix hybrid metadata ownership is tracked by fix(attention): unwrap hybrid backend for FP8 DSA MHA #38508. Neither is included here.
• No full-model accuracy, PD serving, or performance validation has been run for this branch. Main now includes default-on KPool metadata fusion from [GLM-5.3 Flash] Restore and enable KPool metadata fusion #38845; the reported runtime results predate that change. A static merge-tree check against main 4c71d14fbab6864a0dd7d3c77c5326ccc708badd is conflict-free, but this branch has not been rebased or runtime-validated on that base.

Speed Tests and Profiling

No standalone speedup is claimed. The bounded BF16 payload allocation above describes the reader's allocation strategy; it is not an end-to-end memory or latency measurement.

Checklist

• Format and configured pre-commit checks pass.
• Add registered CPU and CUDA tests.
• Retain the existing RoPE and HIP paths.
• Resolve remaining integration dependencies and validate full-model accuracy.
• Measure serving performance on the final integrated source.

Developed with AI assistance.

---

CI States

Latest PR Test (Base): ❌ Run #34802586908
Latest PR Test (Extra): ❌ Run #34802586726
Latest PR Test (AMD ROCm 10): ❌ Run #34802586929