WIP: b12x updates

[askliar]

No description provided.

[gemini-code-assist]

Code Review

This pull request introduces support for FlashInfer B12x fused MoE kernels on SM12x devices, adds a torch.compile-based fast path for BF16 linear layers, and includes a tinygemm BF16 dispatch mechanism. My review identified several performance and correctness issues: the scale factor layout conversion in FlashInferB12xExperts should be cached to avoid per-pass overhead, the dynamo cache limit constants in utils.py need to be moved to the module level to resolve test failures, _apply_tinygemm should avoid repeated zero-tensor allocations, and the test case for inductor backends needs to be corrected to reflect the actual configuration.

[gemini-code-assist]

The layout conversion flashinfer_convert_sf_to_mma_layout is called on every forward pass for both w1 and w2 scale factors. Since expert weights and their scales are typically static after loading, this conversion is a significant performance bottleneck. It should be performed once during weight processing or cached using the tensor's data pointer as a key.

        # Convert swizzled 3D scale factors [E, M, K_sf] to 6D MMA layout
        # expected by the SM12x kernel's _get_weight_views().
        sf_vec_size = 16
        if not hasattr(self, "_w1_sf_mma") or self._w1_sf_mma_key != self.w1_scale.data_ptr():
            E_w1, M_w1, K_sf_w1 = self.w1_scale.shape
            self._w1_sf_mma = flashinfer_convert_sf_to_mma_layout(
                self.w1_scale.reshape(E_w1 * M_w1, K_sf_w1),
                m=M_w1,
                k=K_sf_w1 * sf_vec_size,
                num_groups=E_w1,
                sf_vec_size=sf_vec_size,
            )
            self._w1_sf_mma_key = self.w1_scale.data_ptr()
        w1_sf_mma = self._w1_sf_mma

        if not hasattr(self, "_w2_sf_mma") or self._w2_sf_mma_key != self.w2_scale.data_ptr():
            E_w2, M_w2, K_sf_w2 = self.w2_scale.shape
            self._w2_sf_mma = flashinfer_convert_sf_to_mma_layout(
                self.w2_scale.reshape(E_w2 * M_w2, K_sf_w2),
                m=M_w2,
                k=K_sf_w2 * sf_vec_size,
                num_groups=E_w2,
                sf_vec_size=sf_vec_size,
            )
            self._w2_sf_mma_key = self.w2_scale.data_ptr()
        w2_sf_mma = self._w2_sf_mma

[gemini-code-assist]

The constants _CACHE_LIMIT and _ACCUMULATED_LIMIT are defined as local variables within force_inductor_max_autotune_gemm. However, the tests in tests/model_executor/test_compile_bf16_linear_dispatch.py attempt to access them as module-level attributes of vllm.model_executor.layers.utils using names like _DYNAMO_CACHE_SIZE_LIMIT. This will result in an AttributeError during testing. These should be moved to the module level and renamed to match the test expectations.

_DYNAMO_CACHE_SIZE_LIMIT = 65536
_DYNAMO_ACCUMULATED_CACHE_SIZE_LIMIT = 1048576

[gemini-code-assist]

Allocating a new zero tensor via torch.zeros on every call to _apply_tinygemm when bias is None introduces unnecessary overhead, especially for small batch sizes where tinygemm is intended to provide the most benefit. Consider using a cached zero tensor or modifying the underlying kernel to handle a None bias directly.

[gemini-code-assist]

This assertion will fail because the implementation in vllm/model_executor/layers/utils.py explicitly sets inductor_config.max_autotune_gemm_backends = "ATEN,TRITON", excluding CUTLASS. The implementation comment also notes "no inductor CUTLASS". Please update the test to match the intended implementation.

Suggested change

	assert "CUTLASS" in inductor_config.max_autotune_gemm_backends
	assert "CUTLASS" not in inductor_config.max_autotune_gemm_backends

[mergify]

This pull request has merge conflicts that must be resolved before it can be
merged. Please rebase the PR, @askliar.

https://docs.github.com/en/pull-requests/collaborating-with-pull-requests/working-with-forks/syncing-a-fork