reduce_max/reduce_sum dims=[1] returns zeros for fp32 tiles (fp32_reduce_acc path)

[zoecarver]

Summary

ttl.math.reduce_max and ttl.math.reduce_sum with dims=[1] silently return all zeros when the input tile is fp32. The same kernel on bfloat16 tiles returns correct per-row values. dims=[0] and dims=[0, 1] work correctly in fp32; only dims=[1] is affected.

The bug is in the fp32-accumulation lowering for reduce: passing options=\"--no-ttl-reduce-full-fp32\" to @ttl.operation makes fp32 dims=[1] return correct values.

Environment

• ttlang 1.0.0.dev4 (from the Python error traceback)
• Blackhole hardware (bh-qbae-15 / sterling-all container)
• layout=ttnn.TILE_LAYOUT, tested with both DRAM_MEMORY_CONFIG and L1_MEMORY_CONFIG

Minimal repro

import torch, ttnn, ttl
TILE = 32

@ttl.operation(grid=(1, 1))
def reduce_max_kernel(x, scaler, out):
    x_dfb   = ttl.make_dataflow_buffer_like(x,      shape=(1, 1), block_count=2)
    sc_dfb  = ttl.make_dataflow_buffer_like(scaler, shape=(1, 1), block_count=1)
    red_dfb = ttl.make_dataflow_buffer_like(x,      shape=(1, 1), block_count=2)
    out_dfb = ttl.make_dataflow_buffer_like(out,    shape=(1, 1), block_count=2)

    @ttl.compute()
    def compute():
        sc = sc_dfb.wait(); xb = x_dfb.wait()
        red_dfb.reserve().store(ttl.math.reduce_max(xb, sc, dims=[1]))
        mb = out_dfb.reserve()
        mb.store(ttl.math.broadcast(red_dfb.wait(), mb, dims=[1]))

    @ttl.datamovement()
    def dm_read():
        ttl.copy(scaler[0, 0], sc_dfb.reserve()).wait()
        ttl.copy(x[0, 0],      x_dfb.reserve()).wait()

    @ttl.datamovement()
    def dm_write():
        ttl.copy(out_dfb.wait(), out[0, 0]).wait()

# Input: row r filled with value r -> row max should be r.
for dtype_t, dtype_nn in [(torch.bfloat16, ttnn.bfloat16), (torch.float32, ttnn.float32)]:
    x  = torch.zeros(TILE, TILE, dtype=dtype_t)
    for r in range(TILE): x[r, :] = float(r)
    sc  = torch.ones(TILE, TILE, dtype=dtype_t)
    out = torch.zeros_like(x)
    cfg = dict(dtype=dtype_nn, layout=ttnn.TILE_LAYOUT,
               memory_config=ttnn.DRAM_MEMORY_CONFIG)
    device = ttnn.open_device(device_id=0)
    x_tt  = ttnn.from_torch(x,  device=device, **cfg)
    sc_tt = ttnn.from_torch(sc, device=device, **cfg)
    out_tt = ttnn.from_torch(out, device=device, **cfg)
    reduce_max_kernel(x_tt, sc_tt, out_tt)
    print(dtype_t, ttnn.to_torch(out_tt).float()[:4, 0])
    ttnn.close_device(device)

Expected vs observed

dtype	out[:, 0] first 8
bf16	[0, 1, 2, 3, 4, 5, 6, 7] (correct)
fp32	[0, 0, 0, 0, 0, 0, 0, 0] (wrong)

The same pattern exercised in test/python/simple_reduce_bcast.py works fine but uses bf16 + dims=[0, 1], so does not exercise this path.

Knobs tried

Config	fp32 dims=[1] result
default	zeros
fp32_dest_acc_en=True	zeros
fp32_dest_acc_en=True, dst_full_sync_en=False	zeros
L1_MEMORY_CONFIG instead of DRAM	zeros
options=\"--no-ttl-maximize-dst\"	zeros
options=\"--no-ttl-reduce-full-fp32\"	correct

Nested-with layout (matching simple_reduce_bcast.py) vs flat layout makes no difference; only --no-ttl-reduce-full-fp32 changes the outcome.

Scope

• reduce_max and reduce_sum both affected.
• Only dims=[1] on fp32 tiles. dims=[0] and dims=[0, 1] are correct.
• bf16 tiles unaffected regardless of dims.

Workaround

Pass options=\"--no-ttl-reduce-full-fp32\" to @ttl.operation until the fp32-accumulation reduce lowering is fixed for dims=[1].

[brnorris03]

I checked this against the tt-lang checkout pinned to tt-metal ccb46fcf93890b3bfb4c786de31a490e49586e26.

The failing generated fp32 kernel uses:

reduce_init<PoolType::MAX, ReduceDim::REDUCE_ROW, true>(...);
reduce_tile<PoolType::MAX, ReduceDim::REDUCE_ROW, true>(...);

The same repro passes when --no-ttl-reduce-full-fp32 makes the third template arg false. fp32 REDUCE_COL and REDUCE_SCALAR also pass with true, so the failure appears isolated to ReduceDim::REDUCE_ROW with enforce_fp32_accumulation=true.

Relevant tt-metal/LLK source links:

• reduce_init forwards enforce_fp32_accumulation into both unpack and math LLK init, and sets RISCV_DEBUG_REG_DBG_FEATURE_DISABLE bit 11 when it is true:
  https://github.com/tenstorrent/tt-metal/blob/ccb46fcf93890b3bfb4c786de31a490e49586e26/tt_metal/hw/inc/api/compute/reduce.h#L51-L61

