EpilogueDirectStore warp decomposition uses WarpShape::kN where WarpCount is meant

Description

EpilogueDirectStore's constructor computes the warp decomposition as:

int warp_id_mn = warp_idx % (WarpCount::kM * WarpShape::kN);

(include/cutlass/epilogue/threadblock/epilogue_direct_store.h line 171)

The second factor is a typo: it mixes the warp shape's WarpShape::kN into what should be the warp count. EpilogueBase does this correctly (include/cutlass/epilogue/threadblock/epilogue_base.h lines 219-220):

int warp_mn = warp_idx % (WarpCount::kM * WarpCount::kN);

When WarpCount::kN > WarpShape::kN the modulo truncates warp_idx early and two different warps receive the same (warp_m, warp_n) coordinates. For all currently instantiated direct-store configs WarpCount::kN <= WarpShape::kN, so nothing is observably wrong today; the defect surfaces only for shapes where Shape::kN > WarpShape::kN^2.

Suggested fix

Use warp_idx % WarpCount::kCount (or % (WarpCount::kM * WarpCount::kN) to mirror EpilogueBase).

Start at include/cutlass/epilogue/threadblock/epilogue_direct_store.h line 171 and compare the warp decomposition with include/cutlass/epilogue/threadblock/epilogue_base.h lines 219-220. Update the decomposition to use the warp count, then verify that the direct-store calculation assigns distinct coordinates for the described warp-count configurations.