Masked Softmax

CdlMaskedSoftmax is a softmax with a memory: it computes probabilities over the last axis of a tensor, but lets you say how many of those positions are actually valid and ignores the rest. This is the attention building block from the Dive into Deep Learning book - the exact trick at the heart of every transformer, decoder, and sequence model. If ComfyDL's whole pitch is "build a BERT or Seq2Seq without writing code," this node is one of the pieces that makes that honest.

Why would you ever want to ignore positions? Because real sequences have different lengths. In a batch of sentences, sentence one might be 40 tokens and sentence two 12 - but the tensor is padded to 40. When you compute softmax over the attention scores, you want the padded positions to get essentially zero probability so the model doesn't attend to nothing. That's precisely what the mask does. And with valid_lens left unwired, it degrades gracefully to a plain softmax, so it doubles as the plain softmax node you'd otherwise go hunting for.

How it works

The implementation is the standard masked-softmax recipe from the d2l codebase: positions at or beyond each valid_lens entry get their scores replaced with -1e6 before softmax, which makes their softmax probability effectively zero after normalization. It operates on the last axis. The valid_lens shape decides how the mask is applied:

• A 1-D valid_lens gives one length per sample in the batch - repeated across the rows of that sample.
• A 2-D valid_lens gives a length per row, for the per-row masking you need in attention.

One subtle detail from the source worth respecting: it clones the input tensor before masking, so it never mutates the tensor you fed it. Nice for a debug-heavy graph.

Inputs and output

• X - the tensor to softmax over its last axis, e.g. [batch, seq, num] attention scores.
• valid_lens - optional cdlTensor. Omit it for plain softmax; provide it for masking.

The single output is output, the masked softmax result, same shape as X.

Installing ComfyDL

It ships with the pack:

cd ComfyUI/custom_nodes
git clone https://github.com/Cynthia-lxx/ComfyDL
pip install -r ComfyDL/requirements.txt

Restart ComfyUI. The only extra dependency is matplotlib; torch comes with ComfyUI. ComfyUI Manager users: search "ComfyDL", and if it's missing from the built-in list (the pack isn't published to the official Comfy Registry yet), use Install via Git URL with the repo link.

Common issues

The classic mistake is getting valid_lens semantics backwards - it's the count of valid positions, not the count of masked ones. Feed it 3 and positions 0–2 stay live, everything from 3 onward gets zeroed. The other trap is the shape contract: 1-D lengths apply per sample, 2-D per row, and if you hand it something in between the broadcast will silently do something surprising. And since softmax never outputs exact zeros (the -1e6 just makes them tiny), don't expect the masked rows to be bit-for-bit zero - they'll be vanishingly small, which is all the model cares about. Wire this into an attention or RNN decoder node and the pack's Seq2Seq examples start making a lot of sense.