fix(hip): query the device LDS limit instead of assuming 64 KB

[demandal25]

Summary

Single-decode rejected any configuration needing more than 65536 bytes of shared memory, with the message "exceeds CDNA3 limit of 64KB". That number is correct for CDNA3 and wrong for CDNA4 — measured on an MI350X:

sharedMemPerBlock                = 163840  (160 KB)
maxSharedMemoryPerMultiProcessor = 163840  (160 KB)

So the hard-coded guard rejects decode configurations gfx950 can actually run, at 2.5x the assumed ceiling.

What changed

• include/flashinfer/attention/generic/decode.cuh — the guard compares against getMaxSharedMemPerBlock(dev_id) instead of the literal 65536U, and the error message quotes the queried limit rather than a fixed "64KB", so a future failure names the real ceiling of the device it happened on.
• include/gpu_iface/gpu_runtime_compat.hpp — memoize getMaxSharedMemPerBlock per device id, so moving the guard onto a queried limit does not add a device query to every decode launch.

Architecture / design notes

The helper already exists in include/gpu_iface/gpu_runtime_compat.hpp and is already exercised by three call sites in prefill.cuh, so this adds no new machinery — decode simply was not using it.

gpuGetDevice had to be hoisted. The existing call sits inside the partition-kv else branch, reached only when seq_len > 256 && tmp != nullptr, while the check runs before that. It is now done once above the check and reused below.

Why the helper is now cached

Querying the limit put getMaxSharedMemPerBlock on every SingleDecodeWithKVCacheDispatched call, including the seq_len <= 256 || tmp == nullptr fast path that previously issued no device query at all — and the helper copies out a whole hipDeviceProp_t rather than reading a single attribute. Measured on MI350X / ROCm 7.2:

call	cost
hipGetDeviceProperties, first call	6.359 us
hipGetDeviceProperties, steady state	0.168 us
hipFuncSetAttribute (already on this path)	0.083 us
hipLaunchKernel (already on this path)	1.015 us
getMaxSharedMemPerBlock, memoized	0.002 us

ROCm memoizes device properties internally, so the steady-state cost was never dramatic — roughly 16% of the kernel launch the same dispatch already pays. Caching removes it anyway, along with the 6.4 us first-call hit, and the three prefill.cuh call sites benefit for free.

The cache mirrors getMultiProcessorCount directly above it: thread_local so concurrent callers never race, and 0 treated as "not cached" so a device reporting 0 is simply not memoized rather than poisoning the entry. Device ids outside [0, 64) fall back to querying every time rather than indexing out of bounds.

getMaxSharedMemPerMultiprocessor is deliberately left uncached — identical shape, but it is only reached from pod.cuh / batch_pod.cuh, off this PR's path.

CDNA3 impact

None. gfx942 reports sharedMemPerBlock = 65536, so the comparison is byte-identical to the previous literal and no configuration changes status — only the message wording differs.

Test plan

On gfx950 (MI350X), ROCm 7.2.0, torch 2.9.1:

• Device query measured directly via HIP: sharedMemPerBlock = 163840.
• Full batch-decode suite at HEAD — 1872 passed, 0 failed, 0 errors. This file exercises the single-decode dispatch the guard sits in.
• Memoized value checked against a direct hipGetDeviceProperties query: identical (163840), stable over 1000 calls.
• Full suite on gfx942 hardware (MI300X, ctr-rack31-mi300x-2, commit 3708813f, pinned detached worktree): 27476 passed, 3585 skipped, 0 failed. This covers the second commit's memoization of getMaxSharedMemPerBlock, which matters because that helper is also called from three sites in prefill.cuh — so the change sits on the prefill dispatch path, not only decode.
• pre-commit run -a

On gfx942 (MI300X, ctr-rack31-mi300x-2, same toolchain — torch 2.9.1+rocm7.2.0, HIP 7.2.26015, amd-aiter 0.1.10):

• sharedMemPerBlock = 65536, so the comparison is byte-identical to the previous literal and no configuration changes status on CDNA3.
• Decode suite re-run at HEAD. The gfx942 check above predates the memoization commit. Caching does not change the value returned on any architecture, but CDNA3 is the only architecture where this guard actually rejects anything — on CDNA4 the 160 KB ceiling means it effectively never fires — so the rejection path is worth confirming on real gfx942 hardware before merge.

Side by side, same container stack, architecture the only variable:

	gfx942 (MI300X)	gfx950 (MI350X)
sharedMemPerBlock	65536 (64 KB)	163840 (160 KB)
CUs	304	256
warpSize	64	64

[Copilot]

Pull request overview

Updates the ROCm single-decode shared-memory guard to use the device-reported per-block shared memory limit (LDS) instead of a hard-coded 64 KB assumption, enabling valid gfx950 (CDNA4) configurations that require >64 KB dynamic shared memory.

Changes:

• Replace the smem_size > 65536U check with smem_size > getMaxSharedMemPerBlock(dev_id).
• Improve the error message to report the queried device limit in bytes.
• Hoist gpuGetDevice(&dev_id) so the device ID can be reused later in the partition-KV path.

---

💡 Add a code-review agent skill or configure MCP servers for context-aware, tailored reviews. Learn more in the docs.

[Copilot]

Pull request overview

Copilot reviewed 2 out of 2 changed files in this pull request and generated no new comments.