zero out attn output

fixes and zero out input to fi kernel

vllm/model_executor/layers/attention/mla_attention.py

@@ -518,7 +518,7 @@ class MLAAttention(nn.Module, AttentionLayerBase):
                 self._k_scale,
             )
             if self.attn_backend.accept_output_buffer:
-                output = torch.empty(output_shape, dtype=q.dtype, device=q.device)
+                output = torch.zeros(output_shape, dtype=q.dtype, device=q.device)
                 torch.ops.vllm.unified_mla_attention_with_output(
                     q,
                     kv_c_normed,
@@ -692,6 +692,15 @@ class MLAAttention(nn.Module, AttentionLayerBase):
                 assert attn_metadata.decode is not None
             attn_out, lse = self.impl.forward_mqa(mqa_q, kv_cache, attn_metadata, self)
 
+            # Zero out padded region for CUDA graphs. Padded requests have
+            # seq_lens=0 in the decode metadata (device tensor updated in-place
+            # before each graph replay), so masked_fill_ always launches a kernel
+            # during capture yet only zeros padding slots during replay.
+            if attn_metadata.decode is not None:
+                decode_seq_lens = attn_metadata.decode.seq_lens
+                pad_mask = (decode_seq_lens == 0).view(-1, 1, 1)
+                attn_out[:decode_seq_lens.shape[0]].masked_fill_(pad_mask, 0)
+
             # correct dcp attn_out with lse.
             if self.impl.dcp_world_size > 1:
                 if self.dcp_a2a:

vllm/model_executor/model_loader/default_loader.py

@@ -316,6 +316,13 @@ class DefaultModelLoader(BaseModelLoader):
         if not (model_config.is_moe and parallel_config.enable_expert_parallel):
             return
 
+        # When EPLB is enabled, redundant physical expert slots may map to
+        # logical experts that belong to other ranks in the default partition.
+        # The weight loader needs to see ALL logical expert weights so it can
+        # populate these redundant slots.  Skip the filter entirely.
+        if parallel_config.enable_eplb:
+            return
+
         num_experts = model_config.get_num_experts()
         if num_experts <= 0:
             return

vllm/model_executor/model_loader/ep_weight_filter.py

@@ -73,4 +73,9 @@ def should_skip_weight(
     if eid is None:
         # Not an expert weight (dense / shared-expert / embedding) → keep.
         return False
+    # Only skip heavy weight tensors, never scale/metadata tensors.
+    # Scale tensors are tiny and some backends need them from ALL experts
+    # (e.g. FlashInfer NVFP4 computes a global max of activation scales).
+    if not weight_name.endswith(".weight"):
+        return False
     return eid not in local_expert_ids

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

@@ -181,7 +181,12 @@ class FlashInferMLAImpl(MLACommonImpl[MLACommonMetadata]):
         if self.bmm2_scale is None:
             self.bmm2_scale = layer._v_scale_float
 
+        out = torch.zeros(q.shape[0], q.shape[2], self.kv_lora_rank,
+                  dtype=torch.bfloat16, device=q.device)
+
+        self._workspace_buffer.fill_(0)
         o = trtllm_batch_decode_with_kv_cache_mla(
+            out=out,
             query=q,
             kv_cache=kv_c_and_k_pe_cache.unsqueeze(1),
             workspace_buffer=self._workspace_buffer,