Does Triton have a None Type variable?

[priyammaz]

I am relatively new to this, but am working on a dropout implementation for the forward and backward pass.

This is what I would like to do:

if USE_DROPOUT:
  dropout_seed = SEED + batch_head_offset
  q_offsets_drop = offsets_row[None, :]
  kv_offsets_drop = offsets_col[:, None]
  dropout_offset = q_offsets_drop * SEQ_LEN + kv_offsets_drop
  random_vals = tl.rand(dropout_seed, dropout_offset)
  dropout_mask = random_vals > dropout_p
  P_T_block = tl.where(dropout_mask, P_T_block / (1.0 - dropout_p), 0.0)
  
if USE_DROPOUT:
  dP_T_block = tl.where(dropout_mask, dP_T_block / (1.0 - dropout_p), 0.0)

This throws a compiler error that the dropout_mask (in the second if statement) doesn't exist. I guess it doesnt understand that if a variable is created for one condition, it would exist for all other occurrences of that condition.

So I can solve this in two ways, neither of which I like. I can either recompute the mask again, as its a seeded generator so we get the same results anyway but that adds extra ops we can avoid

if USE_DROPOUT:
  dropout_seed = SEED + batch_head_offset
  q_offsets_drop = offsets_row[None, :]
  kv_offsets_drop = offsets_col[:, None]
  dropout_offset = q_offsets_drop * SEQ_LEN + kv_offsets_drop
  random_vals = tl.rand(dropout_seed, dropout_offset)
  dropout_mask = random_vals > dropout_p
  P_T_block = tl.where(dropout_mask, P_T_block / (1.0 - dropout_p), 0.0)
  
if USE_DROPOUT:
  dropout_seed = SEED + batch_head_offset
  q_offsets_drop = offsets_row[None, :]
  kv_offsets_drop = offsets_col[:, None]
  dropout_offset = q_offsets_drop * SEQ_LEN + kv_offsets_drop
  random_vals = tl.rand(dropout_seed, dropout_offset)
  dropout_mask = random_vals > dropout_p
  dP_T_block = tl.where(dropout_mask, dP_T_block / (1.0 - dropout_p), 0.0)

Or we can always create a variable called dropout_mask at the start like this

dropout_mask = tl.full((BLOCK_SIZE_COL, BLOCK_SIZE_ROW), True, tl.int1)
if USE_DROPOUT:
  dropout_seed = SEED + batch_head_offset
  q_offsets_drop = offsets_row[None, :]
  kv_offsets_drop = offsets_col[:, None]
  dropout_offset = q_offsets_drop * SEQ_LEN + kv_offsets_drop
  random_vals = tl.rand(dropout_seed, dropout_offset)
  dropout_mask = random_vals > dropout_p
  P_T_block = tl.where(dropout_mask, P_T_block / (1.0 - dropout_p), 0.0)
  
if USE_DROPOUT:
  dP_T_block = tl.where(dropout_mask, dP_T_block / (1.0 - dropout_p), 0.0)

But then we are always doing a memory allocation even when not using dropout. I could also just copy paste my kernel (one with dropout and one without) and then handle the if/else on which kernel I want to call which would fix this, but I wanted to know if there is a better way to handle this that I am missing!