XQA always sets dynamic smem size on rank 0

[zackangelo]

flashinfer/csrc/xqa/mha.cu

Lines 2589 to 2594 in 9bf007d

	static uint32_t configureKernel() {
	uint32_t size;
	cudaMemcpyFromSymbol(&size, smemSize, sizeof(smemSize));
	cudaFuncSetAttribute(kernel_mha, cudaFuncAttributeMaxDynamicSharedMemorySize, size);
	return size;
	}

I believe this will always set dynamic shared memory size on GPU 0 when the library loads. This makes it difficult to integrate in multi-GPU environments where a PID has access to more than one GPU because we can't control where this value is applied.

If this value is applied on the incorrect GPU all kernel launches will fail with a CUDA invalid argument error.

[yzh119]

I believe this will always set dynamic shared memory size on GPU 0 when the library loads.

Is it because we make it a static function? How about we inline it instead?

[zackangelo]

I was able to work around this by specifying CUDA_VISIBLE_DEVICES on our GPU worker PIDs. Seems like inline would work better but the memcpy seems undesirable in the kernel run path?

[bledden]

!claim

[flashinfer-bot]

Issue assigned to @bledden.

[saltyminty]

Hi! Following up to see if this issue is still being hit? We weren't able to reproduce it on our end in v0.6.4

[zackangelo]

We worked around it in our stack, so not sure. We can probably go ahead and close it!