update VLLM_USE_OPT_RESHAPE_AND_CACHE to support bf16 and qwen3-dense

vllm/model_executor/model_loader/utils.py

@@ -211,6 +211,9 @@ def _get_model_architecture(model_config: ModelConfig) -> tuple[type[nn.Module],
                         os.environ['VLLM_USE_FUSE_SILU_AND_MUL'] = '1'
                     if not envs.is_set("VLLM_USE_OPT_RESHAPE_AND_CACHE"):
                         os.environ['VLLM_USE_OPT_RESHAPE_AND_CACHE'] = '1'
+                # if architectures in [['Qwen3ForCausalLM']]:
+                #     if not envs.is_set("VLLM_USE_OPT_RESHAPE_AND_CACHE"):
+                #         os.environ['VLLM_USE_OPT_RESHAPE_AND_CACHE'] = '0'
 
             if architectures in [['DeepseekV32ForCausalLM']]:
                 if not envs.is_set("VLLM_USE_V32_ENCODE"):

vllm/v1/attention/backends/flash_attn.py

@@ -898,7 +898,7 @@ class FlashAttentionImpl(AttentionImpl):
         # op uses the slot_mapping's shape to determine the number of
         # actual tokens.
         if current_platform.is_rocm():
-            if envs.VLLM_USE_OPT_RESHAPE_AND_CACHE and key.dtype == value.dtype == torch.float16:
+            if envs.VLLM_USE_OPT_RESHAPE_AND_CACHE:
                 from lightop import reshape_and_cache_cuda
                 reshape_and_cache_cuda(
                     key,