page_attn: make global-max reduction order-invariant

vllm/custom-esimd-kernels-vllm/csrc/eagle/eagle.sycl

@@ -490,6 +490,11 @@ void page_attn_decode(
   torch::Tensor temp_p = torch::zeros({ static_cast<long>(totalSize) },
     torch::device(kv_cache.device()).dtype(torch::kFloat32));
 
+  // The global-max reduction in the phase-1 kernels is an unconditional atomic fmax,
+  // so this slice must start below any real score. The rest of temp_p stays zeroed.
+  temp_p.narrow(0, static_cast<long>(szP + szGroupMax),
+    static_cast<long>(szGlobalMax)).fill_(FP32_MIN);
+
   float* pState = (float*)temp_p.data_ptr();
   float* pGroupMax = pState + szP;
   float* pGlobalMax = pGroupMax + szGroupMax;

vllm/custom-esimd-kernels-vllm/csrc/eagle/page.attn.fp8.h

@@ -366,20 +366,17 @@ ESIMD_INLINE void sdpaDecodeGqa4Phase1Fp8(
         1,
         uint32_t>(pPollP, atomicOffset);
 
-      if (0 == arrivalId) {
-        atomic_update<
-          __ESIMD_NS::atomic_op::fcmpxchg,
-          float,
-          1,
-          uint32_t>(pGlobalMax, atomicOffset, ppMax, zeros);
-      }
-      else {
-        atomic_update<
-          __ESIMD_NS::atomic_op::fmax,
-          float,
-          1,
-          uint32_t>(pGlobalMax, atomicOffset, ppMax);
-      }
+      // Order-invariant global-max reduction: every arrival does an atomic fmax, so
+      // pGlobalMax must be seeded to FP32_MIN (done in page_attn_decode) rather than 0.
+      // See page.attn.h (fp16 path) for the full rationale; the arrival counter is now
+      // unused but kept so the pPollP layout matches the fp16 kernels.
+      (void)arrivalId;
+      (void)zeros;
+      atomic_update<
+        __ESIMD_NS::atomic_op::fmax,
+        float,
+        1,
+        uint32_t>(pGlobalMax, atomicOffset, ppMax);
 
       outputOffset += pStride;
       outputMaxOffset += maxStride;
@@ -908,20 +905,17 @@ ESIMD_INLINE void sdpaDecodeGqa2Phase1Fp8(
         1,
         uint32_t>(pPollP, atomicOffset);
 
-      if (0 == arrivalId) {
-        atomic_update<
-          __ESIMD_NS::atomic_op::fcmpxchg,
-          float,
-          1,
-          uint32_t>(pGlobalMax, atomicOffset, ppMax, zeros);
-      }
-      else {
-        atomic_update<
-          __ESIMD_NS::atomic_op::fmax,
-          float,
-          1,
-          uint32_t>(pGlobalMax, atomicOffset, ppMax);
-      }
+      // Order-invariant global-max reduction: every arrival does an atomic fmax, so
+      // pGlobalMax must be seeded to FP32_MIN (done in page_attn_decode) rather than 0.
+      // See page.attn.h (fp16 path) for the full rationale; the arrival counter is now
+      // unused but kept so the pPollP layout matches the fp16 kernels.
+      (void)arrivalId;
+      (void)zeros;
+      atomic_update<
+        __ESIMD_NS::atomic_op::fmax,
+        float,
+        1,
+        uint32_t>(pGlobalMax, atomicOffset, ppMax);
     }
   }
 }

vllm/custom-esimd-kernels-vllm/csrc/eagle/page.attn.h

@@ -200,20 +200,22 @@ ESIMD_INLINE void sdpaDecodeGqa4Phase1(
         1,
         uint32_t>(pPollP, atomicOffset);
 
-      if (0 == arrivalId) {
-        atomic_update<
-          __ESIMD_NS::atomic_op::fcmpxchg,
-          float,
-          1,
-          uint32_t>(pGlobalMax, atomicOffset, ppMax, zeros);
-      }
-      else {
-        atomic_update<
-          __ESIMD_NS::atomic_op::fmax,
-          float,
-          1,
-          uint32_t>(pGlobalMax, atomicOffset, ppMax);
-      }
+      // Order-invariant global-max reduction: every arrival does an atomic fmax, so
+      // pGlobalMax must be seeded to FP32_MIN (done in page_attn_decode) rather than 0.
+      // The prior fcmpxchg-if-first / fmax-else scheme assumed a 0 seed and had two
+      // problems: (1) with a single 64-tile (maxSeqLen<=64) the first arrival's
+      // fcmpxchg is the only writer, so any seed != 0 leaves pGlobalMax unchanged;
+      // (2) even with a 0 seed it is not order-invariant -- a later arrival's
+      // fmax(0, neg)=0 can mask a negative max before arrival-0's fcmpxchg overwrites
+      // it, losing the true max when the whole head's max is negative. fmax-from-
+      // FP32_MIN removes both (the arrival counter is now unused).
+      (void)arrivalId;
+      (void)zeros;
+      atomic_update<
+        __ESIMD_NS::atomic_op::fmax,
+        float,
+        1,
+        uint32_t>(pGlobalMax, atomicOffset, ppMax);
 
       outputOffset += pStride;
       outputMaxOffset += maxStride;
@@ -766,20 +768,22 @@ ESIMD_INLINE void sdpaDecodeGqa2Phase1(
         1,
         uint32_t>(pPollP, atomicOffset);
 
-      if (0 == arrivalId) {
-        atomic_update<
-          __ESIMD_NS::atomic_op::fcmpxchg,
-          float,
-          1,
-          uint32_t>(pGlobalMax, atomicOffset, ppMax, zeros);
-      }
-      else {
-        atomic_update<
-          __ESIMD_NS::atomic_op::fmax,
-          float,
-          1,
-          uint32_t>(pGlobalMax, atomicOffset, ppMax);
-      }
+      // Order-invariant global-max reduction: every arrival does an atomic fmax, so
+      // pGlobalMax must be seeded to FP32_MIN (done in page_attn_decode) rather than 0.
+      // The prior fcmpxchg-if-first / fmax-else scheme assumed a 0 seed and had two
+      // problems: (1) with a single 64-tile (maxSeqLen<=64) the first arrival's
+      // fcmpxchg is the only writer, so any seed != 0 leaves pGlobalMax unchanged;
+      // (2) even with a 0 seed it is not order-invariant -- a later arrival's
+      // fmax(0, neg)=0 can mask a negative max before arrival-0's fcmpxchg overwrites
+      // it, losing the true max when the whole head's max is negative. fmax-from-
+      // FP32_MIN removes both (the arrival counter is now unused).
+      (void)arrivalId;
+      (void)zeros;
+      atomic_update<
+        __ESIMD_NS::atomic_op::fmax,
+        float,
+        1,
+        uint32_t>(pGlobalMax, atomicOffset, ppMax);
     }
   }
 }