[CI][Bugfix] Fix failing Blackwell test

[MatthewBonanni]

Purpose

Blackwell test is broken on main. This PR fixes the underlying issue. The dtype was changing between tests, causing an assertion failure in the fused MoE buffer. Now, the buffer is rebuilt if the dtype (or device) changes

Test Plan

pytest tests/kernels/moe/test_nvfp4_moe.py

Test Result

Test passes

---

Essential Elements of an Effective PR Description Checklist

• The purpose of the PR, such as "Fix some issue (link existing issues this PR will resolve)".
• The test plan, such as providing test command.
• The test results, such as pasting the results comparison before and after, or e2e results
• (Optional) The necessary documentation update, such as updating supported_models.md and examples for a new model.
• (Optional) Release notes update. If your change is user facing, please update the release notes draft in the Google Doc.

[gemini-code-assist]

Code Review

This pull request fixes a bug in the Blackwell test where a shared buffer was not being recreated when the dtype changed between tests. The fix correctly adds checks for device and dtype to the buffer reallocation logic.

However, the implementation of SharedResizableBuffer.get is not thread-safe. Since the buffers are shared as class attributes, concurrent requests can lead to a race condition, causing memory corruption. I've added a critical review comment with a suggestion to add a lock to ensure thread safety.

[gemini-code-assist]

While this change correctly handles buffer recreation for different dtypes or devices in sequential runs (like tests), it is not thread-safe.

The SharedResizableBuffer instances are used as class attributes in FusedMoEModularKernel, making them singletons shared across all MoE layers. If multiple threads execute MoE layers concurrently, they can race to check and reallocate the buffer in get(), leading to memory corruption or incorrect results.

To ensure thread safety, the read-modify-write operation on self.buffer must be atomic. Please protect it with a threading.Lock.

This would involve:

1. Adding import threading at the top of the file.
2. Initializing a lock in SharedResizableBuffer.__init__: self._lock = threading.Lock().
3. Wrapping the logic in get() with the lock: with self._lock: ....

Example of the final SharedResizableBuffer class:

import threading
from math import prod
...

class SharedResizableBuffer:

    def __init__(self):
        self.buffer = None
        self._lock = threading.Lock()

    def get(self, shape: tuple[int, ...], device: torch.device,
            dtype: torch.dtype):
        with self._lock:
            shape_numel = prod(shape)
            if (self.buffer is None or self.buffer.numel() < shape_numel or
                self.buffer.device != device or self.buffer.dtype != dtype):
                self.buffer = torch.empty(shape_numel, device=device, dtype=dtype)
            return self.buffer[:shape_numel].view(*shape)

Since these changes are outside the diff, I cannot provide a direct code suggestion. This is a critical issue that should be addressed to prevent race conditions in production.

[LucasWilkinson]

Ya seems fine; I guess the worst case is we just end up back in the previous behavior of just allocating every time 👍 thanks for the fix!

[njhill]

@MatthewBonanni do you know how this was introduced in the first place? Did that test group run on the PR in question? is it that it only fails intermittently?

[MatthewBonanni]

@njhill It looks like this was introduced by the DBO PR, #23693 . The blackwell test was not run on that PR: https://buildkite.com/vllm/ci/builds/30923

[njhill]

@njhill It looks like this was introduced by the DBO PR, #23693 . The blackwell test was not run on that PR: https://buildkite.com/vllm/ci/builds/30923

Thanks @MatthewBonanni, I guess that means we should potentially adjust the scoping of the blackwell tests.