Merge branch 'v0.11.0-dev-kvscale' into 'v0.11.0-dev'

V0.11.0 dev kvscale

See merge request dcutoolkit/deeplearing/vllm!378

vllm/attention/layer.py

@@ -282,9 +282,11 @@ class Attention(nn.Module, AttentionLayerBase):
         `vllm.forward_context.get_forward_context().attn_metadata`.
         """
         if self.calculate_kv_scales:
-            attn_metadata = get_forward_context().attn_metadata
-            if attn_metadata.enable_kv_scales_calculation:
-                self.calc_kv_scales(query, key, value)
+            # attn_metadata = get_forward_context().attn_metadata
+            # if attn_metadata.enable_kv_scales_calculation:
+            #     self.calc_kv_scales(query, key, value)
+            torch.ops.vllm.maybe_calc_kv_scales(query, key, value,
+                                                self.layer_name)  
 
         output_dtype = query.dtype
         if self.query_quant is not None:
@@ -570,6 +572,38 @@ def maybe_save_kv_layer_to_connector(
     connector.save_kv_layer(layer_name, kv_cache_layer,
                             attn_metadata[layer_name])
 
+def maybe_calc_kv_scales(
+    query: torch.Tensor,
+    key: torch.Tensor,
+    value: torch.Tensor,
+    layer_name: str,
+) -> None:
+    forward_context: ForwardContext = get_forward_context()
+    self = forward_context.no_compile_layers[layer_name]
+
+    # Only calculate if the layer's calculate_kv_scales flag is True
+    # This flag gets set to False after the first forward pass
+    if not self.calculate_kv_scales:
+        return
+
+    self.calc_kv_scales(query, key, value)
+
+
+def maybe_calc_kv_scales_fake(
+    query: torch.Tensor,
+    key: torch.Tensor,
+    value: torch.Tensor,
+    layer_name: str,
+) -> None:
+    return
+
+
+direct_register_custom_op(
+    op_name="maybe_calc_kv_scales",
+    op_func=maybe_calc_kv_scales,
+    mutates_args=["query", "key", "value"],
+    fake_impl=maybe_calc_kv_scales_fake,
+)
 
 def unified_attention(
     query: torch.Tensor,

vllm/v1/attention/backends/flash_attn.py

@@ -593,9 +593,9 @@ class FlashAttentionImpl(AttentionImpl):
             block_table = attn_metadata.block_table
             scheduler_metadata = attn_metadata.scheduler_metadata
 
-            descale_shape = (cu_seqlens_q.shape[0] - 1, self.num_kv_heads)
-
+            # descale_shape = (cu_seqlens_q.shape[0] - 1, self.num_kv_heads)
             if not current_platform.is_rocm():
+                 descale_shape = (cu_seqlens_q.shape[0] - 1, self.num_kv_heads)
                  flash_attn_varlen_func(
                     q=query[:num_actual_tokens],
                     k=key_cache,
@@ -643,8 +643,9 @@ class FlashAttentionImpl(AttentionImpl):
                     scheduler_metadata=scheduler_metadata,
                     # fa_version=self.vllm_flash_attn_version,
                     # q_descale=layer._q_scale.expand(descale_shape),
-                    # k_descale=layer._k_scale.expand(descale_shape),
-                    # v_descale=layer._v_scale.expand(descale_shape),
+                    q_descale=None,
+                    k_descale=layer._k_scale,
+                    v_descale=layer._v_scale,
                     # num_splits=attn_metadata.max_num_splits,
                     s_aux=self.sinks,
                     is_prefix_cache=True,
@@ -699,8 +700,9 @@ class FlashAttentionImpl(AttentionImpl):
                 prefix_scheduler_metadata=attn_metadata.prefix_scheduler_metadata,
                 suffix_scheduler_metadata=attn_metadata.scheduler_metadata,
                 # q_descale=layer._q_scale,
-                # k_descale=layer._k_scale,
-                # v_descale=layer._v_scale,
+                q_descale=None,
+                k_descale=layer._k_scale,
+                v_descale=layer._v_scale,
             )
         return output
 
@@ -778,6 +780,9 @@ class FlashAttentionImpl(AttentionImpl):
                 # q_descale=layer._q_scale.expand(descale_shape),
                 # k_descale=layer._k_scale.expand(descale_shape),
                 # v_descale=layer._v_scale.expand(descale_shape),
+                q_descale=None,
+                k_descale=layer._k_scale,
+                v_descale=layer._v_scale,
                 is_prefix_cache=False,
             )
 
@@ -931,12 +936,12 @@ def cascade_attention(
             return_softmax_lse=True,
             scheduler_metadata=prefix_scheduler_metadata,
             # fa_version=fa_version,
-            # q_descale=q_descale.expand(descale_shape)
-            # if q_descale is not None else None,
-            # k_descale=k_descale.expand(descale_shape)
-            # if k_descale is not None else None,
-            # v_descale=v_descale.expand(descale_shape)
-            # if v_descale is not None else None,
+            q_descale=q_descale.expand(descale_shape)
+            if q_descale is not None else None,
+            k_descale=k_descale.expand(descale_shape)
+            if k_descale is not None else None,
+            v_descale=v_descale.expand(descale_shape)
+            if v_descale is not None else None,
             is_prefix_cache=True,
         )
 
@@ -984,12 +989,12 @@ def cascade_attention(
             return_softmax_lse=True,
             scheduler_metadata=suffix_scheduler_metadata,
             # fa_version=fa_version,
-            # q_descale=q_descale.expand(descale_shape)
-            # if q_descale is not None else None,
-            # k_descale=k_descale.expand(descale_shape)
-            # if k_descale is not None else None,
-            # v_descale=v_descale.expand(descale_shape)
-            # if v_descale is not None else None,
+            q_descale=q_descale.expand(descale_shape)
+            if q_descale is not None else None,
+            k_descale=k_descale.expand(descale_shape)
+            if k_descale is not None else None,
+            v_descale=v_descale.expand(descale_shape)
+            if v_descale is not None else None,
             is_prefix_cache=True,
         )