• reduce_tile forwards the same template arg into llk_math_reduce:
  https://github.com/tenstorrent/tt-metal/blob/ccb46fcf93890b3bfb4c786de31a490e49586e26/tt_metal/hw/inc/api/compute/reduce.h#L135-L140

• Blackhole REDUCE_ROW calls reduce_row_perform_transpose<enforce_fp32_accumulation, ...>(); col/scalar do not use this row transpose branch:
  https://github.com/tenstorrent/tt-metal/blob/ccb46fcf93890b3bfb4c786de31a490e49586e26/tt_metal/third_party/tt_llk/tt_llk_blackhole/llk_lib/llk_math_reduce.h#L153-L184

• In the Blackhole full-fp32 row transpose branch, the LLK disables FP32 ALU mode, splits the fp32 datum into hi16/lo16 halves, runs MOVD2B / TRNSPSRCB / MOVB2D for each half, then re-enables FP32 mode:
  https://github.com/tenstorrent/tt-metal/blob/ccb46fcf93890b3bfb4c786de31a490e49586e26/tt_metal/third_party/tt_llk/tt_llk_blackhole/llk_lib/llk_math_reduce.h#L26-L64

• The unpack side also has special state for enforce_fp32_accumulation=true: it disables debug feature bit 11 and sets ALU_ACC_CTRL_Zero_Flag_disabled_src for the hi16/lo16 MOVB2D sequence:
  https://github.com/tenstorrent/tt-metal/blob/ccb46fcf93890b3bfb4c786de31a490e49586e26/tt_metal/hw/ckernels/blackhole/metal/llk_api/llk_unpack_AB_reduce_api.h#L19-L30
  https://github.com/tenstorrent/tt-metal/blob/ccb46fcf93890b3bfb4c786de31a490e49586e26/tt_metal/third_party/tt_llk/tt_llk_blackhole/llk_lib/llk_unpack_AB_reduce.h#L119-L126

• reduce_uninit<true> restores debug feature bit 11:
  https://github.com/tenstorrent/tt-metal/blob/ccb46fcf93890b3bfb4c786de31a490e49586e26/tt_metal/third_party/tt_llk/tt_llk_blackhole/llk_lib/llk_math_reduce.h#L320-L330
  https://github.com/tenstorrent/tt-metal/blob/ccb46fcf93890b3bfb4c786de31a490e49586e26/tt_metal/third_party/tt_llk/tt_llk_blackhole/llk_lib/llk_math_common.h#L16-L26

Related upstream issues:

• [Master Issue] Bit 11 tt-llk#1568 is the most relevant public tracker I found. It is the debug-bit-11 master issue, explicitly lists Reduce operation + Float32 as impacted, and notes that the Blackhole reduce workaround with bit 11 was added in tt_llk_blackhole/llk_lib/llk_math_reduce.h:
  [Master Issue] Bit 11 tt-llk#1568

• [subissue-guardrails-LLK] Remove need for DST 11 bit handling tt-metal#36411 tracks removing the dependency on DST bit 11 handling. It describes the current set/clear behavior across threads as non-predictable and proposes explicit SrcA format overrides for MOVD2B/MOVB2D instead:
  [subissue-guardrails-LLK] Remove need for DST 11 bit handling tt-metal#36411

• [ops]: sdpa with high precision compute config and bfp8 data fails with UndefinedBehavior in ttsim tt-metal#40583 is also relevant because it discusses a high-precision SDPA failure where reduce_c calls LLK row-reduce transpose and the enforce_fp32_accumulation state is inconsistent with subsequent MOVD2B use:
  [ops]: sdpa with high precision compute config and bfp8 data fails with UndefinedBehavior in ttsim tt-metal#40583

• [BH] Investigate whether disabling unpack dbg feature is necessary on BH tt-llk#878 is related to the Blackhole side of this state handling. It asks whether Blackhole still needs _llk_unpack_dbg_feature_disable_() and specifically mentions the reduce max row init site:
  [BH] Investigate whether disabling unpack dbg feature is necessary on BH tt-llk#878

I also tested the cleanup hypothesis directly. With the same single-tile fp32 row-reduce repro, the generated kernel was changed to emit reduce_uninit<true>() after reduce_tile<..., true>(), matching the tt-metal template-arg contract for reduce init/tile/uninit. That changed the cleanup sequence, but the default full-fp32 row reduce still produced an all-zero output and failed the torch comparison. Running the same repro with --no-ttl-reduce-full-fp32 changes the emitted reduce calls to reduce_init<..., REDUCE_ROW, false>() / reduce_tile<..., REDUCE_ROW, false>(); with the same Python tensor conversion, DFB movement, output packing, and result check, the device output matches torch.

That does not prove tt-lang has no required mitigation: tt-lang may still need to avoid enforce_fp32_accumulation=true for Blackhole row reduce if that is the supported workaround, or it may need to satisfy some additional tt-metal API precondition if one exists. However, based on the current tt-metal source, tt-lang is emitting the documented reduce API shape for fp32 accumulation: reduce_init<..., true>, reduce_tile<..., true>, and now reduce_uninit<true>. Since changing only the reduce full-fp32 template arg from true to false makes the repro pass, the observed all-zero output points at a bug in the implementation of the Blackhole REDUCE_ROW && enforce_fp32_accumulation LLK math/unpack sequence, not at Python conversion, DFB movement, pack cleanup, or a missing reduce uninit in tt-lang.

Co-authored by Codex