[Bug] Out of shared-memory's hardware limit when serving Qwen3-Next in vLLM

[shh2000]

Describe the bug

How to Reproduce:

1. export USE_FLAGGEMS=1 USE_C_EXTENSION=1
2. vllm serve Qwen/Qwen3-Coder-Next --tensor-parallel-size 4 --max-num-batched-tokens 16384 --max-num-seqs 2048 --served-model-name qwencf --port 9011
3. vllm bench serve would probabilistically trigger out-of-shared-memory exception. especially when concurrency >= 128, but no bugs on CUDA
4. FlagEval benchmarks would consistently but at random locations trigger out-of-shared-memory exception.
   The origin error log is:

[[0;36m(Worker_TP0 pid=46727)^[[0;0m ERROR 02-09 08:08:26 [multiproc_executor.py:824]   File "/usr/local/lib/python3.12/dist-packages/triton/runtime/jit.py", line 774, in run^M
^[[0;36m(Worker_TP0 pid=46727)^[[0;0m ERROR 02-09 08:08:26 [multiproc_executor.py:824]     launch_metadata = kernel.launch_metadata(grid, stream, *bound_args.values())^M
^[[0;36m(Worker_TP0 pid=46727)^[[0;0m ERROR 02-09 08:08:26 [multiproc_executor.py:824]                       ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^M
^[[0;36m(Worker_TP0 pid=46727)^[[0;0m ERROR 02-09 08:08:26 [multiproc_executor.py:824]   File "/usr/local/lib/python3.12/dist-packages/triton/compiler/compiler.py", line 490, in launch_metadata^M
^[[0;36m(Worker_TP0 pid=46727)^[[0;0m ERROR 02-09 08:08:26 [multiproc_executor.py:824]     self._init_handles()^M
^[[0;36m(Worker_TP0 pid=46727)^[[0;0m ERROR 02-09 08:08:26 [multiproc_executor.py:824]   File "/usr/local/lib/python3.12/dist-packages/triton/compiler/compiler.py", line 464, in _init_handles^M
^[[0;36m(Worker_TP0 pid=46727)^[[0;0m ERROR 02-09 08:08:26 [multiproc_executor.py:824]     raise_(OutOfResources(self.metadata.shared, max_shared, "shared memory"))^M
^[[0;36m(Worker_TP0 pid=46727)^[[0;0m ERROR 02-09 08:08:26 [multiproc_executor.py:824]   File "/usr/local/lib/python3.12/dist-packages/triton/compiler/compiler.py", line 456, in raise_^M
^[[0;36m(Worker_TP0 pid=46727)^[[0;0m ERROR 02-09 08:08:26 [multiproc_executor.py:824]     raise err^M
^[[0;36m(Worker_TP0 pid=46727)^[[0;0m ERROR 02-09 08:08:26 [multiproc_executor.py:824] triton.runtime.errors.OutOfResources: out of resource: shared memory, Required: 294936, Hardware limit: 232448. Reducing block sizes or `num_stages` may help.^M

FlagGems+triton3.5.0 would triggered similar problem. What could be the cause?

Environment details

1. H20*4
2. nvcr.io/nvidia/pytorch:25.12-py3
   and install libs one by one:
3. vllm==0.13.0
4. flag-gems 0574a3db0929, compiled with cpp wrapper(libtriton_jit)
5. flagtree==0.4.0+3.5
6. vllm-plugin-fl 37db70d1061b99

[shh2000]

The problem kernel is mm_kernel_general_host_tma:

@libentry()
@libtuner(
    configs=matmul_get_configs(),
    key=["M", "N", "K", "stride_am", "stride_bk", "dtype"],
    strategy=["align32", "align32", "align32", "align32", "align32", "default"],
    warmup=5,
    rep=5,
)
@triton.jit
def mm_kernel_general_host_tma(
    a_desc,
    b_desc,
    c_desc,
    M,
    N,
    K,
    stride_am,
    stride_ak,
    stride_bk,
    stride_bn,
    stride_cm,
    stride_cn,
    BLOCK_M: tl.constexpr,
    BLOCK_N: tl.constexpr,
    BLOCK_K: tl.constexpr,
    GROUP_M: tl.constexpr,
    A_ROW_MAJOR: tl.constexpr,
    B_ROW_MAJOR: tl.constexpr,
    dtype: tl.constexpr,
    enable_warp_specialization=True,
):
    pid = tl.program_id(0)
    grid_m = tl.cdiv(M, BLOCK_M)
    grid_n = tl.cdiv(N, BLOCK_N)

    width = GROUP_M * grid_n
    group_id = pid // width
    group_size = min(grid_m - group_id * GROUP_M, GROUP_M)
    pid_m = group_id * GROUP_M + (pid % group_size)
    pid_n = (pid % width) // (group_size)

    accumulator = tl.zeros((BLOCK_M, BLOCK_N), dtype=tl.float32)
    offset_am = (pid_m * BLOCK_M).to(tl.int32)
    offset_bn = (pid_n * BLOCK_N).to(tl.int32)
    iters = tl.cdiv(K, BLOCK_K)
    for k in range(iters):
        offset_ak = (k * BLOCK_K).to(tl.int32)

        if A_ROW_MAJOR:
            a = a_desc.load([offset_am, offset_ak])
        else:
            a_t = a_desc.load([offset_ak, offset_am])
            a = tl.trans(a_t)

        if B_ROW_MAJOR:
            b = b_desc.load([offset_ak, offset_bn])
        else:
            b_t = b_desc.load([offset_bn, offset_ak])
            b = tl.trans(b_t)

        if a_desc.dtype == tl.float16 or a_desc.dtype == tl.bfloat16:
            accumulator = tl.dot(a, b, acc=accumulator, allow_tf32=False)
        else:
            accumulator = tl.dot(a, b, acc=accumulator, input_precision="tf32x3")

    c = accumulator.to(c_desc.dtype)
    c_desc.store([offset_am, offset_bn], c)

where the autotune config is:

def matmul_get_configs(pre_hook=matmul_tma_set_block_size_hook):
    return [
        triton.Config(
            {"BLOCK_M": BM, "BLOCK_N": BN, "BLOCK_K": BK},
            num_stages=s,
            num_warps=w,
            pre_hook=pre_hook,
        )
        for BM in [32, 64, 128, 256]
        for BN in [32, 64, 128]
        for BK in [32, 64, 128]
        for s in [2, 3, 4]
        for w in [4, 8]
    ]