Fix stale is_sparse_impl reference and trim comment bloat

Signed-off-by: Aakif Nawaz <aakif.nawaz@amd.com>

Signed-off-by: Aakif Nawaz <aakif.nawaz@amd.com>

Author: akii96

vllm/model_executor/layers/attention/mla_attention.py

@@ -557,10 +557,7 @@ def __init__(
         kv_b_proj._cpu_skip_gemm_dispatch = True
         self.use_direct_call = not current_platform.opaque_attention_op()
 
-        # the fused sparse-MLA Q-prep kernel folds the main RoPE into
-        # the decode kernel, so it needs the rotary embedding's cos/sin cache.
-        # Without a rotary embedding there is no RoPE to fold, so disable the
-        # fused path to keep the wrapper / KV-write / decode branches consistent.
+        # Disable fused Q-prep path when no rotary_emb is available.
         self.rotary_emb = rotary_emb
         self._fused_rope_cos_sin: tuple[torch.Tensor, torch.Tensor] | None = None
         if getattr(self.impl, "use_fused_qk_rope_cache", False) and rotary_emb is None:
@@ -978,14 +975,11 @@ def forward_impl(
                 mqa_ql_nope = mqa_ql_nope.transpose(0, 1)
 
             use_fused_qk_rope_cache = (
-                is_sparse_impl
+                self.impl.is_sparse
                 and getattr(self.impl, "use_fused_qk_rope_cache", False)
                 and self.rotary_emb is not None
             )
             if use_fused_qk_rope_cache:
-                # fold RoPE + Q-concat + KV-concat + KV-cache-write
-                # into one kernel. positions is threaded in as a real argument
-                # (survives torch.compile / CUDA graph replay).
                 if positions is None:
                     raise RuntimeError(
                         "Fused MLA Q-prep is enabled but `positions` was not "

vllm/model_executor/layers/mla.py

@@ -127,9 +127,6 @@ def __init__(
             rotary_emb=self.rotary_emb,
         )
 
-        # when the decode impl fuses RoPE into its Q-prep kernel
-        # (fused_qk_rope_concat_and_cache_mla), defer the main RoPE here and
-        # pass positions down to the impl as a forward() argument.
         self._defer_rope_to_fused_kernel = self.rotary_emb is not None and getattr(
             self.mla_attn.impl, "use_fused_qk_rope_cache", False
         )

vllm/v1/attention/backends/mla/rocm_aiter_mla_sparse.py

@@ -709,8 +709,6 @@ def __init__(
             (q_concat_shape, vllm_config.model_config.dtype),
         )
 
-        # fold decode RoPE + Q-concat + KV-concat + KV-cache-write into
-        # a single aiter kernel (fused_qk_rope_concat_and_cache_mla).
         self.use_fused_qk_rope_cache = (
             rocm_aiter_ops.is_fused_mla_qkprep_enabled()
             and self.kv_cache_dtype.startswith("fp8")
@@ -727,8 +725,6 @@ def do_kv_cache_update(
         k_scale: torch.Tensor,
     ) -> None:
         if self.use_fused_qk_rope_cache:
-            # The KV-cache write (concat + RoPE + fp8 quant) is folded into
-            # fused_qk_rope_concat_and_cache_mla
             return
         super().do_kv_cache_update(
             kv_c_normed, k_pe, kv_cache, slot_mapping, kv_cache_dtype, k_scale
@@ -748,20 +744,10 @@ def fused_qk_rope_concat_and_cache(
         sin_cache: torch.Tensor,
         is_neox: bool,
     ) -> torch.Tensor:
-        """Fused decode Q-prep.
-
-        Applies RoPE to ``q_pe``/``k_pe``, concatenates ``ql_nope`` with the
-        RoPE'd ``q_pe`` into ``q_out``, and writes the concatenated,
-        fp8-quantized latent+rope KV row into ``kv_cache``. The fused path only
-        runs with an fp8 KV cache, so ``q_out`` is produced directly in fp8
-        (quantized by the kernel with ``q_scale``).
-        """
+        """Fused RoPE + Q-concat + KV-concat + fp8 KV-cache write."""
         num_tokens, num_q_heads = ql_nope.shape[:2]
-        # The aiter kernel builds the nope half of q_out (and the KV latent) by
-        # copying ql_nope and REQUIRES it to be contiguous: a strided ql_nope
-        # silently yields wrong q_out values. Force contiguity here.
+        # aiter kernel requires contiguous ql_nope
         ql_nope = ql_nope.contiguous()
-        # Emit q_out directly in fp8 when the KV cache is fp8
         q_out_dtype = (
             current_platform.fp8_dtype()
             if self.kv_cache_dtype.startswith("fp8")